Skip to content

V3 gpu flash attention triton solution fix #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Exorust merged 2 commits into from
Sep 8, 2026
Merged

V3 gpu flash attention triton solution fix #29

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
5612d9f
Add flash_attention_triton wrapper that passes checks to the flash-at…
jeniag Sep 7, 2026
37ee00b
Add flash_attention_triton wrapper that passes checks to the flash-at…
jeniag Sep 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
256 changes: 187 additions & 69 deletions v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
"cells": [
{
"cell_type": "markdown",
"id": "a14a63c5",
"metadata": {},
"source": [
"# Implement FlashAttention-2 in Triton — Solution\n",
Expand Down Expand Up @@ -79,9 +80,19 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 2,
"id": "e1ce38d2",
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Triton is available!\n",
"CUDA available: True\n"
]
}
],
"source": [
"import torch\n",
"import torch.nn.functional as F\n",
Expand All @@ -102,9 +113,23 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 3,
"id": "cdf6b637",
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Q shape: torch.Size([2, 4, 128, 64])\n",
"K shape: torch.Size([2, 4, 128, 64])\n",
"V shape: torch.Size([2, 4, 128, 64])\n",
"\n",
"Full attention matrix would be: 2 x 4 x 128 x 128\n",
"= 512.0 KB\n"
]
}
],
"source": [
"# Test data\n",
"torch.manual_seed(42)\n",
Expand All @@ -129,7 +154,8 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 4,
"id": "bdb3477a",
"metadata": {},
"outputs": [],
"source": [
Expand All @@ -152,7 +178,8 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 5,
"id": "072c33a3",
"metadata": {},
"outputs": [],
"source": [
Expand Down Expand Up @@ -232,82 +259,121 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 6,
"id": "422ab2cd",
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Triton flash attention kernel defined.\n"
]
}
],
"source": [
"# Triton kernel reference (requires GPU to actually run)\n",
"# This shows what the GPU kernel would look like\n",
"\n",
"if TRITON_AVAILABLE:\n",
" DEVICE = triton.runtime.driver.active.get_active_torch_device()\n",
" \n",
" @triton.jit\n",
" def flash_attention_kernel(\n",
" Q_ptr, K_ptr, V_ptr, O_ptr,\n",
" stride_qb, stride_qh, stride_qn, stride_qd,\n",
" stride_kb, stride_kh, stride_kn, stride_kd,\n",
" stride_vb, stride_vh, stride_vn, stride_vd,\n",
" stride_ob, stride_oh, stride_on, stride_od,\n",
" N, D: tl.constexpr,\n",
" BLOCK_Q: tl.constexpr, BLOCK_KV: tl.constexpr,\n",
" ):\n",
" \"\"\"\n",
" FlashAttention-2 Triton kernel.\n",
" Each program processes one (batch, head, q_block) tile.\n",
" \"\"\"\n",
" # Program IDs\n",
" Q_block_ptr,\n",
" K_block_ptr,\n",
" V_block_ptr,\n",
" output_ptr,\n",
" Q_strideBH, Q_strideN, Q_strideD,\n",
" K_strideBH, K_strideN, K_strideD,\n",
" V_strideBH, V_strideN, V_strideD,\n",
" output_strideBH, output_strideN, output_strideD,\n",
" N:tl.constexpr, \n",
" D:tl.constexpr, \n",
" BLOCK_Q: tl.constexpr,\n",
" BLOCK_KV: tl.constexpr,\n",
" ): \n",
" scale = 1.0 / tl.sqrt(float(D)) \n",
" n_blocks = tl.cdiv(N , BLOCK_KV) \n",
" running_max = tl.full((BLOCK_Q,1), float('-inf'), dtype=tl.float32)\n",
" running_sum = tl.zeros((BLOCK_Q,1),dtype=tl.float32)\n",
" running_output = tl.zeros((BLOCK_Q,D),dtype=tl.float32)\n",
"\n",
" q_block_idx = tl.program_id(0)\n",
" bh_idx = tl.program_id(1) # combined batch*head index\n",
" \n",
" scale = 1.0 / tl.sqrt(float(D))\n",
" \n",
" # Offsets for this Q block\n",
" q_offset = q_block_idx * BLOCK_Q\n",
" q_range = q_offset + tl.arange(0, BLOCK_Q)\n",
" bh_idx = tl.program_id(1)\n",
" \n",
" q_start = Q_block_ptr + Q_strideBH * bh_idx\n",
" q_range = q_block_idx * BLOCK_Q + tl.arange(0, BLOCK_Q)\n",
" d_range = tl.arange(0, D)\n",
" q_mask = q_range[:, None] < N\n",
" \n",
" # Load Q block into SRAM\n",
" q_ptrs = Q_ptr + bh_idx * stride_qh + q_range[:, None] * stride_qn + d_range[None, :] * stride_qd\n",
" Q_block = tl.load(q_ptrs, mask=q_mask, other=0.0)\n",
" q_offsets = q_range[:,None]*Q_strideN+d_range[None,:]*Q_strideD\n",
" q_mask = q_range[:,None]<N \n",
" \n",
" # Initialize accumulators\n",
" running_max = tl.full([BLOCK_Q, 1], float('-inf'), dtype=tl.float32)\n",
" running_sum = tl.zeros([BLOCK_Q, 1], dtype=tl.float32)\n",
" running_out = tl.zeros([BLOCK_Q, D], dtype=tl.float32)\n",
" Q_block = tl.load(q_start+ q_offsets, mask=q_mask, other=0.0) \n",
" \n",
" # Loop over K,V blocks\n",
" n_kv_blocks = tl.cdiv(N, BLOCK_KV)\n",
" for kv_idx in range(n_kv_blocks):\n",
" kv_offset = kv_idx * BLOCK_KV\n",
" kv_range = kv_offset + tl.arange(0, BLOCK_KV)\n",
" kv_mask = kv_range[None, :] < N\n",
" \n",
" # Load K block\n",
" k_ptrs = K_ptr + bh_idx * stride_kh + kv_range[None, :] * stride_kn + d_range[:, None] * stride_kd\n",
" K_block = tl.load(k_ptrs, mask=kv_mask, other=0.0) # (D, BLOCK_KV)\n",
" for kv_block_idx in range(n_blocks):\n",
" \n",
" k_bh_start = K_block_ptr + K_strideBH * bh_idx\n",
" v_bh_start = V_block_ptr + V_strideBH * bh_idx \n",
" kv_range = kv_block_idx * BLOCK_KV + tl.arange(0, BLOCK_KV)\n",
" d_range = tl.arange(0, D)\n",
" k_offsets = kv_range[:,None]*K_strideN+d_range[None,:]*K_strideD\n",
" v_offsets = kv_range[:,None]*V_strideN+d_range[None,:]*V_strideD\n",
" k_mask = kv_range[:,None]<N \n",
" v_mask = kv_range[:,None]<N \n",
" \n",
" # Compute scores: Q @ K^T\n",
" S = tl.dot(Q_block, K_block) * scale # (BLOCK_Q, BLOCK_KV)\n",
" K_block = tl.load(k_bh_start+ k_offsets, mask=k_mask, other=0.0)\n",
" V_block = tl.load(v_bh_start+ v_offsets, mask=v_mask, other=0.0) \n",
" \n",
" S = tl.dot(Q_block, K_block.trans(1,0)) * scale \n",
" \n",
" # Online softmax\n",
" block_max = tl.max(S, axis=1)[:, None]\n",
" block_max = S.max(axis=1, keep_dims=True)\n",
" new_max = tl.maximum(running_max, block_max)\n",
" correction = tl.exp(running_max - new_max)\n",
" P = tl.exp(S - new_max)\n",
" running_sum = running_sum * correction + tl.sum(P, axis=1)[:, None]\n",
" \n",
" # Load V block and accumulate\n",
" v_ptrs = V_ptr + bh_idx * stride_vh + kv_range[:, None] * stride_vn + d_range[None, :] * stride_vd\n",
" v_mask = kv_range[:, None] < N\n",
" V_block = tl.load(v_ptrs, mask=v_mask, other=0.0)\n",
" running_out = running_out * correction + tl.dot(P.to(V_block.dtype), V_block)\n",
" running_sum = running_sum * correction + P.sum(axis=1, keep_dims=True) \n",
" running_output = running_output * correction + tl.dot(P.to(V_block.dtype), V_block) \n",
" running_max = new_max\n",
" \n",
" # Normalize and store\n",
" result = running_out / running_sum\n",
" o_ptrs = O_ptr + bh_idx * stride_oh + q_range[:, None] * stride_on + d_range[None, :] * stride_od\n",
" tl.store(o_ptrs, result, mask=q_mask)\n",
" \n",
" result = running_output / running_sum \n",
" output_start = output_ptr + output_strideBH * bh_idx \n",
" output_offsets = q_range[:,None]*output_strideN+d_range[None,:]*output_strideD\n",
" \n",
" tl.store(output_start + output_offsets, result, mask=q_mask) \n",
" \n",
" def flash_attention_triton(\n",
" Q, K, V, block_size=32, use_sol=False\n",
" ):\n",
"\n",
" orig_device = Q.device\n",
" Q= Q.to(DEVICE)\n",
" K= K.to(DEVICE)\n",
" V= V.to(DEVICE)\n",
" \n",
" B, H, N, D = Q.shape \n",
" \n",
" # Output accumulator\n",
" output = torch.zeros_like(Q)\n",
" \n",
" # Number of blocks\n",
" n_q_blocks = math.ceil(N / block_size) \n",
" \n",
" grid = (n_q_blocks,B*H)\n",
" print(f\"{B=},{H=},{N=},{D=}, {grid=} {(output.stride(1), output.stride(2), output.stride(3))=}\")\n",
" krnl = flash_attention_kernel_from_sol if use_sol else flash_attention_kernel\n",
" krnl[grid](\n",
" Q,\n",
" K,\n",
" V,\n",
" output,\n",
" Q.stride(1), Q.stride(2), Q.stride(3),\n",
" K.stride(1), K.stride(2), K.stride(3),\n",
" V.stride(1), V.stride(2), V.stride(3),\n",
" output.stride(1), output.stride(2), output.stride(3),\n",
" N=int(N), \n",
" D=int(D), \n",
" BLOCK_Q=int(block_size),\n",
" BLOCK_KV=int(block_size),\n",
" ) \n",
" return output.to(orig_device)\n",
" print(\"Triton flash attention kernel defined.\")\n",
"else:\n",
" print(\"Triton not available. The kernel code above shows the GPU implementation.\")\n",
Expand All @@ -316,9 +382,45 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 7,
"id": "45628148",
"metadata": {},
"outputs": [],
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"============================================================\n",
"VALIDATING FLASH ATTENTION\n",
"============================================================\n",
"\n",
"--- Correctness Test ---\n",
"Max absolute error: 3.58e-07\n",
"PASSED\n",
"\n",
"--- Block Size Robustness Test ---\n",
" block_size=16: max_err=3.87e-07 PASSED\n",
" block_size=32: max_err=3.58e-07 PASSED\n",
" block_size=64: max_err=3.28e-07 PASSED\n",
"\n",
"--- Non-Divisible Sequence Length Test ---\n",
" seq_len=50, block_size=16: max_err=2.98e-07 PASSED\n",
"\n",
"--- Memory Analysis ---\n",
" Standard attention peak: 128x128 = 16384 elements\n",
" Flash attention peak: 32x32 = 1024 elements per tile\n",
" Memory reduction: 16.0x\n",
"\n",
"--- Scaling Analysis ---\n",
" N= 256: standard= 0.2 MB, flash_tile=4.0 KB, ratio=64x\n",
" N= 1024: standard= 4.0 MB, flash_tile=4.0 KB, ratio=1024x\n",
" N= 4096: standard= 64.0 MB, flash_tile=4.0 KB, ratio=16384x\n",
" N= 8192: standard= 256.0 MB, flash_tile=4.0 KB, ratio=65536x\n",
"\n",
"All tests passed!\n"
]
}
],
"source": [
"# Validation\n",
"print(\"=\" * 60)\n",
Expand Down Expand Up @@ -377,19 +479,35 @@
"\n",
"print(\"\\nAll tests passed!\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "59e24d81-2de9-4f9b-8896-cb7095ea8604",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"version": "3.10.0"
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.3"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
}