From 5612d9f80ef693c997877d02f271cb089f2a823e Mon Sep 17 00:00:00 2001
From: "test@example.com" <26772842+jeniag@users.noreply.github.com>
Date: Mon, 7 Sep 2026 11:57:06 -0700
Subject: [PATCH 1/2] Add flash_attention_triton wrapper that passes checks to
 the flash-attention-triton question solution notebook

---
 .../flash-attention-triton_SOLN.ipynb         | 253 +++++++++++++-----
 1 file changed, 184 insertions(+), 69 deletions(-)

diff --git a/v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb b/v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb
index ba0e100..054fc39 100644
--- a/v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb
+++ b/v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb
@@ -2,6 +2,7 @@
  "cells": [
   {
    "cell_type": "markdown",
+   "id": "a14a63c5",
    "metadata": {},
    "source": [
     "# Implement FlashAttention-2 in Triton — Solution\n",
@@ -79,9 +80,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
+   "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",
@@ -102,9 +113,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
+   "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",
@@ -129,7 +154,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
+   "id": "bdb3477a",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -152,7 +178,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
+   "id": "072c33a3",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -232,82 +259,126 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 5,
+   "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",
+    "        q_offsets = q_range[:,None]*Q_strideN+d_range[None,:]*Q_strideD\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_block = tl.load(q_start+ q_offsets, mask=q_mask, other=0.0)        \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",
-    "        \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",
+    "            #tl.static_print(str((running_output.shape, correction.shape, P.shape, V_block.shape)))\n",
+    "            running_output = running_output * correction  + tl.dot(P.to(V_block.dtype), V_block)   \n",
+    "            # if bh_idx==7 and q_block_idx==3 and kv_block_idx==1:\n",
+    "            #     tl.device_print(\"running_sum\",running_sum.sum())\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",
+    "        # if bh_idx==7 and q_block_idx==3:\n",
+    "        #     tl.device_print(\"result\",result.sum(axis=None,keep_dims=False))\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",
@@ -316,9 +387,45 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
+   "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",
@@ -381,15 +488,23 @@
  ],
  "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
-}
\ No newline at end of file
+}

From 37ee00be0a5052893fbc6f8dd54c6401b49c1507 Mon Sep 17 00:00:00 2001
From: "test@example.com" <26772842+jeniag@users.noreply.github.com>
Date: Mon, 7 Sep 2026 12:01:03 -0700
Subject: [PATCH 2/2] Add flash_attention_triton wrapper that passes checks to
 the flash-attention-triton question solution notebook

---
 .../flash-attention-triton_SOLN.ipynb         | 31 ++++++++++---------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb b/v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb
index 054fc39..daa867d 100644
--- a/v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb
+++ b/v3/gpu-systems/flash-attention-triton/flash-attention-triton_SOLN.ipynb
@@ -80,7 +80,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "id": "e1ce38d2",
    "metadata": {},
    "outputs": [
@@ -113,7 +113,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "id": "cdf6b637",
    "metadata": {},
    "outputs": [
@@ -154,7 +154,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "id": "bdb3477a",
    "metadata": {},
    "outputs": [],
@@ -178,7 +178,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "id": "072c33a3",
    "metadata": {},
    "outputs": [],
@@ -259,7 +259,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "id": "422ab2cd",
    "metadata": {},
    "outputs": [
@@ -329,16 +329,11 @@
     "            correction = tl.exp(running_max - new_max)\n",
     "            P = tl.exp(S - new_max)\n",
     "            \n",
-    "            running_sum = running_sum * correction + P.sum(axis=1, keep_dims=True)\n",
-    "            #tl.static_print(str((running_output.shape, correction.shape, P.shape, V_block.shape)))\n",
-    "            running_output = running_output * correction  + tl.dot(P.to(V_block.dtype), V_block)   \n",
-    "            # if bh_idx==7 and q_block_idx==3 and kv_block_idx==1:\n",
-    "            #     tl.device_print(\"running_sum\",running_sum.sum())\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",
-    "        result = running_output / running_sum\n",
-    "        # if bh_idx==7 and q_block_idx==3:\n",
-    "        #     tl.device_print(\"result\",result.sum(axis=None,keep_dims=False))\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",
@@ -387,7 +382,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 7,
    "id": "45628148",
    "metadata": {},
    "outputs": [
@@ -484,6 +479,14 @@
     "\n",
     "print(\"\\nAll tests passed!\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "59e24d81-2de9-4f9b-8896-cb7095ea8604",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {