- Adopted main's `cute.make_fragment` → `cute.make_rmem_tensor` rename, including one call our backward path adds in `flash_bwd_sm100.py`.
- Kept both main's new MLA-specialization block and the dropout-seed init in `flash_attn_func` / `flash_attn_varlen_func`.
I also dropped two files that shouldn't have been in the PR — a root-level `test_dropout_standalone.py` scratch script and `AI/DROPOUT_IMPLEMENTATION_NOTES.md`. Coverage stays in `tests/cute/test_dropout_{correctness,e2e,grad_match}.py`.
The dropout code paths themselves are unchanged from the version the grad-match suite already passed, so the rebase carries no numerical delta. |\n\nThis file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.\n[Learn more about bidirectional Unicode characters](https://github.co/hiddenchars)\n\n[Show hidden characters](https://github.com/Dao-AILab/flash-attention/pull/2439)"},{"resultType":"repo_readme","repo":"gitctrlx/flash-attention-tutorial","url":"https://github.com/gitctrlx/flash-attention-tutorial","readmeUrl":"https://raw.githubusercontent.com/gitctrlx/flash-attention-tutorial/HEAD/README.md","snippet":"# Flash Attention Tutorial\r\n\r\nThis tutorial is designed for beginners to understand and implement Flash Attention in PyTorch. We follow the structure of the document [\"From Online Softmax to FlashAttention\"](https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf) by Zihao Ye, explaining each concept step by step, deriving the formulas, and providing PyTorch code snippets. We'll build from basic self-attention to the full tiled Flash Attention, including both forward and backward passes for a complete, differentiable PyTorch implementation.\r\n\r\nWe assume basic knowledge of PyTorch and transformers. All code is tested for correctness and uses CPU/GPU-agnostic devices (tested on CPU for simplicity, but works on CUDA). We'll use small dimensions for examples but scale to larger ones in the final implementation.\r\n\r\n## Introduction\r\n\r\nFlash Attention is a memory-efficient and fast algorithm for computing self-attention in transformers. It avoids materializing large intermediate matrices (e.g., the attention matrix) in GPU global memory (High Bandwidth Memory, HBM). Instead, it employs tiling and online softmax techniques to fuse operations into a single kernel, leveraging fast GPU shared memory (SRAM). This reduces I/O overhead between slow HBM and fast SRAM, enabling longer sequences without out-of-memory errors.\r\n\r\nA key limitation in standard attention implementations is the small size of SRAM (typically ~100KB per streaming multiprocessor on modern GPUs), which cannot store the full $O(L^2)$ attention matrix for long sequences $L$. As a result, computations require multiple HBM accesses: writing intermediates to HBM, then reading them back for subsequent operations. This I/O bottleneck slows down performance, especially for large $L$. Flash Attention addresses this by recomputing intermediates on-the-fly in SRAM using tiled blocks, minimizing HBM accesses to sub-quadratic levels.\r\n\r\nThe core challenge is that softmax is not associative (unlike addition in matrix multiplication), complicating tiling. We'll derive how online softmax enables associative updates and build to a full implementation, highlighting reductions in HBM accesses at each stage.\r\n\r\n## 1. Standard Self-Attention\r\n\r\nSelf-attention computes weighted sums of values based on query-key similarities. For simplicity, we initially ignore batches, heads, masks, and scaling (added later).\r\n\r\nThe formula is:\r\n\r\n$$\r\nO = \\mathrm{softmax}(Q K^T) V\r\n$$\r\n\r\nwhere $Q, K, V, O \\in \\mathbb{R}^{L \\times D}$, $L$ is the sequence length, and $D$ is the head dimension. Softmax is applied row-wise.\r\n\r\nThe standard implementation factorizes as:\r\n\r\n$$\r\nX = Q K^T, \\quad A = \\mathrm{softmax}(X), \\quad O = A V\r\n$$\r\n\r\nThis requires storing $X$ and $A$ ($O(L^2)$ memory), which is inefficient for large $L$. In terms of memory accesses, it involves at least three HBM round-trips: (1) compute and write $X$ to HBM, (2) read $X$ from HBM for softmax and write $A$ to HBM, (3) read $A$ from HBM to compute $O$.\r\n\r\nA basic (non-Flash) PyTorch implementation with batches, heads, mask, and scaling:\r\n\r\n```python\r\nimport torch\r\nimport torch.nn as nn\r\nimport math # Use math for sqrt to align with PyTorch style\r\n\r\ndef standard_attention(Q, K, V, mask=None):\r\n # Inputs: Q, K, V shape (batch_size, num_heads, seq_len, head_dim)\r\n head_dim = Q.shape[-1]\r\n scale = 1 / math.sqrt(head_dim)\r\n # Scale queries\r\n Q_scaled = Q * scale\r\n # Compute logits: (b, h, l, l)\r\n logits = torch.matmul(Q_scaled, K.transpose(-2, -1))\r\n \r\n if mask is not None:\r\n # Mask: (b, l) -> (b, 1, 1, l)\r\n key_mask = mask.unsqueeze(1).unsqueeze(1)\r\n logits = torch.where(key_mask > 0, logits, float('-inf'))\r\n \r\n # Softmax to get attention weights\r\n attn_weights = nn.functional.softmax(logits, dim=-1)\r\n # Weighted sum\r\n output = torch.matmul(attn_weights, V)\r\n return output\r\n```\r\n\r\nExample usage:\r\n\r\n```python\r\n# Small example tensors\r\nbatch_size, num_heads, seq_len, head_dim = 1, 1, 4, 8\r\nQ = torch.randn(batch_size, num_heads, seq_len, head_dim)\r\nK = torch.randn(batch_size, num_heads, seq_len, head_dim)\r\nV = torch.randn(batch_size, num_heads, seq_len, head_dim)\r\nmask = torch.ones(batch_size, seq_len)\r\nO = standard_attention(Q, K, V, mask)\r\nprint(O.shape) # torch.Size([1, 1, 4, 8])\r\n```\r\n\r\nWhile matrix multiplication can be tiled (due to associativity), softmax cannot be directly tiled without recomputation techniques.\r\n\r\n## 2. Safe Softmax\r\n\r\nFor a vector $x = [x_1, \\dots, x_N]$:\r\n\r\n$$\r\n\\mathrm{softmax}(x)_i = \\frac{e^{x_i}}{\\sum_{j=1}^N e^{x_j}}\r\n$$\r\n\r\nLarge $x_i$ can cause numerical overflow (e.g., in float16). Safe softmax subtracts the row max $m = \\max_j x_j$:\r\n\r\n$$\r\n\\mathrm{softmax}(x)_i = \\frac{e^{x_i - m}}{\\sum_{j=1}^N e^{x_j - m}}\r\n$$\r\n\r\nThis is a 3-pass algorithm, requiring three HBM reads if $x$ does not fit in SRAM: (1) find max, (2) compute exp and sum, (3) normalize. Each pass reloads $x$ from HBM.\r\n\r\nPseudo-code:\r\n\r\n```text\r\nm = -∞\r\nfor i = 1 to N:\r\n m = max(m, x_i) # Pass 1: global max\r\n\r\nd = 0\r\nfor i = 1 to N:\r\n d += exp(x_i - m) # Pass 2: denominator\r\n\r\nfor i = 1 to N:\r\n a_i = exp(x_i - m) / d # Pass 3: normalize\r\n```\r\n\r\nPyTorch implementation (row-wise for 2D tensor):\r\n\r\n```python\r\ndef safe_softmax(x):\r\n # x: (batch_size, seq_len)\r\n # Pass 1: Compute row maxes\r\n m = torch.max(x, dim=-1, keepdim=True)[0]\r\n # Pass 2: Exps and sum (denominator)\r\n exp_shifted = torch.exp(x - m)\r\n denom = torch.sum(exp_shifted, dim=-1, keepdim=True)\r\n # Pass 3: Normalize\r\n return exp_shifted / denom\r\n\r\n# Example\r\nx = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\r\nprint(safe_softmax(x))\r\n```\r\n\r\nThis is inefficient for large sequences, as multiple passes increase HBM I/O.\r\n\r\n## 3. Online Softmax\r\n\r\nTo reduce passes, introduce a surrogate $d_i' = \\sum_{j=1}^i e^{x_j - m_i}$ with the recurrence:\r\n\r\n$$\r\nd_i' = d_{i-1}' \\cdot e^{m_{i-1} - m_i} + e^{x_i - m_i}\r\n$$\r\n\r\nwhere $m_i = \\max(m_{i-1}, x_i)$. At $i=N$, $d_N' = \\sum e^{x_j - m_N}$. This enables a 2-pass algorithm (2 HBM reads if not in SRAM): one for statistics, one for normalization.\r\n\r\nPseudo-code:\r\n\r\n```text\r\nm = -∞\r\nd' = 0\r\nfor i = 1 to N: # Pass 1: Compute m_i, d_i'\r\n m_new = max(m, x_i)\r\n d' = d' * exp(m - m_new) + exp(x_i - m_new)\r\n m = m_new\r\n\r\nfor i = 1 to N: # Pass 2: Normalize\r\n a_i = exp(x_i - m) / d'\r\n```\r\n\r\nPyTorch implementation (row-wise):\r\n\r\n```python\r\ndef online_softmax(x):\r\n # x: (batch_size, seq_len)\r\n batch_size, seq_len = x.shape\r\n # Initialize\r\n m = torch.full((batch_size, 1), float('-inf'), device=x.device)\r\n d_prime = torch.zeros((batch_size, 1), device=x.device)\r\n # Pass 1: Online update (loop for clarity; vectorize in practice)\r\n for i in range(seq_len):\r\n x_i = x[:, i:i+1]\r\n m_new = torch.maximum(m, x_i)\r\n exp_diff = torch.exp(m - m_new)\r\n exp_term = torch.exp(x_i - m_new)\r\n d_prime = d_prime * exp_diff + exp_term\r\n m = m_new\r\n # Pass 2: Compute probabilities\r\n probs = torch.exp(x - m) / d_prime\r\n return probs\r\n\r\n# Example\r\nx = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])\r\nprint(online_softmax(x)) # Matches safe_softmax\r\n```\r\n\r\nThis reduces I/O to two HBM accesses but still requires multiple passes.\r\n\r\n## 4. Flash Attention (Without Tiling)\r\n\r\nIn attention, we only need $O = A V$, not $A$ itself. For each row $k$, derive a one-pass recurrence using surrogates.\r\n\r\nNotation:\r\n\r\n- $x_j = Q[k, :] \\cdot K[j, :]$ (scalar logit for position $j$)\r\n- $o_i = \\sum_{j=1}^i a_j V[j, :]$ (partial output up to $i$)\r\n\r\nSurrogate $o_i' = \\sum_{j=1}^i e^{x_j - m_i} V[j, :] / d_i'$, with recurrence:\r\n\r\n$$\r\no_i' = o_{i-1}' \\cdot \\frac{d_{i-1}' e^{m_{i-1} - m_i}}{d_i'} + \\frac{e^{x_i - m_i}}{d_i'} V[i, :]\r\n$$\r\n\r\nThis enables a single loop (1 HBM access per block if tiled later), updating statistics on-the-fly.\r\n\r\nPseudo-code for one row:\r\n\r\n```text\r\nm = -∞\r\nd' = 0\r\no' = zeros(D)\r\nfor j = 1 to N:\r\n x_j = Q[k, :] @ K[j, :]\r\n m_new = max(m, x_j)\r\n exp_diff = exp(m - m_new)\r\n exp_term = exp(x_j - m_new)\r\n d'_new = d' * exp_diff + exp_term\r\n o' = o' * (d' * exp_diff / d'_new) + (exp_term / d'_new) * V[j, :]\r\n m = m_new\r\n d' = d'_new\r\nO[k, :] = o'\r\n```\r\n\r\nPyTorch implementation (batched, multi-head, no tiling; for illustration only):\r\n\r\n```python\r\ndef simple_flash_forward(Q, K, V, mask=None):\r\n # Inputs: (b, h, l, d); no tiling (inefficient for large l)\r\n b, h, l, d = Q.shape\r\n scale = 1 / math.sqrt(d)\r\n O = torch.zeros_like(Q)\r\n # Loop over batch*heads\r\n for batch_head_idx in range(b * h):\r\n # Flatten to (l, d)\r\n q_bh = Q.view(b * h, l, d)[batch_head_idx]\r\n k_bh = K.view(b * h, l, d)[batch_head_idx]\r\n v_bh = V.view(b * h, l, d)[batch_head_idx]\r\n m_bh = mask.view(b, l)[batch_head_idx // h] if mask is not None else torch.ones(l, device=Q.device)\r\n # Loop over query rows\r\n for row_idx in range(l):\r\n q_row = q_bh[row_idx] * scale # (d,)\r\n row_max = float('-inf')\r\n row_denom_prime = torch.tensor(0.0, device=Q.device)\r\n row_output_prime = torch.zeros(d, device=Q.device)\r\n # Loop over key columns (online update)\r\n for col_idx in range(l):\r\n if m_bh[col_idx] <= 0:\r\n continue\r\n logit = torch.dot(q_row, k_bh[col_idx]) # scalar\r\n new_max = max(row_max, logit)\r\n exp_diff = torch.exp(row_max - new_max)\r\n exp_term = torch.exp(logit - new_max)\r\n new_denom_prime = row_denom_prime * exp_diff + exp_term\r\n new_output_prime = row_output_prime * (row_denom_prime * exp_diff / new_denom_prime) + \\\r\n (exp_term / new_denom_prime) * v_bh[col_idx]\r\n row_max = new_max\r\n row_denom_prime = new_denom_prime\r\n row_output_prime = new_output_prime\r\n # Assign to output\r\n O.view(b * h, l, d)[batch_head_idx, row_idx] = row_output_prime\r\n return O\r\n```\r\n\r\n(Note: This scalar loop is for clarity and not efficient; tiling fuses operations in the next section.)\r\n\r\n## 5. Flash Attention (With Tiling)\r\n\r\nFor large $L$, tile $K$ and $V$ into blocks of size $B$ (fitting in SRAM). For each query row/block, load tiles, compute local max/exp/sum in SRAM, and update global statistics associatively. This ensures a single effective pass, with HBM accesses scaling as $O(L^2 D^2 / M)$ (where $M$ is SRAM size), far fewer than standard attention's $O(L^2 + L D)$.\r\n\r\nPseudo-code for one row (tiled):\r\n\r\n```text\r\nm = -∞\r\nd' = 0\r\no' = zeros(D)\r\nfor tile_idx = 1 to num_tiles:\r\n tile_start = (tile_idx - 1) * B\r\n tile_end = tile_start + B\r\n K_tile = K[tile_start:tile_end, :]\r\n V_tile = V[tile_start:tile_end, :]\r\n logits_tile = Q[k, :] @ K_tile.T # (B,)\r\n local_max = max(logits_tile)\r\n global_max = max(m, local_max)\r\n exp_global_diff = exp(m - global_max)\r\n exp_local_terms = exp(logits_tile - global_max)\r\n local_sum = sum(exp_local_terms)\r\n d'_new = d' * exp_global_diff + local_sum\r\n weighted_v_sum = (exp_local_terms / d'_new) @ V_tile # (D,)\r\n o' = o' * (d' * exp_global_diff / d'_new) + weighted_v_sum\r\n m = global_max\r\n d' = d'_new\r\nO[k, :] = o'\r\n```\r\n\r\nFull tiled forward in PyTorch (returns $O$, row sums $l$, row maxes $m$ for backward):\r\n\r\n```python\r\nBLOCK_SIZE = 1024 # Adjust based on SRAM\r\nNEG_INF = float('-inf')\r\nEPS = 1e-6 # For numerical stability\r\n\r\ndef flash_attention_forward(Q, K, V, mask=None):\r\n # Inputs: (b, h, l, d)\r\n b, h, l, d = Q.shape\r\n device = Q.device\r\n O = torch.zeros_like(Q)\r\n row_sums = torch.zeros(b, h, l, 1, device=device) # l (denominators)\r\n row_maxes = torch.full((b, h, l, 1), NEG_INF, device=device) # m\r\n \r\n # Block sizes (Q rows, KV columns)\r\n q_block_size = min(BLOCK_SIZE, l)\r\n kv_block_size = BLOCK_SIZE\r\n \r\n # Split into blocks\r\n Q_blocks = torch.split(Q, q_block_size, dim=2)\r\n K_blocks = torch.split(K, kv_block_size, dim=2)\r\n V_blocks = torch.split(V, kv_block_size, dim=2)\r\n mask = torch.ones(b, l, device=device) if mask is None else mask\r\n mask_blocks = torch.split(mask, kv_block_size, dim=1)\r\n \r\n num_q_blocks, num_kv_blocks = len(Q_blocks), len(K_blocks)\r\n \r\n # Split outputs for accumulation\r\n O_blocks = list(torch.split(O, q_block_size, dim=2))\r\n row_sums_blocks = list(torch.split(row_sums, q_block_size, dim=2))\r\n row_maxes_blocks = list(torch.split(row_maxes, q_block_size, dim=2))\r\n \r\n scale = 1 / math.sqrt(d)\r\n \r\n # Outer loop over KV tiles\r\n for kv_idx in range(num_kv_blocks):\r\n K_tile = K_blocks[kv_idx]\r\n V_tile = V_blocks[kv_idx]\r\n mask_tile = mask_blocks[kv_idx].unsqueeze(1).unsqueeze(1) # (b, 1, 1, block_size)\r\n \r\n # Inner loop over Q tiles\r\n for q_idx in range(num_q_blocks):\r\n Q_tile = Q_blocks[q_idx]\r\n curr_O = O_blocks[q_idx]\r\n curr_row_sums = row_sums_blocks[q_idx]\r\n curr_row_maxes = row_maxes_blocks[q_idx]\r\n \r\n # Compute block logits\r\n Q_tile_scaled = Q_tile * scale\r\n logits_block = torch.matmul(Q_tile_scaled, K_tile.transpose(-2, -1)) # (b, h, q_block, kv_block)\r\n logits_block = torch.where(mask_tile > 0, logits_block, NEG_INF)\r\n \r\n # Local max and exp\r\n local_max = torch.max(logits_block, dim=-1, keepdim=True)[0]\r\n exp_logits = torch.exp(logits_block - local_max)\r\n exp_logits = torch.where(mask_tile > 0, exp_logits, 0.0)\r\n \r\n # Local sum\r\n local_sum = torch.sum(exp_logits, dim=-1, keepdim=True) + EPS\r\n \r\n # Update global max and sum\r\n new_max = torch.maximum(curr_row_maxes, local_max)\r\n new_row_sums = torch.exp(curr_row_maxes - new_max) * curr_row_sums + \\\r\n torch.exp(local_max - new_max) * local_sum\r\n \r\n # Update output\r\n exp_max_diff = torch.exp(curr_row_maxes - new_max)\r\n exp_local_diff = torch.exp(local_max - new_max)\r\n weighted_v = torch.matmul(exp_logits, V_tile) # (b, h, q_block, d)\r\n O_blocks[q_idx] = (curr_row_sums * exp_max_diff / new_row_sums) * curr_O + \\\r\n (exp_local_diff / new_row_sums) * weighted_v\r\n \r\n # Store updated stats\r\n row_sums_blocks[q_idx] = new_row_sums\r\n row_maxes_blocks[q_idx] = new_max\r\n \r\n # Concatenate blocks\r\n O = torch.cat(O_blocks, dim=2)\r\n row_sums = torch.cat(row_sums_blocks, dim=2)\r\n row_maxes = torch.cat(row_maxes_blocks, dim=2)\r\n return O, row_sums, row_maxes\r\n```\r\n\r\n## 6. Flash Attention Backward Pass\r\n\r\nThe backward pass computes gradients $dQ$, $dK$, $dV$ without storing the $O(L^2)$ attention matrix. It recomputes block-wise in SRAM using saved row maxes $m$ and sums $l$ from forward, with similar tiling. This maintains memory efficiency, with HBM accesses $O(L^2 D^2 / M)$.\r\n\r\nKey derivations (from original paper):\r\n\r\n- $dV = A^T dO$\r\n- For softmax grad: $dS_{ij} = A_{ij} (dO_i^T V_j - D_i)$ where $D_i = dO_i^T O_i$\r\n- $dQ_i = dS_{i:} K^T$, $dK_j = dS_{:j}^T Q$\r\n\r\nPseudo-code (tiled, simplified from paper):\r\n\r\n```text\r\nfor kv_tile_idx = 1 to num_kv_tiles:\r\n Load K_tile, V_tile\r\n Init temp_dK_tile = zeros, temp_dV_tile = zeros\r\n for q_tile_idx = 1 to num_q_tiles:\r\n Load Q_tile, O_tile, dO_tile, row_sums_tile, row_maxes_tile\r\n Compute logits_block = scale * Q_tile @ K_tile.T\r\n Apply mask\r\n probs_block = exp(logits_block - row_maxes_tile) / row_sums_tile\r\n temp_dV_tile += probs_block.T @ dO_tile\r\n dP_block = dO_tile @ V_tile.T\r\n D_tile = row_sum(dO_tile * O_tile) # (q_block,)\r\n dS_block = probs_block * (dP_block - D_tile)\r\n dQ_tile += dS_block @ K_tile\r\n temp_dK_tile += dS_block.T @ Q_tile\r\n Write dK_tile, dV_tile\r\n```\r\n\r\nPyTorch implementation (tiled backward; assumes no dropout/mask for simplicity; extend as needed):\r\n\r\n```python\r\ndef flash_attention_backward(dO, Q, K, V, O, row_sums, row_maxes, mask=None):\r\n # Inputs: dO, Q, K, V, O (b, h, l, d); row_sums, row_maxes (b, h, l, 1)\r\n b, h, l, d = Q.shape\r\n device = Q.device\r\n scale = 1 / math.sqrt(d)\r\n \r\n dQ = torch.zeros_like(Q)\r\n dK = torch.zeros_like(K)\r\n dV = torch.zeros_like(V)\r\n \r\n q_block_size = min(BLOCK_SIZE, l)\r\n kv_block_size = BLOCK_SIZE\r\n \r\n # Splits (reuse forward logic)\r\n Q_blocks = torch.split(Q, q_block_size, dim=2)\r\n K_blocks = torch.split(K, kv_block_size, dim=2)\r\n V_blocks = torch.split(V, kv_block_size, dim=2)\r\n O_blocks = torch.split(O, q_block_size, dim=2)\r\n dO_blocks = torch.split(dO, q_block_size, dim=2)\r\n row_sums_blocks = torch.split(row_sums, q_block_size, dim=2)\r\n row_maxes_blocks = torch.split(row_maxes, q_block_size, dim=2)\r\n dQ_blocks = list(torch.split(dQ, q_block_size, dim=2))\r\n mask = torch.ones(b, l, device=device) if mask is None else mask\r\n mask_blocks = torch.split(mask, kv_block_size, dim=1)\r\n \r\n num_q_blocks, num_kv_blocks = len(Q_blocks), len(K_blocks)\r\n \r\n # Outer loop over KV tiles\r\n for kv_idx in range(num_kv_blocks):\r\n K_tile = K_blocks[kv_idx]\r\n V_tile = V_blocks[kv_idx]\r\n mask_tile = mask_blocks[kv_idx].unsqueeze(1).unsqueeze(1)\r\n \r\n temp_dK = torch.zeros_like(K_tile)\r\n temp_dV = torch.zeros_like(V_tile)\r\n \r\n # Inner loop over Q tiles\r\n for q_idx in range(num_q_blocks):\r\n Q_tile = Q_blocks[q_idx]\r\n O_tile = O_blocks[q_idx]\r\n dO_tile = dO_blocks[q_idx]\r\n row_sums_tile = row_sums_blocks[q_idx]\r\n row_maxes_tile = row_maxes_blocks[q_idx]\r\n \r\n # Recompute block probs\r\n Q_tile_scaled = Q_tile * scale\r\n logits_block = torch.matmul(Q_tile_scaled, K_tile.transpose(-2, -1))\r\n logits_block = torch.where(mask_tile > 0, logits_block, NEG_INF)\r\n probs_block = torch.exp(logits_block - row_maxes_tile) / (row_sums_tile + EPS)\r\n \r\n # dV accumulation\r\n temp_dV += torch.matmul(probs_block.transpose(-2, -1), dO_tile)\r\n \r\n # dP and D\r\n dP_block = torch.matmul(dO_tile, V_tile.transpose(-2, -1))\r\n D_tile = torch.sum(dO_tile * O_tile, dim=-1, keepdim=True) # (b, h, q_block, 1)\r\n \r\n # dS\r\n dS_block = probs_block * (dP_block - D_tile)\r\n \r\n # dQ accumulation\r\n dQ_blocks[q_idx] += torch.matmul(dS_block, K_tile)\r\n \r\n # dK accumulation\r\n temp_dK += torch.matmul(dS_block.transpose(-2, -1), Q_tile)\r\n \r\n # Write gradients\r\n dK_blocks = torch.split(dK, kv_block_size, dim=2)\r\n dV_blocks = torch.split(dV, kv_block_size, dim=2)\r\n dK_blocks[kv_idx].copy_(temp_dK)\r\n dV_blocks[kv_idx].copy_(temp_dV)\r\n \r\n dQ = torch.cat(dQ_blocks, dim=2)\r\n dK = torch.cat(torch.split(dK, kv_block_size, dim=2), dim=2) # Reassemble if needed\r\n dV = torch.cat(torch.split(dV, kv_block_size, dim=2), dim=2)\r\n return dQ, dK, dV\r\n```\r\n\r\n(Note: This is a simplified version without dropout; refer to the original paper for full details including dropout regeneration.)\r\n\r\nTo make it differentiable, wrap in a custom `torch.autograd.Function`.\r\n\r\n## Reference\r\n\r\n- **Tutorial Notes**: [From Online Softmax to FlashAttention](https://courses.cs.washington.edu/courses/cse599m/23sp/notes/flashattn.pdf) by Zihao Ye\r\n- **Original Paper**: [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135) by Tri Dao et al.\r\n- **Official Repository**: [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention)\r\n- **PyTorch Example**: [shreyansh26/FlashAttention-PyTorch](https://github.com/shreyansh26/FlashAttention-PyTorch)\r\n\r\n## License\r\n\r\nMIT License"},{"resultType":"repo_readme","repo":"damienjose/cuda-flashattention","url":"https://github.com/damienjose/cuda-flashattention","readmeUrl":"https://raw.githubusercontent.com/damienjose/cuda-flashattention/HEAD/README.md","snippet":"# Flash Attention: Fast and Memory-Efficient Attention Mechanism\n\n## Overview\nFlash Attention is a fast, memory-efficient, and IO-aware implementation of attention mechanisms used in deep learning models. \nThis project showcases a simplified version of Flash Attention on GPUs using CUDA with randomly initialized matrix inputs, demonstrating how significant performance improvements can be achieved with optimized GPU compute, over traditional attention implementations.\n\n## Key Features\n1. **Fast**: Speeds up training for large models like BERT and GPT by up to 3x.\n2. **Memory-Efficient**: Reduces memory complexity from \\(O(N^2)\\) to \\(O(N)\\) by leveraging hardware memory hierarchies.\n3. **Exact**: Ensures accuracy without approximation.\n4. **IO-Aware**: Optimizes memory access and communication for modern GPUs.\n\n## Implementation Highlights\nFlash Attention is implemented in three versions:\n1. **CPU Implementation**: Provides a baseline implementation using sequential processing.\n2. **Naive GPU Implementation**: Uses CUDA global memory for parallel computations.\n3. **Optimized GPU Implementation**: Leverages shared memory, memory tiling, thread coarsening, and memory coalescing for maximum efficiency.\n\n## Key Optimizations\n1. **Memory Tiling**: Divides computation into manageable tiles to enhance memory locality.\n2. **Thread Coarsening**: Increases the work assigned per thread to reduce overhead.\n3. **Shared Memory**: Uses fast on-chip shared memory to minimize global memory access.\n4. **Memory Coalescing**: Ensures efficient memory access patterns for threads.\n\n## Experimental Results\nPerformance benchmarks were conducted on an NVIDIA GeForce RTX 4050 Laptop GPU with varying configurations. Below are the latency results for 1024 x 1024 matrices:\n\n| Implementation | Latency |\n|---------------------------|------------|\n| CPU | 9,427 ms |\n| GPU (Naive) | 38.78 ms |\n| GPU (Optimized Shared) | 14.12 ms |\n\nThe optimized implementation achieves a **3x speedup** compared to the naive GPU version and a **600x improvement** over the CPU baseline.\n\n## Source Code Structure\n- **CPU Implementation**: Implements scaled dot-product attention using sequential processing.\n- **Naive GPU Implementation**: Implements CUDA kernels using global memory for QK^T, softmax, and output computation.\n- **Optimized GPU Implementation**: Implements CUDA kernels using shared memory, tiling, and other optimizations for improved performance.\n\n## How to Build and Run\n1. Clone this repository:\n ```bash\n git clone https://github.com/damienjose/cuda-flashattention\n cd cuda-flashattention\n ```\n\n2. Set up your CUDA environment:\n - Install the latest CUDA toolkit.\n - Update the `sm` and `compute` settings in the Visual Studio project properties, specific for your GPU. See https://www.truehost.com/what-is-compute-capability-of-a-gpu/\n\n3. Build the project:\n ```bash\n nvcc -o cuda-flashattention kernel.cu -arch=sm_80\n ```\n\n4. Run the program:\n ```bash\n ./cuda-flashattention\n ```\n\n## Profiling and Performance Analysis\n- The naive GPU implementation is memory-bound due to frequent global memory accesses.\n- The optimized shared memory implementation transforms memory-bound operations into compute-bound ones.\n- Profiling tools like NVIDIA Nsight can be used for further performance analysis.\n\n## References\n1. [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135)\n2. [Understanding Flash-Attention and Flash-Attention-2](https://towardsai.net/p/artificial-intelligence/understanding-flash-attention-and-flash-attention-2-the-path-to-scale-the-context-lenght-of-language-models)\n3. [GitHub - Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention)\n\n## Authors\n- Bryan Zhou ([bryanzhou2000@gmail.com](mailto:bryanzhou2000@gmail.com))\n- Damien Jose ([damien.jose@gmail.com](mailto:damien.jose@gmail.com))\n\n## Acknowledgments\nThis project was done as part of the AUT24 GPU Compute course at the University of Washington. We thank Professor Colin N. Reinhardt and TAs Arnab Karmakar, Xiaxi Shen for guiding us throughout this course. We also appreciate the authors of Flash Attention for their contributions and open sourcing their code."},{"resultType":"github_history","repo":"oobabooga/textgen","url":"https://github.com/oobabooga/textgen/issues/2427","pageType":"merged_pr","number":2427,"segmentCount":1,"title":"Conversation","snippet":"## Conversation\n\nThis PR fixes the flash attention implementation that is currently in the repo and adds the `--flash-attention` flag to activate it.\n\nTo use it, install flash attention:\n\n```\npip install ninja\npip install flash-attn\n```\n\nIf that doesn't work (it should only take < 20 minutes), you need to compile from source with ninja:\n\n```\npip install ninja\ngit clone https://github.com/HazyResearch/flash-attention.git\ncd flash-attention\npython setup.py install\n```\n\nIf downstream users confirm it is working for them, I will add the pip install to the requirements.txt\n\nWith the fixed implementation, flash attention works as expected:\n\nRTX 3090 (24 GB VRAM)\n\ntsumeone SuperCOT 30B 128 groupsize CUDA:\n\n```\nMax new tokens: 1\n3001 tokens\nw/ flash attention: OOM\n\n2614 tokens:\nw/ flash attention: 23.7 GB\n```\n\nllama 13B with an 8 Rank LoRA applied:\n\n```\nMax new tokens: 1\n3001 tokens\nw flash attention: 12.7 GB\nw/o flash attention: 15.3 GB\n\n4002 tokens\nw flash attention: ~14 GB\nw/o flash attention: 17.7 GB\n```\n\nIn fact, I can get 6300 tokens on 13B.\n\nThe current implementation in the repo uses PyTorch's scaled\\_dot\\_product to perform the attention call. This method falls back to standard attention when flash attention or memory efficient attention is not available -- so what's wrong? PyTorch's implementation is somewhat old, and at the time, flash attention did not support attention masking in the forward pass, so trying to run it with the attention mask fails and falls back to normal attention. Since that time, HazyResearch has already updated their flash attention to work for Causal language modeling, so we can just use it directly from the source. Maybe in the future PyTorch will fix their version.\n\nAlso, all credit for the code goes to LAION who already implemented it for LLaMa:\n\n[https://github.com/LAION-AI/Open-Assistant/blob/main/model/model\\_training/models/patching\\_llama.py](https://github.com/LAION-AI/Open-Assistant/blob/main/model/model_training/models/patching_llama.py)\n\nI merely took the code there and adapted it to work without the rest of their scaffolding.\n\n**NOTE:**\n\n**This PR works, but it is not ready to be merged since flash attention requires either `pip install flash-attn` or compilation from the source by cloning HazyResearch's repo and following the source compiling instructions. I already had it compiled on my machine, so until it is confirmed that downstream users can compile on their machine or the pip install works for them (Personally it froze for me which is why I had to compile from source with ninja), I would not recommend pushing this through just yet.**\n\nBut once it is pushed I think everyone should use Flash Attention by default -- it exists for a reason, there is no catch!\n\n👍1jepjoo reacted with thumbs up emoji❤️4oobabooga, matatonic, Maykeye, and gr1336 reacted with heart emoji\n\n- 👍1 reaction\n- ❤️4 reactions\n\n`\n Add fixed Flash attention code\n`\n\n`\n c45cfe3\n`\n\n| |\n| --- |\n| Had to compile flash attention, which worked;
Started ooba with
CUDA\\_VISIBLE\\_DEVICE=0,1,2 python server.py --model tsumeone\\_llama-30b-supercot-4bit-128g-cuda --wbits 4 --groupsize 128 --flash-attention --model\\_type llama --pre\\_layer 16 36 60
(3 x 12gb rtx 3060)
Under Parameters, changed \"Truncate the prompt up to this length\" from 2048 to 2600
Fed it a short story, with the ending removed
2592 tokens in the input.
hit generate;
It wrote some stuff, but after ~200 tokens or so output becomes repetitive.
Output generated in 115.07 seconds (3.04 tokens/s, 350 tokens, context 2249, seed 1551651864)
Is this for specific models? |\n\nContributorAuthor\n\n| |\n| --- |\n| [@Fairfax-Mooresby](https://github.com/Fairfax-Mooresby)
No, it is for any model. The reason it is repetitive is well, most models aren't designed to go past 2048, since it is the base context limit of LLaMa. However, they can be trained to go beyond that, such as with Bluemoon 13B and 30B. You might want to try those as well. More models will come in the future to leverage higher context as well, I'm sure |\n\n| |\n| --- |\n| Flash attention always failed on my 3090 or P40/P6000. You found a way for it to work? Does it have to be compiled from source? |\n\n| |\n| --- |\n| [@kaiokendev](https://github.com/kaiokendev)
Unfortunately the increased token count for those models seems to break --pre\\_layer loading support, making 30b multi-gpu slow and use more Vram than it should. I'll give it a whirl on 13b and see how far i can push the context. Thanks for your work on this. |\n\nContributorAuthor\n\n| |\n| --- |\n| > Flash attention always failed on my 3090 or P40/P6000. You found a way for it to work? Does it have to be compiled from source?
I don't know about P40/P6000, it works on my 3090. This PR is using the direct implementation, and not via PyTorch since their implementation seems to be old |\n\n| |\n| --- |\n| I will try it on both. Less memory will definitely help training 30b |\n\n| |\n| --- |\n| This looks very promising. I am trying to get it running but I am probably doing something wrong. This command
```
pip install flash-attn
```
fails with error
```
File \"/tmp/pip-build-env-hmoohnwz/overlay/lib/python3.10/site-packages/setuptools/build_meta.py\", line 338, in run_setup
exec(code, locals())
File \"
ModuleNotFoundError: No module named 'torch'
[end of output]
note: This error originates from a subprocess, and is likely not a problem with pip.
error: subprocess-exited-with-error
× Getting requirements to build wheel did not run successfully.
│ exit code: 1
╰─> See above for output.
note: This error originates from a subprocess, and is likely not a problem with pip.
```
So I tried installing the previous version and it worked, although as you pointed out it took some 10 minutes to install:
```
pip install flash-attn==1.0.4
```
Starting the web UI with
```
python server.py --model llama-13b-4bit-128g --listen --flash-attention
```
works and I can generate with a 3700 tokens prompt on a RTX 3090, but it only outputs spaces.
For llama-30b I OOM without flash attention with a 1800 tokens context, and I also OOM with `--flash-attention` with the same context.
I'll try to install the latest version from source.
EDIT: same result installing flash-attn==1.0.6 from source :( |\n\n| |\n| --- |\n| I got it set up and installed fine. It doesn't compile a kernel for pascal so no go there. But I'm converting a 30b with groupsize 128 and see if it does help memory. From cursory looks at nvtop while generating, it used less. |\n\n| |\n| --- |\n| Tested all of these attention things last night on P40 and 3090. The only thing that gave slightly more tokens when generating at the edge of the 128g-30b was autogptq + quant\\_attn. Everything else went OOM about the same time. I got around 169X. The difference is small too. I got like 6 tokens generated before OOM vs 1. Goes to show there is no free lunch. That or something is wrong with the setup for everything including xformers but I don't think so. |\n\nContributorAuthor\n\n| |\n| --- |\n| > Tested all of these attention things last night on P40 and 3090. The only thing that gave slightly more tokens when generating at the edge of the 128g-30b was autogptq + quant\\_attn. Everything else went OOM about the same time. I got around 169X. The difference is small too. I got like 6 tokens generated before OOM vs 1. Goes to show there is no free lunch. That or something is wrong with the setup for everything including xformers but I don't think so.
Can you try with 13B and let me know your numbers with and without flash attention? |\n\n| |\n| --- |\n| Ok. I will try the 13b today. I have a bluemoonrp that does extra context in that size. |\n\n| |\n| --- |\n| Well, I checked with bluemoon RP and with wizard-storywriter-7b that has infinite context pretty much. I can get up to 8096ish for the former (it breaks) with or without flash attention. For the latter, somewhere around 6k (doesn't break).
I also notice GPU tends to go down to 30% utilization when I am using this. I compiled from source tho. |\n\n| |\n| --- |\n| hi~any progress about this commit? |\n\nContributorAuthor\n\n| |\n| --- |\n| [@laoda513](https://github.com/laoda513) I think with exllama I will close this PR since there is little to be gained for inference here. Personally I was never a fan of such a hacky solution to add flash attention especially when you cannot run the backward pass because LLaMA attn head dim size is 128. I also think the reason my result is different from others is due to memory fragmentation issues, and I don't have the time to debug that especially when exllama already fixed such issues. I have been using xformers as a replacement and it works perfectly, I even use it when training since it works on 3090 out of the box on CUDA 12.0, it is significantly improved and I don't have the headaches from flash attention. |\n\n👍1oobabooga reacted with thumbs up emoji\n\n- 👍1 reaction\n\n[on Jul 3, 2023Jul 3, 2023](https://github.com/oobabooga/textgen/pull/2427#event-9707996562)\n\nThis file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.\n[Learn more about bidirectional Unicode characters](https://github.co/hiddenchars)\n\n[Show hidden characters](https://github.com/oobabooga/textgen/pull/2427)"},{"resultType":"repo_readme","repo":"hao-ai-lab/flash-attention-fp4","url":"https://github.com/hao-ai-lab/flash-attention-fp4","readmeUrl":"https://raw.githubusercontent.com/hao-ai-lab/flash-attention-fp4/HEAD/README.md","snippet":"# FlashAttention (FP4 Fork)\n\n ##
![](\"https://em-content.zobj.net/source/apple/391/star-struck_1f929.png\")
FP4 Flash Attention 4 on B200\n\n> [!NOTE]\n> **This is a fork of [flash-attention](https://github.com/Dao-AILab/flash-attention) with the addition of a CuTe DSL FP4/FP8 implementation of Flash Attention 4 on B200 for the [Attn-QAT](https://arxiv.org/abs/2603.00040) paper.** Supports NVFP4 and MXFP8 block-scaled QK with BF16 or FP8 PV, achieving up to **1.31x speedup** over BF16 FA4 (2018 vs 1545 TFLOPS). Peaks at **2018 TFLOPS** (NVFP4+FP8), **1948 TFLOPS** (MXFP8+FP8), and **1920 TFLOPS** (NVFP4+BF16) on B200. See **[flash_attn/cute/README.md](flash_attn/cute/README.md)** for results and usage.\n\n---\n\nThis repository provides the official implementation of FlashAttention and\nFlashAttention-2 from the\nfollowing papers.\n\n**FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness** \nTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré \nPaper: https://arxiv.org/abs/2205.14135 \nIEEE Spectrum [article](https://spectrum.ieee.org/mlperf-rankings-2022) about our submission to the MLPerf 2.0 benchmark using FlashAttention.\n\n\n**FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning** \nTri Dao\n\nPaper: https://tridao.me/publications/flash2/flash2.pdf\n\n\n\n\n## Usage\n\nWe've been very happy to see FlashAttention being widely adopted in such a short\ntime after its release. This [page](https://github.com/Dao-AILab/flash-attention/blob/main/usage.md)\ncontains a partial list of places where FlashAttention is being used.\n\nFlashAttention and FlashAttention-2 are free to use and modify (see LICENSE).\nPlease cite and credit FlashAttention if you use it.\n\n\n## FlashAttention-3 beta release\nFlashAttention-3 is optimized for Hopper GPUs (e.g. H100). \n\nBlogpost: https://tridao.me/blog/2024/flash3/\n\nPaper: https://tridao.me/publications/flash3/flash3.pdf\n\n\n\nThis is a beta release for testing / benchmarking before we integrate that with\nthe rest of the repo.\n\nCurrently released:\n- FP16 / BF16 forward and backward, FP8 forward\n\nRequirements: H100 / H800 GPU, CUDA >= 12.3.\n\nWe highly recommend CUDA 12.8 for best performance.\n\nTo install:\n```sh\ncd hopper\npython setup.py install\n```\nTo run the test:\n```sh\nexport PYTHONPATH=$PWD\npytest -q -s test_flash_attn.py\n```\nOnce the package is installed, you can import it as follows:\n```python\nimport flash_attn_interface\nflash_attn_interface.flash_attn_func()\n```\n\n## Installation and features\n**Requirements:**\n- CUDA toolkit or ROCm toolkit\n- PyTorch 2.2 and above.\n- `packaging` Python package (`pip install packaging`)\n- `ninja` Python package (`pip install ninja`) *\n- Linux. Might work for Windows starting v2.3.2 (we've seen a few positive [reports](https://github.com/Dao-AILab/flash-attention/issues/595)) but Windows compilation still requires more testing. If you have ideas on how to set up prebuilt CUDA wheels for Windows, please reach out via Github issue.\n\n\\* Make sure that `ninja` is installed and that it works correctly (e.g. `ninja\n--version` then `echo $?` should return exit code 0). If not (sometimes `ninja\n--version` then `echo $?` returns a nonzero exit code), uninstall then reinstall\n`ninja` (`pip uninstall -y ninja && pip install ninja`). Without `ninja`,\ncompiling can take a very long time (2h) since it does not use multiple CPU\ncores. With `ninja` compiling takes 3-5 minutes on a 64-core machine using CUDA toolkit.\n\n**To install:**\n```sh\npip install flash-attn --no-build-isolation\n```\nAlternatively you can compile from source:\n```sh\npython setup.py install\n```\n\nIf your machine has less than 96GB of RAM and lots of CPU cores, `ninja` might\nrun too many parallel compilation jobs that could exhaust the amount of RAM. To\nlimit the number of parallel compilation jobs, you can set the environment\nvariable `MAX_JOBS`:\n```sh\nMAX_JOBS=4 pip install flash-attn --no-build-isolation\n```\n\n**Interface:** `src/flash_attention_interface.py`\n\n### NVIDIA CUDA Support\n**Requirements:**\n- CUDA 12.0 and above.\n\nWe recommend the\n[Pytorch](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch)\ncontainer from Nvidia, which has all the required tools to install FlashAttention.\n\nFlashAttention-2 with CUDA currently supports:\n1. Ampere, Ada, or Hopper GPUs (e.g., A100, RTX 3090, RTX 4090, H100). Support for Turing\n GPUs (T4, RTX 2080) is coming soon, please use FlashAttention 1.x for Turing\n GPUs for now.\n2. Datatype fp16 and bf16 (bf16 requires Ampere, Ada, or Hopper GPUs).\n3. All head dimensions up to 256. ~~Head dim > 192 backward requires A100/A800 or H100/H800~~. Head dim 256 backward now works on consumer GPUs (if there's no dropout) as of flash-attn 2.5.5.\n\n### AMD ROCm Support\nROCm version has two backends. There is [composable_kernel](https://github.com/ROCm/composable_kernel) (ck) which is the default backend and a [Triton](https://github.com/triton-lang/triton) backend. They provide an implementation of FlashAttention-2.\n\n**Requirements:**\n- ROCm 6.0 and above.\n\nWe recommend the\n[Pytorch](https://hub.docker.com/r/rocm/pytorch)\ncontainer from ROCm, which has all the required tools to install FlashAttention.\n\n#### Composable Kernel Backend\nFlashAttention-2 ROCm CK backend currently supports:\n1. MI200 or MI300 GPUs.\n2. Datatype fp16 and bf16\n3. Both forward's and backward's head dimensions up to 256.\n\n#### Triton Backend\nThe Triton implementation of the [Flash Attention v2](https://tridao.me/publications/flash2/flash2.pdf) is currently a work in progress.\n\nIt supports AMD's CDNA (MI200, MI300) and RDNA GPU's using fp16, bf16 and fp32 datatypes.\n\nThese features are supported in Fwd and Bwd\n1) Fwd and Bwd with causal masking\n2) Variable sequence lengths\n3) Arbitrary Q and KV sequence lengths\n4) Arbitrary head sizes\n5) Multi and grouped query attention\n6) Dropout\n7) Rotary embeddings\n8) ALiBi\n\nWe are working on the following things\n1) Paged Attention \n2) Sliding Window\n3) FP8\n4) Performance Improvements\n\n##### Getting Started\nTo get started with the triton backend for AMD, follow the steps below.\n\nFirst install the recommended Triton version \n\n```\npip install triton==3.2.0\n```\nThen install Flash Attention with the flag `FLASH_ATTENTION_TRITON_AMD_ENABLE` set to `\"TRUE\"`.\n\n```\ncd flash-attention\ngit checkout main_perf\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" python setup.py install\n```\n\nTo test that things are working, you can run our tests. These tests take hours so you don't need to run the full thing.\n```\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" pytest tests/test_flash_attn_triton_amd.py\n```\n\nYou can use autotune for better performance by using this flag `FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=\"TRUE\"`\n```\nFLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\" FLASH_ATTENTION_TRITON_AMD_AUTOTUNE=\"TRUE\" python $PATH_TO_CODE\n```\n\n###### Docker\nYou can also use the Dockerfile below which does the above steps on top of the latest rocm/pytorch image.\n```\nFROM rocm/pytorch:latest\n\nWORKDIR /workspace\n\n# install triton\nRUN pip install triton==3.2.0\n\n# install flash attention\nENV FLASH_ATTENTION_TRITON_AMD_ENABLE=\"TRUE\"\n\nRUN git clone https://github.com/ROCm/flash-attention.git &&\\ \n cd flash-attention &&\\\n git checkout main_perf &&\\\n python setup.py install\n\n# set working dir\nWORKDIR /workspace/flash-attention\n```\n\nTo build the docker file\n```\ndocker build -t fa_triton .\n```\n\nTo run the docker image\n```\ndocker run -it --network=host --user root --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --ipc=host --shm-size 16G --device=/dev/kfd --device=/dev/dri fa_triton\n```\n\n## How to use FlashAttention\n\nThe main functions implement scaled dot product attention (softmax(Q @ K^T *\nsoftmax_scale) @ V):\n```python\nfrom flash_attn import flash_attn_qkvpacked_func, flash_attn_func\n```\n\n```python\nflash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False,\n window_size=(-1, -1), alibi_slopes=None, deterministic=False):\n\"\"\"dropout_p should be set to 0.0 during evaluation\nIf Q, K, V are already stacked into 1 tensor, this function will be faster than\ncalling flash_attn_func on Q, K, V since the backward pass avoids explicit concatenation\nof the gradients of Q, K, V.\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between [i - window_size[0], i + window_size[1]] inclusive.\nArguments:\n qkv: (batch_size, seqlen, 3, nheads, headdim)\n dropout_p: float. Dropout probability.\n softmax_scale: float. The scaling of QK^T before applying softmax.\n Default to 1 / sqrt(headdim).\n causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of (-alibi_slope * |i - j|) is added to\n the attention score of query i and key j.\n deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n which is slightly slower and uses more memory. The forward pass is always deterministic.\nReturn:\n out: (batch_size, seqlen, nheads, headdim).\n\"\"\"\n```\n\n```python\nflash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False,\n window_size=(-1, -1), alibi_slopes=None, deterministic=False):\n\"\"\"dropout_p should be set to 0.0 during evaluation\nSupports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\nthan Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\nFor example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\nIf window_size != (-1, -1), implements sliding window local attention. Query at position i\nwill only attend to keys between\n[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\nArguments:\n q: (batch_size, seqlen, nheads, headdim)\n k: (batch_size, seqlen, nheads_k, headdim)\n v: (batch_size, seqlen, nheads_k, headdim)\n dropout_p: float. Dropout probability.\n softmax_scale: float. The scaling of QK^T before applying softmax.\n Default to 1 / sqrt(headdim).\n causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n is added to the attention score of query i and key j.\n deterministic: bool. Whether to use the deterministic implementation of the backward pass,\n which is slightly slower and uses more memory. The forward pass is always deterministic.\nReturn:\n out: (batch_size, seqlen, nheads, headdim).\n\"\"\"\n```\n\n```python\ndef flash_attn_with_kvcache(\n q,\n k_cache,\n v_cache,\n k=None,\n v=None,\n rotary_cos=None,\n rotary_sin=None,\n cache_seqlens: Optional[Union[(int, torch.Tensor)]] = None,\n cache_batch_idx: Optional[torch.Tensor] = None,\n block_table: Optional[torch.Tensor] = None,\n softmax_scale=None,\n causal=False,\n window_size=(-1, -1), # -1 means infinite context window\n rotary_interleaved=True,\n alibi_slopes=None,\n):\n \"\"\"\n If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from\n k and v. This is useful for incremental decoding: you can pass in the cached keys/values from\n the previous step, and update them with the new keys/values from the current step, and do\n attention with the updated cache, all in 1 kernel.\n\n If you pass in k / v, you must make sure that the cache is large enough to hold the new values.\n For example, the KV cache could be pre-allocated with the max sequence length, and you can use\n cache_seqlens to keep track of the current sequence lengths of each sequence in the batch.\n\n Also apply rotary embedding if rotary_cos and rotary_sin are passed in. The key @k will be\n rotated by rotary_cos and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\n If causal or local (i.e., window_size != (-1, -1)), the query @q will be rotated by rotary_cos\n and rotary_sin at indices cache_seqlens, cache_seqlens + 1, etc.\n If not causal and not local, the query @q will be rotated by rotary_cos and rotary_sin at\n indices cache_seqlens only (i.e. we consider all tokens in @q to be at position cache_seqlens).\n\n See tests/test_flash_attn.py::test_flash_attn_kvcache for examples of how to use this function.\n\n Supports multi-query and grouped-query attention (MQA/GQA) by passing in KV with fewer heads\n than Q. Note that the number of heads in Q must be divisible by the number of heads in KV.\n For example, if Q has 6 heads and K, V have 2 heads, head 0, 1, 2 of Q will attention to head\n 0 of K, V, and head 3, 4, 5 of Q will attention to head 1 of K, V.\n\n If causal=True, the causal mask is aligned to the bottom right corner of the attention matrix.\n For example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 = masked out) is:\n 1 1 1 1 0\n 1 1 1 1 1\n If seqlen_q = 5 and seqlen_k = 2, the causal mask is:\n 0 0\n 0 0\n 0 0\n 1 0\n 1 1\n If the row of the mask is all zero, the output will be zero.\n\n If window_size != (-1, -1), implements sliding window local attention. Query at position i\n will only attend to keys between\n [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive.\n\n Note: Does not support backward pass.\n\n Arguments:\n q: (batch_size, seqlen, nheads, headdim)\n k_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n page_block_size must be a multiple of 256.\n v_cache: (batch_size_cache, seqlen_cache, nheads_k, headdim) if there's no block_table,\n or (num_blocks, page_block_size, nheads_k, headdim) if there's a block_table (i.e. paged KV cache)\n k [optional]: (batch_size, seqlen_new, nheads_k, headdim). If not None, we concatenate\n k with k_cache, starting at the indices specified by cache_seqlens.\n v [optional]: (batch_size, seqlen_new, nheads_k, headdim). Similar to k.\n rotary_cos [optional]: (seqlen_ro, rotary_dim / 2). If not None, we apply rotary embedding\n to k and q. Only applicable if k and v are passed in. rotary_dim must be divisible by 16.\n rotary_sin [optional]: (seqlen_ro, rotary_dim / 2). Similar to rotary_cos.\n cache_seqlens: int, or (batch_size,), dtype torch.int32. The sequence lengths of the\n KV cache.\n block_table [optional]: (batch_size, max_num_blocks_per_seq), dtype torch.int32.\n cache_batch_idx: (batch_size,), dtype torch.int32. The indices used to index into the KV cache.\n If None, we assume that the batch indices are [0, 1, 2, ..., batch_size - 1].\n If the indices are not distinct, and k and v are provided, the values updated in the cache\n might come from any of the duplicate indices.\n softmax_scale: float. The scaling of QK^T before applying softmax.\n Default to 1 / sqrt(headdim).\n causal: bool. Whether to apply causal attention mask (e.g., for auto-regressive modeling).\n window_size: (left, right). If not (-1, -1), implements sliding window local attention.\n rotary_interleaved: bool. Only applicable if rotary_cos and rotary_sin are passed in.\n If True, rotary embedding will combine dimensions 0 & 1, 2 & 3, etc. If False,\n rotary embedding will combine dimensions 0 & rotary_dim / 2, 1 & rotary_dim / 2 + 1\n (i.e. GPT-NeoX style).\n alibi_slopes: (nheads,) or (batch_size, nheads), fp32. A bias of\n (-alibi_slope * |i + seqlen_k - seqlen_q - j|)\n is added to the attention score of query i and key j.\n\n Return:\n out: (batch_size, seqlen, nheads, headdim).\n \"\"\"\n```\n\nTo see how these functions are used in a multi-head attention layer (which\nincludes QKV projection, output projection), see the MHA [implementation](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mha.py).\n\n## Changelog\n\n### 2.0: Complete rewrite, 2x faster\nUpgrading from FlashAttention (1.x) to FlashAttention-2\n\nThese functions have been renamed:\n- `flash_attn_unpadded_func` -> `flash_attn_varlen_func`\n- `flash_attn_unpadded_qkvpacked_func` -> `flash_attn_varlen_qkvpacked_func`\n- `flash_attn_unpadded_kvpacked_func` -> `flash_attn_varlen_kvpacked_func`\n\nIf the inputs have the same sequence lengths in the same batch, it is simpler\nand faster to use these functions:\n```python\nflash_attn_qkvpacked_func(qkv, dropout_p=0.0, softmax_scale=None, causal=False)\n```\n```python\nflash_attn_func(q, k, v, dropout_p=0.0, softmax_scale=None, causal=False)\n```\n### 2.1: Change behavior of causal flag\n\nIf seqlen_q != seqlen_k and causal=True, the causal mask is aligned to the\nbottom right corner of the attention matrix, instead of the top-left corner.\n\nFor example, if seqlen_q = 2 and seqlen_k = 5, the causal mask (1 = keep, 0 =\nmasked out) is: \nv2.0: \n 1 0 0 0 0 \n 1 1 0 0 0 \nv2.1: \n 1 1 1 1 0 \n 1 1 1 1 1 \n\nIf seqlen_q = 5 and seqlen_k = 2, the causal mask is: \nv2.0: \n 1 0 \n 1 1 \n 1 1 \n 1 1 \n 1 1 \nv2.1: \n 0 0 \n 0 0 \n 0 0 \n 1 0 \n 1 1 \nIf the row of the mask is all zero, the output will be zero.\n\n### 2.2: Optimize for inference\n\nOptimize for inference (iterative decoding) when query has very small sequence\nlength (e.g., query sequence length = 1). The bottleneck here is to load KV\ncache as fast as possible, and we split the loading across different thread\nblocks, with a separate kernel to combine results.\n\nSee the function `flash_attn_with_kvcache` with more features for inference\n(perform rotary embedding, updating KV cache inplace).\n\nThanks to the xformers team, and in particular Daniel Haziza, for this\ncollaboration.\n\n### 2.3: Local (i.e., sliding window) attention\n\nImplement sliding window attention (i.e., local attention). Thanks to [Mistral\nAI](https://mistral.ai/) and in particular Timothée Lacroix for this\ncontribution. Sliding window was used in the [Mistral 7B](https://mistral.ai/news/announcing-mistral-7b/) model.\n\n### 2.4: ALiBi (attention with linear bias), deterministic backward pass.\n\nImplement ALiBi (Press et al., 2021). Thanks to Sanghun Cho from Kakao Brain for this contribution.\n\nImplement deterministic backward pass. Thanks to engineers from [Meituan](www.meituan.com) for this contribution.\n\n### 2.5: Paged KV cache.\n\nSupport paged KV cache (i.e., [PagedAttention](https://arxiv.org/abs/2309.06180)).\nThanks to @beginlner for this contribution.\n\n### 2.6: Softcapping.\n\nSupport attention with softcapping, as used in Gemma-2 and Grok models.\nThanks to @Narsil and @lucidrains for this contribution.\n\n### 2.7: Compatibility with torch compile\n\nThanks to @ani300 for this contribution.\n\n## Performance\n\nWe present expected speedup (combined forward + backward pass) and memory savings from using FlashAttention against PyTorch standard attention, depending on sequence length, on different GPUs (speedup depends on memory bandwidth - we see more speedup on slower GPU memory).\n\nWe currently have benchmarks for these GPUs:\n* [A100](#a100)\n* [H100](#h100)\n\n\n\n### A100\n\nWe display FlashAttention speedup using these parameters:\n* Head dimension 64 or 128, hidden dimension 2048 (i.e. either 32 or 16 heads).\n* Sequence length 512, 1k, 2k, 4k, 8k, 16k.\n* Batch size set to 16k / seqlen.\n\n#### Speedup\n\n\n\n#### Memory\n\n\n\nWe show memory savings in this graph (note that memory footprint is the same no matter if you use dropout or masking).\nMemory savings are proportional to sequence length -- since standard attention has memory quadratic in sequence length, whereas FlashAttention has memory linear in sequence length.\nWe see 10X memory savings at sequence length 2K, and 20X at 4K.\nAs a result, FlashAttention can scale to much longer sequence lengths.\n\n### H100\n\n\n\n## Full model code and training script\n\nWe have released the full GPT model\n[implementation](https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/models/gpt.py).\nWe also provide optimized implementations of other layers (e.g., MLP, LayerNorm,\ncross-entropy loss, rotary embedding). Overall this speeds up training by 3-5x\ncompared to the baseline implementation from Huggingface, reaching up to 225\nTFLOPs/sec per A100, equivalent to 72% model FLOPs utilization (we don't need\nany activation checkpointing).\n\nWe also include a training\n[script](https://github.com/Dao-AILab/flash-attention/tree/main/training) to\ntrain GPT2 on Openwebtext and GPT3 on The Pile.\n\n## Triton implementation of FlashAttention\n\nPhil Tillet (OpenAI) has an experimental implementation of FlashAttention in Triton:\nhttps://github.com/openai/triton/blob/master/python/tutorials/06-fused-attention.py\n\nAs Triton is a higher-level language than CUDA, it might be easier to understand\nand experiment with. The notations in the Triton implementation are also closer\nto what's used in our paper.\n\nWe also have an experimental implementation in Triton that support attention\nbias (e.g. ALiBi):\nhttps://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/flash_attn_triton.py\n\n\n## Tests\nWe test that FlashAttention produces the same output and gradient as a reference\nimplementation, up to some numerical tolerance. In particular, we check that the\nmaximum numerical error of FlashAttention is at most twice the numerical error\nof a baseline implementation in Pytorch (for different head dimensions, input\ndtype, sequence length, causal / non-causal).\n\nTo run the tests:\n```sh\npytest -q -s tests/test_flash_attn.py\n```\n## When you encounter issues\n\nThis new release of FlashAttention-2 has been tested on several GPT-style\nmodels, mostly on A100 GPUs.\n\nIf you encounter bugs, please open a GitHub Issue!\n\n## Tests\nTo run the tests:\n```sh\npytest tests/test_flash_attn_ck.py\n```\n\n## Citation\nIf you use this codebase, or otherwise found our work valuable, please cite:\n```\n@inproceedings{dao2022flashattention,\n title={Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},\n author={Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\\'e}, Christopher},\n booktitle={Advances in Neural Information Processing Systems (NeurIPS)},\n year={2022}\n}\n@inproceedings{dao2023flashattention2,\n title={Flash{A}ttention-2: Faster Attention with Better Parallelism and Work Partitioning},\n author={Dao, Tri},\n booktitle={International Conference on Learning Representations (ICLR)},\n year={2024}\n}\n```"},{"resultType":"repo_readme","repo":"alexzhang13/flashattention2-custom-mask","url":"https://github.com/alexzhang13/flashattention2-custom-mask","readmeUrl":"https://raw.githubusercontent.com/alexzhang13/flashattention2-custom-mask/HEAD/README.md","snippet":"# FlashAttention2 with Custom Masks 🎭\n**Note: This is an unofficial implementation of FlashAttention2.**\n\nFor efficiency purposes, the standard implementations of FlashAttention currently do not support **arbitrary custom masks**. \nTheir implementation of specific masks like causal masking for language modeling are implemented using branch logic to save memory. This repository is just a modified version of the tutorial Triton implementation of FlashAttention2 that allows the user\nto define a (batch of) custom mask. It modifies both the forward and backwards pass to handle custom masking (you can define a different mask per head and batch).\n \nOriginal Triton code: [https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html](https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html)\n\nSee the original thread: [https://github.com/Dao-AILab/flash-attention/issues/352](https://github.com/Dao-AILab/flash-attention/issues/352)\n\n## Quick Install\nCreate a Python environment (>=3.8) and install through pip:\n```\npip install flashattention2-custom-mask\n```\n\n## Example Setup\nThe relevant libraries needed to use the custom-mask FlashAttention2 kernel are below:\n```\npip install triton>=3.0.0\npip install torch\n```\n\n#### For Viewing Benchmarking Results\nOther libraries for evaluating the performance of the models is below. These are primarily for `test_benchmark.py`, which verifies the correctness of the implementation.\n```\npip install pytest\npip install matplotlib\npip install pandas\n```\nTo compare with the official FlashAttention and `xformers.ops.memory_efficient_attention` implementations, make sure to install both libraries separately (follow the instructions on these repositories).\n```\npip install flash-attn --no-build-isolation\npip3 install -U xformers --index-url https://download.pytorch.org/whl/cu121\n```\n\n## Testing Correctness\nThere are two `pytest` functions in `test_benchmark.py`, one that tests whether a reference implementation of multi-head attention with a causal mask matches the Triton version in both the forward pass and backwards pass gradients. The second tests whether the same implementation with **random** masks matches the Triton version. You can modify these tests to do more rigorous correctness tests and check with `pytest`.\n\n## Simple Example\nYou can insert this module into your standard attention pipeline.\n```python\nfrom fa2_custom_mask import flash_attention_custom_mask\n\nB, H, L, D = 4, 16, 4096, 64\nsm_scale = 1 / (D ** 0.5)\n\nfp32_q = torch.randn(B, H, L, D).float().cuda()\nfp32_k = torch.randn(B, H, L, D).float().cuda()\nfp32_v = torch.randn(B, H, L, D).float().cuda()\nmask = torch.randint(0, 2, (B, 1, L, L)).int().cuda()\nmask = torch.broadcast_to(mask, (B, H, L, L))\n\nout = flash_attention_custom_mask(fp32_q, fp32_k, fp32_v, mask=mask, sm_scale=sm_scale)\n...\nout.backward(loss)\n```\n\n## Benchmarking\nSimple benchmark against the base Triton implementation. In our custom mask version, we pass in the canonical causal mask as input (hence storing in global device memory). Running `test_benchmark.py`,\nwith batch size=4, # heads=16, hidden dim=64, and sequence length `N_CTX` ranging from 256 to 16384 in powers of 2. You can replicate the experiments by running\n```\npytest\npython test_benchmark.py\n```\n\n#### Causal Masks and No Masks Comparisons \nWe compare against the original experiments and original implementation, as well as the official FlashAttention and xformers implementation (note: there seems to be a versioning issue, so it's using a different implementation. I corrected the version in the later benchmarking experiments). \n\n \n#### Causal Masks and No Masks Comparisons (with Correct xfrormers version)\nWe compare against the original experiments and original implementation, as well as the xformers implementation. Notably, the original implementation does well for causal masking because of some pipelining tricks and ability to not have to store masks.\n\n#### Custom Masking Comparison\nWe compare directly to the [xformers memory efficient attention](https://facebookresearch.github.io/xformers/components/ops.html) which allows for custom masking. We generate random masks (fixed across the head dimension).\n\n\n\n## Notes and Bugs\n1. This implementation only works on Ampere devices and up. I originally tried running it on a V100 (Volta) and it failed. \n2. You need to be on `triton>=3.0.0`, or it'll complain about permutation indices on the value vector pointer. The `torch` and `flash-attn` libraries may force you to install `triton=2.x.x`, but you can just re-install `triton>=3.0.0` and it should work. I may fix this manually in the future.\n * This is oddly specific, but I'm not able to have `flash-attn` and `xformers` at the same time. I had to run them separately and generate the plots.\n3. TODO: Add benchmarking for peak memory consumption and other efficiency metrics.\n\nIf time permits, I'm interested in making this implementation generalizable / changing the CUDA implementation for FA3 (if it's necessary of course). I also probably will run some more realistic workloads and see what happens."},{"resultType":"github_history","repo":"keras-team/keras","url":"https://github.com/keras-team/keras/issues/19418","pageType":"issue","number":19418,"segmentCount":1,"title":"Add Flash Attention\\#19418","snippet":"# Add Flash Attention\\#19418\n\nAssignees\n\n## Description\n\n[innat](https://github.com/innat)\n\n**Describe**\n\nFlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.\n\nPaper [https://arxiv.org/abs/2205.14135](https://arxiv.org/abs/2205.14135)\n\nCited by: 671\n\n**Implementation**\n\n- PyTorch: [https://github.com/Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention)\n- Jax: [https://github.com/lucidrains/flash-attention-jax](https://github.com/lucidrains/flash-attention-jax)\n- TensorFlow (with custom ops): [https://github.com/intelligent-machine-learning/dlrover/tree/master/tfplus/tfplus/flash\\_attn](https://github.com/intelligent-machine-learning/dlrover/tree/master/tfplus/tfplus/flash_attn)\n\nHuggingface [https://huggingface.co/docs/text-generation-inference/en/conceptual/flash\\_attention](https://huggingface.co/docs/text-generation-inference/en/conceptual/flash_attention)\n\n**Others**\n\nHas version 2 of it.\n\n[https://arxiv.org/abs/2307.08691](https://arxiv.org/abs/2307.08691)\n\nFlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning\n\n## Activity\n\nadded\n\n[on Apr 1, 2024on Apr 1, 2024](https://github.com/keras-team/keras/issues/19418#event-12315688029)\n\nassigned\n\n[SuryanarayanaY](https://github.com/SuryanarayanaY)\n\n[on Apr 1, 2024on Apr 1, 2024](https://github.com/keras-team/keras/issues/19418#event-12315689496)\n\n### fchollet commented on Apr 2, 2024on Apr 2, 2024\n\n[fchollet](https://github.com/fchollet)\n\n[on Apr 2, 2024on Apr 2, 2024](https://github.com/keras-team/keras/issues/19418#issuecomment-2031040424)\n\nLast edited by fchollet\n\nFor JAX, we may want to rely on Pallas. For TF, since we can't rely on custom ops, we may have to skip support.\n\nPresumably we should add it in the form of a new backend op, `ops.nn.flash_attention`.\n\n👍React with 👍2innat and QuantHao\n\nadded\n\n[on Apr 2, 2024on Apr 2, 2024](https://github.com/keras-team/keras/issues/19418#event-12320898187)\n\nmentioned this [on Apr 6, 2024on Apr 6, 2024](https://github.com/keras-team/keras/issues/19418#event-1375581833)\n\n- [🚀 Contributing to Keras 🚀 #18442](https://github.com/keras-team/keras/issues/18442)\n\nmentioned this [on Nov 28, 2024on Nov 28, 2024](https://github.com/keras-team/keras/issues/19418#event-1751964412)\n\n- [How to enable Flash-Attn in the PyTorch backend. #20556](https://github.com/keras-team/keras/issues/20556)\n\n### cpuimage commented on May 29, 2025on May 29, 2025\n\n[cpuimage](https://github.com/cpuimage)\n\n[on May 29, 2025on May 29, 2025](https://github.com/keras-team/keras/issues/19418#issuecomment-2919730207)\n\nHere's my implementation of chunked flash attention for Keras: [chunked-flash-attention-keras](https://github.com/cpuimage/chunked-flash-attention-keras). Feedback welcome!\n\nadded a commit that references this issue [3w agoon Jun 25, 2026](https://github.com/keras-team/keras/issues/19418#event-27225769288)\n\n[feat(css): allow scoping css to importers exports (#19418)\\\n\\\nCo-authored-by: bluwy - A PR should do only one thing, smaller PRs enable faster reviews.
- Every PR should include unit tests and end-to-end tests to ensure it works and is not broken by other future PRs.
- Write the commit message by fulfilling the PR description to help reviewer and future developers understand.
If CI fails, you can run linting and testing checks locally according [Contributing](https://docs.vllm.ai/projects/ascend/zh-cn/latest/developer_guide/contribution/index.html) and [Testing](https://docs.vllm.ai/projects/ascend/zh-cn/latest/developer_guide/contribution/testing.html). |\n\n**[gemini-code-assist](https://github.com/apps/gemini-code-assist) Bot**\n\nreviewed\n\n[last monthJun 8, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#pullrequestreview-4447658606)\n\n[View reviewed changes](https://github.com/vllm-project/vllm-ascend/pull/10181/files)\n\n### **[gemini-code-assist](https://github.com/apps/gemini-code-assist) Bot** left a comment\n\nThe reason will be displayed to describe this comment to others. [Learn more](https://docs.github.com/articles/managing-disruptive-comments/#hiding-a-comment).\n\n## Code Review\n\n**Suggested PR Title:**\n\n```\n[Attention][BugFix] Use torch_npu._npu_flash_attention_v3 with compressed mask on Ascend 310P\n```\n\n**Suggested PR Summary:**\n\n```\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use `torch_npu._npu_flash_attention_v3` instead of the older `_npu_flash_attention`. It introduces a fixed sequence length of 2048 (`FA_V3_COMPRESSED_MASK_SEQ_LEN`) for the 310P compressed-mask path and configures the flash attention call with `mask_type=MASK_TYPE_NORM_COMPRESS_SELF_ATTENTION` (value 3).\n\nHowever, the current implementation calls `_npu_flash_attention_v3` using positional arguments, which is fragile and incorrect due to parameter mismatch on actual NPU hardware (e.g., the 4th positional argument is `pse_shift`, not `atten_mask`). We recommend updating the call to use keyword arguments and updating the corresponding unit tests to verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/` that mock `_npu_flash_attention_v3` and verify the arguments passed during prefill and encoder attention.\n```\n\nComment thread[vllm\\_ascend/\\_310p/attention/attention\\_v1.py](https://github.com/vllm-project/vllm-ascend/pull/10181/files#diff-84a9cbcfe0590214c45c32c3c28ccbea2aab6e6bd6e9e4df39e31605fdc48fd8)Show resolvedHide resolved\n\nComment thread[tests/ut/\\_310p/attention/test\\_attention\\_v1\\_310.py](https://github.com/vllm-project/vllm-ascend/pull/10181/files#diff-998eaafe5ae4cebce27b44aff4ec43a89d19d6e8ce69ddca7996762cc0c43b2c)\nOutdated\nShow resolvedHide resolved\n\nComment thread[tests/ut/\\_310p/attention/test\\_attention\\_v1\\_310.py](https://github.com/vllm-project/vllm-ascend/pull/10181/files#diff-998eaafe5ae4cebce27b44aff4ec43a89d19d6e8ce69ddca7996762cc0c43b2c)\nOutdated\nShow resolvedHide resolved\n\nrequested review from\n[LCAIZJ](https://github.com/LCAIZJ) and\n[Yikun](https://github.com/Yikun)\n\nas [code owners](https://github.com/vllm-project/vllm-ascend/blob/4f86c9a864ea2600a5ac2de587eda66fffba16fc/.github/CODEOWNERS#L23) [last monthJune 8, 2026 13:25](https://github.com/vllm-project/vllm-ascend/pull/10181#event-26472552638)\n\nadded\nthe\nlabel\n\n[last monthJun 8, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#event-26472622185)\n\n`\n [Feature][310p] support compressed-mask for PA-split-fuse op\n` …\n\n`\n 59bf202\n`\n\n```\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\n```\n\nthe\n6-8-FA-compressed-mask\nbranch\nfrom\n[`cdff5c9`](https://github.com/vllm-project/vllm-ascend/commit/cdff5c980e52513902dfba5c52cff70b8686bed6) to\n[`59bf202`](https://github.com/vllm-project/vllm-ascend/commit/59bf20207c0a3caa45b39422a1093b72ec3af5d0) [Compare](https://github.com/vllm-project/vllm-ascend/compare/cdff5c980e52513902dfba5c52cff70b8686bed6..59bf20207c0a3caa45b39422a1093b72ec3af5d0) [last monthJune 9, 2026 02:49](https://github.com/vllm-project/vllm-ascend/pull/10181#event-26501264051)\n\nremoved\nthe\nlabel\n\n[last monthJun 9, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#event-26501272975)\n\nadded\nthe\nlabel\n\n[last monthJun 9, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#event-26502232794)\n\n`\n Merge branch 'main' into 6-8-FA-compressed-mask\n`\n\non Jun 8, 2026, 11:28 PM\n\n`\n d7a5eb9\n`\n\nchanged the title\n~~\\[Feature\\]\\[310p\\] support compressed mask for 310P flashattention~~\\[Feature\\]\\[310p\\] support compressed mask for 310P flashattention and pagedattention-split-fuse [last monthJun 9, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#event-26502599223)\n\nadded\nthe\nlabel\n\n[last monthJun 9, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#event-26503086035)\n\n**[ZT-AIA](https://github.com/ZT-AIA)**\n\napproved these changes\n\n[last monthJun 9, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#pullrequestreview-4456882285)\n\n[View reviewed changes](https://github.com/vllm-project/vllm-ascend/pull/10181/files/d7a5eb945ecf4efe812998891e5c065f61083792)\n\nHide detailsView details[](https://github.com/wangxiyuan)[wangxiyuan](https://github.com/wangxiyuan)\n\nmerged commit [`fc0b9e3`](https://github.com/vllm-project/vllm-ascend/commit/fc0b9e3540bc85dfe355ebb244c80965ab39a2b2)\ninto\n\nvllm-project:main[last monthJun 9, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#event-26515231988)\n\n16 checks passed\n\n[Zhutianyi7230](https://github.com/Zhutianyi7230)\n\npushed a commit\nto Zhutianyi7230/vllm-ascend-zty\nthat referenced\nthis pull request\n\n[3 weeks agoJun 12, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#ref-commit-2b86de4)\n\n`\n [Feature][310p] support compressed mask for 310P flashattention and p…\n`…\n\n`\n 2b86de4\n`\n\n```\n…agedattention-split-fuse (vllm-project#10181)\n\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use\n`torch_npu._npu_flash_attention_v3` instead of the older\n`_npu_flash_attention` and `torch_npu._npu_paged_attention_splitfuse_v2`\ninstead of the older `torch_npu._npu_paged_attention_splitfuse` . It\nintroduces a fixed sequence length of 2048 for the 310P compressed-mask\npath and configures the flash attention call with `mask_type`.\n\nHowever, the current implementation calls `_npu_flash_attention_v3`\nusing positional arguments, which is fragile and incorrect due to\nparameter mismatch on actual NPU hardware (e.g., the 4th positional\nargument is `pse_shift`, not `atten_mask`). We recommend updating the\ncall to use keyword arguments and updating the corresponding unit tests\nto verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/`\nthat mock `_npu_flash_attention_v3` and verify the arguments passed\nduring prefill and encoder attention.\n\n- vLLM version: v0.21.0\n- vLLM main:\nvllm-project/vllm@9090368\n\n---------\n\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\nSigned-off-by: zhutianyi <1589841300@qq.com>\n```\n\n[LostFox11](https://github.com/LostFox11)\n\npushed a commit\nto LostFox11/vllm-ascend\nthat referenced\nthis pull request\n\n[3 weeks agoJun 15, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#ref-commit-ada1de5)\n\n`\n [Feature][310p] support compressed mask for 310P flashattention and p…\n`…\n\n`\n ada1de5\n`\n\n```\n…agedattention-split-fuse (vllm-project#10181)\n\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use\n`torch_npu._npu_flash_attention_v3` instead of the older\n`_npu_flash_attention` and `torch_npu._npu_paged_attention_splitfuse_v2`\ninstead of the older `torch_npu._npu_paged_attention_splitfuse` . It\nintroduces a fixed sequence length of 2048 for the 310P compressed-mask\npath and configures the flash attention call with `mask_type`.\n\nHowever, the current implementation calls `_npu_flash_attention_v3`\nusing positional arguments, which is fragile and incorrect due to\nparameter mismatch on actual NPU hardware (e.g., the 4th positional\nargument is `pse_shift`, not `atten_mask`). We recommend updating the\ncall to use keyword arguments and updating the corresponding unit tests\nto verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/`\nthat mock `_npu_flash_attention_v3` and verify the arguments passed\nduring prefill and encoder attention.\n\n- vLLM version: v0.21.0\n- vLLM main:\nvllm-project/vllm@9090368\n\n---------\n\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\n```\n\n[LostFox11](https://github.com/LostFox11)\n\npushed a commit\nto LostFox11/vllm-ascend\nthat referenced\nthis pull request\n\n[3 weeks agoJun 15, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#ref-commit-2e98639)\n\n`\n [Feature][310p] support compressed mask for 310P flashattention and p…\n`…\n\n`\n 2e98639\n`\n\n```\n…agedattention-split-fuse (vllm-project#10181)\n\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use\n`torch_npu._npu_flash_attention_v3` instead of the older\n`_npu_flash_attention` and `torch_npu._npu_paged_attention_splitfuse_v2`\ninstead of the older `torch_npu._npu_paged_attention_splitfuse` . It\nintroduces a fixed sequence length of 2048 for the 310P compressed-mask\npath and configures the flash attention call with `mask_type`.\n\nHowever, the current implementation calls `_npu_flash_attention_v3`\nusing positional arguments, which is fragile and incorrect due to\nparameter mismatch on actual NPU hardware (e.g., the 4th positional\nargument is `pse_shift`, not `atten_mask`). We recommend updating the\ncall to use keyword arguments and updating the corresponding unit tests\nto verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/`\nthat mock `_npu_flash_attention_v3` and verify the arguments passed\nduring prefill and encoder attention.\n\n- vLLM version: v0.21.0\n- vLLM main:\nvllm-project/vllm@9090368\n\n---------\n\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\n```\n\n[ader47](https://github.com/ader47)\n\npushed a commit\nto ader47/vllm-ascend\nthat referenced\nthis pull request\n\n[2 weeks agoJun 18, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#ref-commit-101e5c7)\n\n`\n [Feature][310p] support compressed mask for 310P flashattention and p…\n`…\n\n`\n 101e5c7\n`\n\n```\n…agedattention-split-fuse (vllm-project#10181)\n\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use\n`torch_npu._npu_flash_attention_v3` instead of the older\n`_npu_flash_attention` and `torch_npu._npu_paged_attention_splitfuse_v2`\ninstead of the older `torch_npu._npu_paged_attention_splitfuse` . It\nintroduces a fixed sequence length of 2048 for the 310P compressed-mask\npath and configures the flash attention call with `mask_type`.\n\nHowever, the current implementation calls `_npu_flash_attention_v3`\nusing positional arguments, which is fragile and incorrect due to\nparameter mismatch on actual NPU hardware (e.g., the 4th positional\nargument is `pse_shift`, not `atten_mask`). We recommend updating the\ncall to use keyword arguments and updating the corresponding unit tests\nto verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/`\nthat mock `_npu_flash_attention_v3` and verify the arguments passed\nduring prefill and encoder attention.\n\n- vLLM version: v0.21.0\n- vLLM main:\nvllm-project/vllm@9090368\n\n---------\n\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\n```\n\n[pisceskkk](https://github.com/pisceskkk)\n\npushed a commit\nto pisceskkk/vllm-ascend\nthat referenced\nthis pull request\n\n[2 weeks agoJun 23, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#ref-commit-35930d9)\n\n`\n [Feature][310p] support compressed mask for 310P flashattention and p…\n`…\n\n`\n 35930d9\n`\n\n```\n…agedattention-split-fuse (vllm-project#10181)\n\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use\n`torch_npu._npu_flash_attention_v3` instead of the older\n`_npu_flash_attention` and `torch_npu._npu_paged_attention_splitfuse_v2`\ninstead of the older `torch_npu._npu_paged_attention_splitfuse` . It\nintroduces a fixed sequence length of 2048 for the 310P compressed-mask\npath and configures the flash attention call with `mask_type`.\n\nHowever, the current implementation calls `_npu_flash_attention_v3`\nusing positional arguments, which is fragile and incorrect due to\nparameter mismatch on actual NPU hardware (e.g., the 4th positional\nargument is `pse_shift`, not `atten_mask`). We recommend updating the\ncall to use keyword arguments and updating the corresponding unit tests\nto verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/`\nthat mock `_npu_flash_attention_v3` and verify the arguments passed\nduring prefill and encoder attention.\n\n- vLLM version: v0.21.0\n- vLLM main:\nvllm-project/vllm@9090368\n\n---------\n\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\n```\n\n[CXY-Katrina](https://github.com/CXY-Katrina)\n\npushed a commit\nto CXY-Katrina/vllm-ascend-zhx\nthat referenced\nthis pull request\n\n[last weekJun 27, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#ref-commit-056e97c)\n\n`\n [Feature][310p] support compressed mask for 310P flashattention and p…\n`…\n\n`\n 056e97c\n`\n\n```\n…agedattention-split-fuse (vllm-project#10181)\n\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use\n`torch_npu._npu_flash_attention_v3` instead of the older\n`_npu_flash_attention` and `torch_npu._npu_paged_attention_splitfuse_v2`\ninstead of the older `torch_npu._npu_paged_attention_splitfuse` . It\nintroduces a fixed sequence length of 2048 for the 310P compressed-mask\npath and configures the flash attention call with `mask_type`.\n\nHowever, the current implementation calls `_npu_flash_attention_v3`\nusing positional arguments, which is fragile and incorrect due to\nparameter mismatch on actual NPU hardware (e.g., the 4th positional\nargument is `pse_shift`, not `atten_mask`). We recommend updating the\ncall to use keyword arguments and updating the corresponding unit tests\nto verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/`\nthat mock `_npu_flash_attention_v3` and verify the arguments passed\nduring prefill and encoder attention.\n\n- vLLM version: v0.21.0\n- vLLM main:\nvllm-project/vllm@9090368\n\n---------\n\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\n```\n\n[HaoxinZong](https://github.com/HaoxinZong)\n\npushed a commit\nto HaoxinZong/vllm-ascend\nthat referenced\nthis pull request\n\n[last weekJun 27, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#ref-commit-d4be74b)\n\n`\n [Feature][310p] support compressed mask for 310P flashattention and p…\n`…\n\n`\n d4be74b\n`\n\n```\n…agedattention-split-fuse (vllm-project#10181)\n\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use\n`torch_npu._npu_flash_attention_v3` instead of the older\n`_npu_flash_attention` and `torch_npu._npu_paged_attention_splitfuse_v2`\ninstead of the older `torch_npu._npu_paged_attention_splitfuse` . It\nintroduces a fixed sequence length of 2048 for the 310P compressed-mask\npath and configures the flash attention call with `mask_type`.\n\nHowever, the current implementation calls `_npu_flash_attention_v3`\nusing positional arguments, which is fragile and incorrect due to\nparameter mismatch on actual NPU hardware (e.g., the 4th positional\nargument is `pse_shift`, not `atten_mask`). We recommend updating the\ncall to use keyword arguments and updating the corresponding unit tests\nto verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/`\nthat mock `_npu_flash_attention_v3` and verify the arguments passed\nduring prefill and encoder attention.\n\n- vLLM version: v0.21.0\n- vLLM main:\nvllm-project/vllm@9090368\n\n---------\n\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\n```\n\n[pisceskkk](https://github.com/pisceskkk)\n\npushed a commit\nto pisceskkk/vllm-ascend\nthat referenced\nthis pull request\n\n[last weekJun 29, 2026](https://github.com/vllm-project/vllm-ascend/pull/10181#ref-commit-a37999e)\n\n`\n [Feature][310p] support compressed mask for 310P flashattention and p…\n`…\n\n`\n a37999e\n`\n\n```\n…agedattention-split-fuse (vllm-project#10181)\n\n### What this PR does / why we need it?\n\nThis PR updates the Ascend 310P attention implementation to use\n`torch_npu._npu_flash_attention_v3` instead of the older\n`_npu_flash_attention` and `torch_npu._npu_paged_attention_splitfuse_v2`\ninstead of the older `torch_npu._npu_paged_attention_splitfuse` . It\nintroduces a fixed sequence length of 2048 for the 310P compressed-mask\npath and configures the flash attention call with `mask_type`.\n\nHowever, the current implementation calls `_npu_flash_attention_v3`\nusing positional arguments, which is fragile and incorrect due to\nparameter mismatch on actual NPU hardware (e.g., the 4th positional\nargument is `pse_shift`, not `atten_mask`). We recommend updating the\ncall to use keyword arguments and updating the corresponding unit tests\nto verify these keyword arguments.\n\n### Does this PR introduce _any_ user-facing change?\n\nNo. This is an internal backend update for Ascend 310P hardware support.\n\n### How was this patch tested?\n\nThe patch includes updated unit tests in `tests/ut/_310p/attention/`\nthat mock `_npu_flash_attention_v3` and verify the arguments passed\nduring prefill and encoder attention.\n\n- vLLM version: v0.21.0\n- vLLM main:\nvllm-project/vllm@9090368\n\n---------\n\nSigned-off-by: Tflowers-0129 <2906339855@qq.com>\n```\n\nThis file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.\n[Learn more about bidirectional Unicode characters](https://github.co/hiddenchars)\n\n[Show hidden characters](https://github.com/vllm-project/vllm-ascend/pull/10181)"}]}