tutorial/5 - skinny flash attention.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "torch.manual_seed(456)\n",
    "\n",
    "nrows, BLOCK_DHEAD = 16, 8\n",
    "BLOCK_N = 2\n",
    "n_segments = 4  # rename to segment number\n",
    "assert nrows % n_segments == 0\n",
    "nrows_per_segment = nrows // n_segments\n",
    "assert nrows_per_segment % BLOCK_N == 0\n",
    "nrows_per_block = nrows_per_segment // BLOCK_N\n",
    "\n",
    "Q_mat = torch.rand((1, BLOCK_DHEAD))\n",
    "K_mat = torch.rand((nrows, BLOCK_DHEAD))\n",
    "V_mat = torch.rand((nrows, BLOCK_DHEAD))\n",
    "O_mat = torch.zeros_like(Q_mat)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# GM variables\n",
    "d_j_split = torch.zeros((n_segments, nrows_per_block, Q_mat.shape[0]))\n",
    "l_j_split = torch.zeros((n_segments, nrows_per_block, Q_mat.shape[0]))\n",
    "acc_splitted = torch.zeros((n_segments, nrows_per_block, *O_mat.shape))\n",
    "\n",
    "q = Q_mat  # load from GM\n",
    "for segment_index in range(0, n_segments):  # <- this loop is parallelized, it is implicit in the Triton implementation\n",
    "    for block_index_N, block_start_N in enumerate(range(0, nrows_per_segment, BLOCK_N)):  # <- this loop is parallelized, it is implicit in the Triton implementation\n",
    "\n",
    "        block_start_N += segment_index * nrows_per_segment\n",
    "        block_end_N = block_start_N + BLOCK_N\n",
    "\n",
    "        k = K_mat[block_start_N:block_end_N, :]  # load from GM\n",
    "        qk = q @ k.T\n",
    "        l_j = torch.max(qk, dim=1).values\n",
    "        l_j_split[segment_index, block_index_N, :] = l_j  # saving to GM\n",
    "        numerators = torch.exp(qk - l_j[:, None])  # safe softmax numerator\n",
    "        d_j = torch.sum(numerators, dim=1)\n",
    "        d_j_split[segment_index, block_index_N, :] = d_j  # saving to GM\n",
    "\n",
    "        v = V_mat[block_start_N:block_end_N, :]  # load from GM\n",
    "        o_segment = numerators @ v\n",
    "\n",
    "        acc_splitted[segment_index, block_index_N, :, :] = o_segment  # saving to GM\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.4563, 0.6146, 0.5144, 0.3769, 0.5411, 0.6278, 0.4206, 0.4573]])\n",
      "tensor([[0.4563, 0.6146, 0.5144, 0.3769, 0.5411, 0.6278, 0.4206, 0.4573]])\n"
     ]
    }
   ],
   "source": [
    "# shared memory variables\n",
    "l_i = torch.zeros((Q_mat.shape[0],)) - float(\"inf\")\n",
    "d_i = torch.zeros((Q_mat.shape[0],))\n",
    "acc = torch.zeros_like(O_mat)\n",
    "\n",
    "\n",
    "for block_index_N in range(0, nrows // (n_segments * BLOCK_N)):\n",
    "    for segment_index in range(0, n_segments):\n",
    "\n",
    "        acc_i = acc_splitted[segment_index, block_index_N]\n",
    "        # l_j = l_j_split[segment_index][block_index_N]\n",
    "        l_j = l_j_split[segment_index, block_index_N]\n",
    "        d_j = d_j_split[segment_index][block_index_N]\n",
    "\n",
    "        l_new = torch.maximum(l_i, l_j)\n",
    "        alpha = torch.exp(l_i - l_new)\n",
    "        beta = torch.exp(l_j - l_new)\n",
    "        d_new = alpha * d_i + beta * d_j\n",
    "\n",
    "        p_scale = beta / d_new\n",
    "\n",
    "        acc_i *= p_scale[:, None]\n",
    "        acc_scale = d_i / d_new * alpha\n",
    "        # scaling factor is applied to the exported matrix\n",
    "        acc = acc * acc_scale[:, None]\n",
    "        acc += acc_i  # accumulating in shared memory\n",
    "        d_i = d_new\n",
    "        l_i = l_new\n",
    "\n",
    "O_mat = acc  # write to GM\n",
    "\n",
    "print(O_mat)\n",
    "print((torch.nn.functional.softmax(Q_mat @ K_mat.T, dim=1)) @ V_mat)\n",
    "assert torch.allclose(O_mat, (torch.nn.functional.softmax(Q_mat @ K_mat.T, dim=1)) @ V_mat)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.4 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "3.8.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}