native-sparse-attention-pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "36351d69-f1d8-4dc2-9dcc-c6150f58305d",
   "metadata": {},
   "source": [
    "# Native-Sparse-Attention\n",
    "\n",
    "author: [dhcode-cpp](https://github.com/dhcode-cpp)\n",
    "\n",
    "blog: [【手撕NSA】DeepSeek新作-原生稀疏注意力-超长文(附代码)](https://zhuanlan.zhihu.com/p/24841366485)\n",
    "\n",
    "1. Compress Attention\n",
    "2. Selection Attention\n",
    "3. Sliding Window Attenion\n",
    "4. Gated Aggregation\n",
    "5. Stride sletection attention\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bdbe0abf-e68a-4539-af9b-fa8de9b67985",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36100c16-1bae-48d2-b69d-314bb8756db6",
   "metadata": {},
   "source": [
    "## Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "290207c6-4aeb-4a0c-9394-1b11b42f97fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = 32 # token ids \n",
    "l = 8 # block\n",
    "d = 8 # sliding stride\n",
    "block_nums = t // l\n",
    "dim = 16 # embedding dimension\n",
    "heads = 4\n",
    "head_dim = dim//heads\n",
    "batch_size = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "3d430aca-12dc-49db-a5f9-17553de1ed07",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.randn(batch_size, t, dim)\n",
    "\n",
    "Wq = torch.randn(dim, dim)\n",
    "Wk = torch.randn(dim, dim)\n",
    "Wv = torch.randn(dim, dim)\n",
    "\n",
    "Q = X @ Wq\n",
    "K = X @ Wk\n",
    "V = X @ Wv\n",
    "\n",
    "# skip apply rope"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fb2f67fa-4a71-4611-b811-29e8236dd49f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 16])\n",
      "torch.Size([1, 32, 16])\n",
      "torch.Size([1, 32, 16])\n"
     ]
    }
   ],
   "source": [
    "print(Q.shape)\n",
    "print(K.shape)\n",
    "print(V.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e3e90e6-313d-462c-9e5e-70de459176a8",
   "metadata": {},
   "source": [
    "## Attention with different KV-len"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebac0ffd-d6ec-423a-8af7-80b34c96a463",
   "metadata": {},
   "source": [
    "## Token Compression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1ba786d1-ea69-4252-91da-cf1bdecef1de",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6\n",
      "tensor([ 1,  5,  9, 13, 17, 21])\n",
      "tensor([ 8, 12, 16, 20, 24, 28])\n"
     ]
    }
   ],
   "source": [
    "d = 4\n",
    "max_idx = round(( t - l ) / d)\n",
    "print(max_idx)\n",
    "print(torch.arange(max_idx) * d + 1)\n",
    "print(torch.arange(max_idx) * d + l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ee4464bb-26bc-499b-a91c-8f9cc757f017",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4\n",
      "tensor([ 1,  9, 17, 25])\n",
      "tensor([ 8, 16, 24, 32])\n"
     ]
    }
   ],
   "source": [
    "d = l\n",
    "max_idx = round(( t ) / d)\n",
    "print(max_idx)\n",
    "print(torch.arange(max_idx) * d + 1)\n",
    "print(torch.arange(max_idx) * d + l)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ff27d753-9762-4daa-be59-4ee312cf4739",
   "metadata": {},
   "outputs": [],
   "source": [
    "W_K_cmp = torch.randn(l, 1)\n",
    "W_V_cmp = torch.randn(l, 1)\n",
    "W_pe = torch.randn(l, dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ed893036-f661-485f-81b1-317a97774d9a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 16])\n",
      "torch.Size([1, 4, 16])\n"
     ]
    }
   ],
   "source": [
    "K_cmp = []\n",
    "V_cmp = []\n",
    "for i in range(max_idx):\n",
    "    cur_K = K[:, i * d + 0: i * d + l , :] + W_pe.unsqueeze(0)\n",
    "    cur_V = V[:, i * d + 0: i * d + l , :] + W_pe.unsqueeze(0)\n",
    "    cur_K = cur_K.transpose(1, 2) @ W_K_cmp \n",
    "    cur_V = cur_V.transpose(1, 2) @ W_V_cmp\n",
    "    K_cmp.append(cur_K)\n",
    "    V_cmp.append(cur_V)\n",
    "\n",
    "K_cmp = torch.cat(K_cmp, dim = 2).transpose(1,2)\n",
    "V_cmp = torch.cat(V_cmp, dim = 2).transpose(1,2)\n",
    "print(K_cmp.shape)\n",
    "print(V_cmp.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "9c61bd97-9da2-4b22-bc04-e8e5b5f5664f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 32, 4])\n"
     ]
    }
   ],
   "source": [
    "# multi-head attn\n",
    "Q_mha = Q.view(1, t, heads, head_dim).transpose(1,2)\n",
    "K_cmp_mha = K_cmp.view(1, block_nums, heads, head_dim).transpose(1,2)\n",
    "V_cmp_mha = V_cmp.view(1, block_nums, heads, head_dim).transpose(1,2)\n",
    "score_cmp = Q_mha @ K_cmp_mha.transpose(2,3) # bs, head, q_len, k_cmp_len\n",
    "print(score_cmp.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "f8f8588f-afd6-400d-9511-354c58818db1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 32, 4])\n",
      "torch.Size([1, 32, 16])\n"
     ]
    }
   ],
   "source": [
    "p_cmp = F.softmax(score_cmp, dim = -1) \n",
    "o_cmp = p_cmp @ V_cmp_mha\n",
    "print(o_cmp.shape)\n",
    "\n",
    "o_cmp = o_cmp.transpose(2, 1).reshape(batch_size, t, dim)\n",
    "print(o_cmp.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f001f394-39a0-44cb-98a7-57658e174045",
   "metadata": {},
   "source": [
    "## Token Selection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "614e156d-359c-460b-9a2f-4b4bfc7cfdbf",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 32, 4])\n",
      "torch.Size([1, 32, 4])\n"
     ]
    }
   ],
   "source": [
    "print(p_cmp.shape)\n",
    "p_slc = p_cmp.sum(dim = 1)\n",
    "print(p_slc.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e058a6cb-741b-41c1-9b8c-41bb59aa6745",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([1, 0])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 32, 2])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "select_top_k = 2\n",
    "value, idx = torch.topk(p_slc, dim = 2, k = select_top_k)\n",
    "print(idx[0,0,:])\n",
    "idx.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1970036f-a955-4914-afad-3e437ad6fe2f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 16, 16])\n",
      "torch.Size([1, 32, 16, 16])\n"
     ]
    }
   ],
   "source": [
    "idx_slc_start = idx * d\n",
    "idx_slc_end = idx * d + l\n",
    "K_slc = torch.randn(batch_size, t, d * select_top_k, dim)\n",
    "V_slc = torch.randn(batch_size, t, d * select_top_k, dim)\n",
    "for i in range(batch_size):\n",
    "    for j in range(t):\n",
    "        for k in range(select_top_k):\n",
    "            K_slc[i, j, k * d : k * d + l, :] = K[i, idx_slc_start[i, j, k ] :  idx_slc_end[i, j, k ] , :]\n",
    "            V_slc[i, j, k * d : k * d + l, :] = V[i, idx_slc_start[i, j, k ] :  idx_slc_end[i, j, k ] , :]\n",
    "print(K_slc.shape)\n",
    "print(V_slc.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2cb3166a-2400-4286-bce1-e35610b18c3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 1, 16, 4])\n",
      "torch.Size([1, 32, 1, 16, 4])\n"
     ]
    }
   ],
   "source": [
    "# shared head KV\n",
    "# IN GQA Group: [1-head KV & N-head Q] ----repeat kv-head---> [N-head KV & N-head Q]\n",
    "\n",
    "V_slc_mha = V_slc.view(batch_size, t, select_top_k * d, heads, head_dim).transpose(2,3)\n",
    "V_slc = V_slc_mha.sum(dim = 2, keepdim = True)\n",
    "print(V_slc.shape) # bs, seq_len, head, select_seq_len, head_dim\n",
    "\n",
    "K_slc_mha = K_slc.view(batch_size, t, select_top_k * d, heads, head_dim).transpose(2,3)\n",
    "K_slc = K_slc_mha.sum(dim = 2, keepdim = True)\n",
    "print(V_slc.shape) # bs, seq_len, head, select_seq_len, head_dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "381f1729-bd1c-4d2a-8bde-9b722f7ad1e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 32, 4])\n",
      "torch.Size([1, 4, 4])\n",
      "torch.Size([1, 1, 16, 4])\n",
      "torch.Size([1, 4, 16, 4])\n",
      "torch.Size([1, 4, 16, 4])\n",
      "torch.Size([1, 4, 1, 16])\n",
      "torch.Size([1, 4, 16, 4])\n",
      "torch.Size([1, 1, 16])\n"
     ]
    }
   ],
   "source": [
    "# debug Q-1 and KV-16 attention\n",
    "print(Q_mha.shape) # bs, head, seq, head_dim\n",
    "print(Q_mha[:, :, 5, :].shape) # t=5\n",
    "print(K_slc[:, 5, :, :, :].shape) # t=5\n",
    "\n",
    "print(Q_mha[:, :, 5, :].unsqueeze(dim = 2).repeat(1, 1, select_top_k * d, 1).shape) # t=5\n",
    "print(K_slc[:, 5, :, :, :].repeat(1, heads, 1, 1).shape) # t=5\n",
    "\n",
    "Q_slc_j = Q_mha[:, :, 5, :].unsqueeze(dim = 2)\n",
    "K_slc_j = K_slc[:, 5, :, :, :].repeat(1, heads, 1, 1)\n",
    "\n",
    "attn_score_j = Q_slc_j @ K_slc_j.transpose(2,3)\n",
    "print(attn_score_j.shape) # bs, head, seq_q, seq_slc_k\n",
    "\n",
    "V_slc_j = V_slc[:, 5, :, :, :].repeat(1, heads, 1, 1)\n",
    "print(V_slc_j.shape)\n",
    "\n",
    "o_j = (attn_score_j @ V_slc_j).transpose(1,2).view(batch_size, 1, dim)\n",
    "print(o_j.shape) # bs, j, dim"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "07283b6c-3037-492a-b08a-028ff31e1742",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 16])\n"
     ]
    }
   ],
   "source": [
    "o_slc = torch.zeros(batch_size, t, dim)\n",
    "for j in range(t):\n",
    "    Q_mha[:, :, j, :].unsqueeze(dim = 2)\n",
    "    K_slc_j = K_slc[:, j, :, :, :].repeat(1, heads, 1, 1)\n",
    "    V_slc_j = V_slc[:, j, :, :, :].repeat(1, heads, 1, 1)\n",
    "    \n",
    "    attn_score_j = Q_slc_j @ K_slc_j.transpose(2,3)\n",
    "    p_slc_j = F.softmax(attn_score_j, dim = -1) \n",
    "    # print(p_slc.shape)\n",
    "\n",
    "    o_slc_j = p_slc_j @ V_slc_j # bs, seq, dim   \n",
    "    # print(o_slc_j.shape)\n",
    "\n",
    "    o_slc_j = o_slc_j.transpose(1,2).view(batch_size, 1, dim)\n",
    "    o_slc[:, j, :] = o_slc_j\n",
    "    \n",
    "print(o_slc.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad8ef5b3-c69c-4509-a47e-2edea52e3fe9",
   "metadata": {},
   "source": [
    "### Token Selection details\n",
    "\n",
    "1. NSA using GQA, so we have many group(>1)\n",
    "2. every group has indipendent KV-Selection\n",
    "3. In group, caluculative n-heads-Q and 1-heads-KV attention\n",
    "4. In group 1-heads-kv repeat to n-heads-kv, but in NSA kernel, the 1-heads-kv send to SRAM shared memory. this procedure make less meomery asscess."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5edae1df-bf98-4d67-81ca-ef7b6a24e359",
   "metadata": {},
   "source": [
    "## sliding window attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "143360d0-88bf-453d-a9e3-e40bb38db17e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[1., 0., 0., 0., 0., 0., 0.],\n",
      "        [1., 1., 0., 0., 0., 0., 0.],\n",
      "        [1., 1., 1., 0., 0., 0., 0.],\n",
      "        [0., 1., 1., 1., 0., 0., 0.],\n",
      "        [0., 0., 1., 1., 1., 0., 0.],\n",
      "        [0., 0., 0., 1., 1., 1., 0.],\n",
      "        [0., 0., 0., 0., 1., 1., 1.]])\n"
     ]
    }
   ],
   "source": [
    "# built sliding window attention\n",
    "def get_window_mask(seq_len, window):\n",
    "    mask = torch.ones(seq_len, seq_len)\n",
    "    mask = torch.tril(mask)\n",
    "    win_mask = torch.ones(seq_len - window, seq_len - window)\n",
    "    win_mask = 1.0 - torch.tril(win_mask)\n",
    "    mask[window:, :seq_len - window] = win_mask\n",
    "    return mask\n",
    "print(get_window_mask(7, 3)) # test\n",
    "window_mask = get_window_mask(t, 8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "37efa397-d24b-40b2-b20b-aae53eb97d23",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[1.0644e-15, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,\n",
      "          0.0000e+00, 0.0000e+00],\n",
      "         [1.0635e-09, 5.6887e-15, 0.0000e+00,  ..., 0.0000e+00,\n",
      "          0.0000e+00, 0.0000e+00],\n",
      "         [1.0207e-12, 5.2067e-10, 3.8002e-10,  ..., 0.0000e+00,\n",
      "          0.0000e+00, 0.0000e+00],\n",
      "         ...,\n",
      "         [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 4.3458e-08,\n",
      "          0.0000e+00, 0.0000e+00],\n",
      "         [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 4.9893e-18,\n",
      "          4.5925e-41, 0.0000e+00],\n",
      "         [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 5.0657e-10,\n",
      "          1.3605e-06, 3.3918e-06]]])\n",
      "torch.Size([1, 32, 16])\n"
     ]
    }
   ],
   "source": [
    "# simplify multihead attention\n",
    "S = Q @ K.transpose(1,2) / math.sqrt(dim)\n",
    "S = F.softmax(S, dim = -1)\n",
    "S = S * window_mask #\n",
    "print(S)\n",
    "o_win = S @ V\n",
    "print(o_win.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7a4c1b1-e5e8-4faf-99e3-b33952c6739b",
   "metadata": {},
   "source": [
    "## Gated Aggregation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "9a48287f-a31b-4dc1-a05f-d9681d0560ad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 3])\n"
     ]
    }
   ],
   "source": [
    "W_gated = torch.randn(dim, 3) # 3: cmp, slc, win\n",
    "gate = X @ W_gated\n",
    "gate = F.sigmoid(gate)\n",
    "print(gate.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0d180fc2-250e-4d06-ad14-deca69fef6fb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 32, 16])\n"
     ]
    }
   ],
   "source": [
    "o_list = [o_cmp, o_slc, o_win]\n",
    "o_star = torch.zeros(batch_size, t, dim)\n",
    "for i in range(3):\n",
    "    o_star += gate[:, :, i].unsqueeze(2) * o_list[i]\n",
    "print(o_star.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21c174aa-0917-41ef-ad8d-4c65de78dc2e",
   "metadata": {},
   "source": [
    "## stride sletection attention"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bed7e63d-a8b3-4d64-bda4-ae4a373c34ee",
   "metadata": {},
   "source": [
    "$$\n",
    "\\mathbf{p}_t^\\text{slc}[j] = \\sum_{m=0}^{\\frac{l'}{d}-1}\\sum_{n=0}^{\\frac{l}{d} -1} \\mathbf{p}_t^\\text{cmp}\\left[\\frac{l'}{d}j+m +n \\right],\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "463b5566-9395-40e3-a9e2-12f2d0fa8089",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "2\n"
     ]
    }
   ],
   "source": [
    "d = 8\n",
    "t = 512 + d\n",
    "l_cmp = 16\n",
    "l_slc = 8 # from paper：“l‘ denote the selection block size”, or setting l_slc = {4, 8, 16, 32, ...}\n",
    "m_max = l_slc // d\n",
    "n_max = l_cmp // d\n",
    "print(m_max)\n",
    "print(n_max)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "b4f20651-da5a-4197-96a5-837abc3b9d24",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "32\n"
     ]
    }
   ],
   "source": [
    "# original is t token, compress -> t_cmp token.\n",
    "t_cmp = (t - d) // l_cmp\n",
    "print(t_cmp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c5fc0d1e-bdea-4f5a-83e9-9044fd322969",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.4276,  1.5061,  0.4480,  0.7817, -0.2250,  0.4883, -0.0217, -1.8421,\n",
      "        -2.7458, -3.0163, -2.0566, -0.7802,  0.8661, -0.6864, -0.7723,  0.7143,\n",
      "        -1.6664, -0.5904,  3.4035,  3.3638, -0.3060, -0.9972, -0.0107,  0.6096,\n",
      "        -0.0962, -2.1244, -1.7315,  0.7371,  0.3672, -0.5261,  1.1155,  0.6176])\n"
     ]
    }
   ],
   "source": [
    "t_cmp = (t - d) // l_cmp\n",
    "\n",
    "p_cmp = torch.randn(t_cmp)\n",
    "p_slc = torch.zeros_like(p_cmp)\n",
    "j_factor = l_slc // d \n",
    "\n",
    "for j in range(t_cmp):\n",
    "    for m in range(m_max):\n",
    "        for n in range(n_max):\n",
    "            idx = j_factor * j + m + n\n",
    "            if idx >= t_cmp:\n",
    "                continue\n",
    "            else:\n",
    "                p_slc[j] += p_cmp[idx]\n",
    "\n",
    "print(p_slc)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "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",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}