key_value_cache.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks at the key-value implementation in the `CausalSelfAttention` class. Illustrates how it works and the outputs of the self-attention mechanism."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mingpt.model import GPT\n",
    "from mingpt.attention import CausalSelfAttention\n",
    "from mingpt.utils import set_seed\n",
    "import torch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We look at a small example with batch of 2 sequences of length 8 with dimension 6. For now, we will just use 1 head for simplicity.\n",
    "\n",
    "We generate a random input tensor - dim 1 of the sequence, this is like the token embeddings (with positional embeddings) that we input to the self-attention mechanism."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 8, 6])"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set_seed(42)\n",
    "batch_size = 2\n",
    "seq_len = 8\n",
    "dim = 6\n",
    "\n",
    "input = torch.rand((batch_size, seq_len, dim))\n",
    "input.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = GPT.get_default_config()\n",
    "\n",
    "config.n_layer = 1\n",
    "config.n_head = 1\n",
    "config.n_embd = dim\n",
    "config.block_size = seq_len"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will just pass this through a forward pass of the `CausalSelfAttention` class and look at the outputs when we do not use any key-value caching. Here, we just input the full sequence (as if we were going to process the entire sequence to predict the next token)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(42)\n",
    "mha = CausalSelfAttention(config)\n",
    "mha.eval()\n",
    "\n",
    "without_kv_cache = mha.forward(input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
       "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
       "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076],\n",
       "         [-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309],\n",
       "         [-0.1518,  0.3811,  0.3394, -0.0629, -0.0441,  0.0369],\n",
       "         [-0.1075,  0.3083,  0.3525, -0.0363, -0.0332,  0.0558],\n",
       "         [-0.1208,  0.3341,  0.3497, -0.0366, -0.0378,  0.0537],\n",
       "         [-0.1245,  0.3554,  0.3567, -0.0293, -0.0346,  0.0549]],\n",
       "\n",
       "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
       "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
       "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795],\n",
       "         [-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177],\n",
       "         [-0.0364,  0.0500,  0.4691, -0.0181,  0.0731,  0.2813],\n",
       "         [-0.0682,  0.1153,  0.4586, -0.0224,  0.0610,  0.2626],\n",
       "         [-0.0420,  0.0875,  0.4597, -0.0096,  0.0579,  0.2689],\n",
       "         [-0.0447,  0.0996,  0.4243, -0.0101,  0.0275,  0.2254]]],\n",
       "       grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "without_kv_cache"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we use the key-value caching mechanism (kind of). If we just set `use_kv_cache` to `True`, the key-value tensors are computed for the first time and then cached. In a decoding setting, we would be able to re-use these cached key-value tensors for the later token generation steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_seed(42)\n",
    "mha = CausalSelfAttention(config)\n",
    "mha.eval()\n",
    "\n",
    "with_kv_cache = mha.forward(input, use_kv_cache=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
       "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
       "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076],\n",
       "         [-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309],\n",
       "         [-0.1518,  0.3811,  0.3394, -0.0629, -0.0441,  0.0369],\n",
       "         [-0.1075,  0.3083,  0.3525, -0.0363, -0.0332,  0.0558],\n",
       "         [-0.1208,  0.3341,  0.3497, -0.0366, -0.0378,  0.0537],\n",
       "         [-0.1245,  0.3554,  0.3567, -0.0293, -0.0346,  0.0549]],\n",
       "\n",
       "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
       "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
       "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795],\n",
       "         [-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177],\n",
       "         [-0.0364,  0.0500,  0.4691, -0.0181,  0.0731,  0.2813],\n",
       "         [-0.0682,  0.1153,  0.4586, -0.0224,  0.0610,  0.2626],\n",
       "         [-0.0420,  0.0875,  0.4597, -0.0096,  0.0579,  0.2689],\n",
       "         [-0.0447,  0.0996,  0.4243, -0.0101,  0.0275,  0.2254]]],\n",
       "       grad_fn=<ViewBackward0>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with_kv_cache"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the outputs of the self-attention mechanism in both cases, but this is expected to be the same for the first pass. All we have done differently is to cache the key-value tensors for the later steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.equal(without_kv_cache, with_kv_cache)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can look into the key-value cache here now too. From inspecting the shape, we have actually created a cache which is a bit larger than the amount of things in the cache (we only have two batches rather than 64 which we have just set as a `max_batch_size` argument above in the config)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 1, 8, 6])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mha.cache_k.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-0.0403,  0.4043, -0.2079, -0.1077, -0.4971, -0.2688],\n",
       "          [ 0.0716,  0.6593, -0.5755,  0.0521, -0.2243, -0.1871],\n",
       "          [-0.0142,  0.5693, -0.3494,  0.0481, -0.2169,  0.0308],\n",
       "          [-0.1485,  0.2542, -0.0300, -0.3365, -0.4056, -0.2569],\n",
       "          [-0.1228,  0.1589, -0.3271, -0.0833, -0.1632, -0.2081],\n",
       "          [-0.2708,  0.2987, -0.0424, -0.2511, -0.1694, -0.0641],\n",
       "          [ 0.0767,  0.5388, -0.3960,  0.0277, -0.3429, -0.2034],\n",
       "          [ 0.1095,  0.6960, -0.4371,  0.0652, -0.3131, -0.1237]]],\n",
       "\n",
       "\n",
       "        [[[-0.2830,  0.1052,  0.3896, -0.4659, -0.3163,  0.0543],\n",
       "          [ 0.0539,  0.5208, -0.5527,  0.1603, -0.2434, -0.1303],\n",
       "          [-0.1671,  0.2920, -0.2535, -0.0939, -0.1222, -0.0482],\n",
       "          [-0.1294,  0.3796,  0.0153, -0.0825, -0.1500,  0.3572],\n",
       "          [-0.2505,  0.1359,  0.4108, -0.3808, -0.2291,  0.1948],\n",
       "          [ 0.1183,  0.5452, -0.2858,  0.0636, -0.3351, -0.0121],\n",
       "          [-0.1156,  0.4128, -0.0587, -0.2097, -0.2028, -0.1035],\n",
       "          [-0.1900,  0.1918, -0.3232, -0.1516, -0.2844, -0.4461]]],\n",
       "\n",
       "\n",
       "        [[[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
       "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]],\n",
       "       grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mha.cache_k[:3]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So how would we actually use key-value caching in practice?\n",
    "\n",
    "So initially when we have an input sequence, we will feed all of that into the model to generate the next token. This is the `pre-fill` stage as we are able to process several input tokens in one pass. Subsequently, we will only have one new token generated at a time. At each stage, we compute the key-value tensors for the new token and add them to the cache. This means that we don't input the full sequence each time (like we do without key-value caching) but just the new token. We only output the attention output for the new token as that is all is needed for generation.\n",
    "\n",
    "For the `forward` method, there's actually another argument `start_pos` which tells the model where to start processing the sequence from. In the initial (called `pre-fill` stage when we process an input sequence for the first time) generation for the first new token, we set `start_pos` to 0 to fill the cache. From there, we would set `start_pos` to the next position in the sequence to generate the next token and then increment it for the next token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "prefill input\n",
      "tensor([[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009],\n",
      "         [0.2566, 0.7936, 0.9408, 0.1332, 0.9346, 0.5936],\n",
      "         [0.8694, 0.5677, 0.7411, 0.4294, 0.8854, 0.5739]],\n",
      "\n",
      "        [[0.7539, 0.1952, 0.0050, 0.3068, 0.1165, 0.9103],\n",
      "         [0.6440, 0.7071, 0.6581, 0.4913, 0.8913, 0.1447],\n",
      "         [0.5315, 0.1587, 0.6542, 0.3278, 0.6532, 0.3958]]])\n",
      "current key-cache:\n",
      "tensor([[[[-0.0403,  0.4043, -0.2079, -0.1077, -0.4971, -0.2688],\n",
      "          [ 0.0716,  0.6593, -0.5755,  0.0521, -0.2243, -0.1871],\n",
      "          [-0.0142,  0.5693, -0.3494,  0.0481, -0.2169,  0.0308],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]],\n",
      "\n",
      "\n",
      "        [[[-0.2830,  0.1052,  0.3896, -0.4659, -0.3163,  0.0543],\n",
      "          [ 0.0539,  0.5208, -0.5527,  0.1603, -0.2434, -0.1303],\n",
      "          [-0.1671,  0.2920, -0.2535, -0.0939, -0.1222, -0.0482],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]],\n",
      "       grad_fn=<SliceBackward0>)\n",
      "prefill output\n",
      "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
      "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
      "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076]],\n",
      "\n",
      "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
      "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
      "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  3\n",
      "input\n",
      "tensor([[[0.2666, 0.6274, 0.2696, 0.4414, 0.2969, 0.8317]],\n",
      "\n",
      "        [[0.9147, 0.2036, 0.2018, 0.2018, 0.9497, 0.6666]]])\n",
      "current key-cache:\n",
      "tensor([[[[-0.0403,  0.4043, -0.2079, -0.1077, -0.4971, -0.2688],\n",
      "          [ 0.0716,  0.6593, -0.5755,  0.0521, -0.2243, -0.1871],\n",
      "          [-0.0142,  0.5693, -0.3494,  0.0481, -0.2169,  0.0308],\n",
      "          [-0.1485,  0.2542, -0.0300, -0.3365, -0.4056, -0.2569],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]],\n",
      "\n",
      "\n",
      "        [[[-0.2830,  0.1052,  0.3896, -0.4659, -0.3163,  0.0543],\n",
      "          [ 0.0539,  0.5208, -0.5527,  0.1603, -0.2434, -0.1303],\n",
      "          [-0.1671,  0.2920, -0.2535, -0.0939, -0.1222, -0.0482],\n",
      "          [-0.1294,  0.3796,  0.0153, -0.0825, -0.1500,  0.3572],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]],\n",
      "       grad_fn=<SliceBackward0>)\n",
      "output\n",
      "tensor([[[-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309]],\n",
      "\n",
      "        [[-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  4\n",
      "input\n",
      "tensor([[[0.1053, 0.2695, 0.3588, 0.1994, 0.5472, 0.0062]],\n",
      "\n",
      "        [[0.9811, 0.0874, 0.0041, 0.1088, 0.1637, 0.7025]]])\n",
      "current key-cache:\n",
      "tensor([[[[-0.0403,  0.4043, -0.2079, -0.1077, -0.4971, -0.2688],\n",
      "          [ 0.0716,  0.6593, -0.5755,  0.0521, -0.2243, -0.1871],\n",
      "          [-0.0142,  0.5693, -0.3494,  0.0481, -0.2169,  0.0308],\n",
      "          [-0.1485,  0.2542, -0.0300, -0.3365, -0.4056, -0.2569],\n",
      "          [-0.1228,  0.1589, -0.3271, -0.0833, -0.1632, -0.2081],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]],\n",
      "\n",
      "\n",
      "        [[[-0.2830,  0.1052,  0.3896, -0.4659, -0.3163,  0.0543],\n",
      "          [ 0.0539,  0.5208, -0.5527,  0.1603, -0.2434, -0.1303],\n",
      "          [-0.1671,  0.2920, -0.2535, -0.0939, -0.1222, -0.0482],\n",
      "          [-0.1294,  0.3796,  0.0153, -0.0825, -0.1500,  0.3572],\n",
      "          [-0.2505,  0.1359,  0.4108, -0.3808, -0.2291,  0.1948],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]],\n",
      "       grad_fn=<SliceBackward0>)\n",
      "output\n",
      "tensor([[[-0.1518,  0.3811,  0.3394, -0.0629, -0.0441,  0.0369]],\n",
      "\n",
      "        [[-0.0364,  0.0500,  0.4691, -0.0181,  0.0731,  0.2813]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  5\n",
      "input\n",
      "tensor([[[0.9516, 0.0753, 0.8860, 0.5832, 0.3376, 0.8090]],\n",
      "\n",
      "        [[0.6790, 0.9155, 0.2418, 0.1591, 0.7653, 0.2979]]])\n",
      "current key-cache:\n",
      "tensor([[[[-0.0403,  0.4043, -0.2079, -0.1077, -0.4971, -0.2688],\n",
      "          [ 0.0716,  0.6593, -0.5755,  0.0521, -0.2243, -0.1871],\n",
      "          [-0.0142,  0.5693, -0.3494,  0.0481, -0.2169,  0.0308],\n",
      "          [-0.1485,  0.2542, -0.0300, -0.3365, -0.4056, -0.2569],\n",
      "          [-0.1228,  0.1589, -0.3271, -0.0833, -0.1632, -0.2081],\n",
      "          [-0.2708,  0.2987, -0.0424, -0.2511, -0.1694, -0.0641],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]],\n",
      "\n",
      "\n",
      "        [[[-0.2830,  0.1052,  0.3896, -0.4659, -0.3163,  0.0543],\n",
      "          [ 0.0539,  0.5208, -0.5527,  0.1603, -0.2434, -0.1303],\n",
      "          [-0.1671,  0.2920, -0.2535, -0.0939, -0.1222, -0.0482],\n",
      "          [-0.1294,  0.3796,  0.0153, -0.0825, -0.1500,  0.3572],\n",
      "          [-0.2505,  0.1359,  0.4108, -0.3808, -0.2291,  0.1948],\n",
      "          [ 0.1183,  0.5452, -0.2858,  0.0636, -0.3351, -0.0121],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]],\n",
      "       grad_fn=<SliceBackward0>)\n",
      "output\n",
      "tensor([[[-0.1075,  0.3083,  0.3525, -0.0363, -0.0332,  0.0558]],\n",
      "\n",
      "        [[-0.0682,  0.1153,  0.4586, -0.0224,  0.0610,  0.2626]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  6\n",
      "input\n",
      "tensor([[[0.5779, 0.9040, 0.5547, 0.3423, 0.6343, 0.3644]],\n",
      "\n",
      "        [[0.8035, 0.3813, 0.7860, 0.1115, 0.2477, 0.6524]]])\n",
      "current key-cache:\n",
      "tensor([[[[-0.0403,  0.4043, -0.2079, -0.1077, -0.4971, -0.2688],\n",
      "          [ 0.0716,  0.6593, -0.5755,  0.0521, -0.2243, -0.1871],\n",
      "          [-0.0142,  0.5693, -0.3494,  0.0481, -0.2169,  0.0308],\n",
      "          [-0.1485,  0.2542, -0.0300, -0.3365, -0.4056, -0.2569],\n",
      "          [-0.1228,  0.1589, -0.3271, -0.0833, -0.1632, -0.2081],\n",
      "          [-0.2708,  0.2987, -0.0424, -0.2511, -0.1694, -0.0641],\n",
      "          [ 0.0767,  0.5388, -0.3960,  0.0277, -0.3429, -0.2034],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]],\n",
      "\n",
      "\n",
      "        [[[-0.2830,  0.1052,  0.3896, -0.4659, -0.3163,  0.0543],\n",
      "          [ 0.0539,  0.5208, -0.5527,  0.1603, -0.2434, -0.1303],\n",
      "          [-0.1671,  0.2920, -0.2535, -0.0939, -0.1222, -0.0482],\n",
      "          [-0.1294,  0.3796,  0.0153, -0.0825, -0.1500,  0.3572],\n",
      "          [-0.2505,  0.1359,  0.4108, -0.3808, -0.2291,  0.1948],\n",
      "          [ 0.1183,  0.5452, -0.2858,  0.0636, -0.3351, -0.0121],\n",
      "          [-0.1156,  0.4128, -0.0587, -0.2097, -0.2028, -0.1035],\n",
      "          [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000]]]],\n",
      "       grad_fn=<SliceBackward0>)\n",
      "output\n",
      "tensor([[[-0.1208,  0.3341,  0.3497, -0.0366, -0.0378,  0.0537]],\n",
      "\n",
      "        [[-0.0420,  0.0875,  0.4597, -0.0096,  0.0579,  0.2689]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  7\n",
      "input\n",
      "tensor([[[0.7104, 0.9464, 0.7890, 0.2814, 0.7886, 0.5895]],\n",
      "\n",
      "        [[0.6057, 0.3725, 0.7980, 0.8399, 0.1374, 0.2331]]])\n",
      "current key-cache:\n",
      "tensor([[[[-0.0403,  0.4043, -0.2079, -0.1077, -0.4971, -0.2688],\n",
      "          [ 0.0716,  0.6593, -0.5755,  0.0521, -0.2243, -0.1871],\n",
      "          [-0.0142,  0.5693, -0.3494,  0.0481, -0.2169,  0.0308],\n",
      "          [-0.1485,  0.2542, -0.0300, -0.3365, -0.4056, -0.2569],\n",
      "          [-0.1228,  0.1589, -0.3271, -0.0833, -0.1632, -0.2081],\n",
      "          [-0.2708,  0.2987, -0.0424, -0.2511, -0.1694, -0.0641],\n",
      "          [ 0.0767,  0.5388, -0.3960,  0.0277, -0.3429, -0.2034],\n",
      "          [ 0.1095,  0.6960, -0.4371,  0.0652, -0.3131, -0.1237]]],\n",
      "\n",
      "\n",
      "        [[[-0.2830,  0.1052,  0.3896, -0.4659, -0.3163,  0.0543],\n",
      "          [ 0.0539,  0.5208, -0.5527,  0.1603, -0.2434, -0.1303],\n",
      "          [-0.1671,  0.2920, -0.2535, -0.0939, -0.1222, -0.0482],\n",
      "          [-0.1294,  0.3796,  0.0153, -0.0825, -0.1500,  0.3572],\n",
      "          [-0.2505,  0.1359,  0.4108, -0.3808, -0.2291,  0.1948],\n",
      "          [ 0.1183,  0.5452, -0.2858,  0.0636, -0.3351, -0.0121],\n",
      "          [-0.1156,  0.4128, -0.0587, -0.2097, -0.2028, -0.1035],\n",
      "          [-0.1900,  0.1918, -0.3232, -0.1516, -0.2844, -0.4461]]]],\n",
      "       grad_fn=<SliceBackward0>)\n",
      "output\n",
      "tensor([[[-0.1245,  0.3554,  0.3567, -0.0293, -0.0346,  0.0549]],\n",
      "\n",
      "        [[-0.0447,  0.0996,  0.4243, -0.0101,  0.0275,  0.2254]]],\n",
      "       grad_fn=<ViewBackward0>)\n"
     ]
    }
   ],
   "source": [
    "set_seed(42)\n",
    "mha = CausalSelfAttention(config)\n",
    "mha.eval()\n",
    "\n",
    "with_cache_sequential = []\n",
    "\n",
    "# prefill the cache with the first 3 tokens\n",
    "print(\"prefill input\")\n",
    "print(input[:, :3])\n",
    "output = mha.forward(input[:, :3], use_kv_cache=True, start_pos=0)\n",
    "print(\"current key-cache:\")\n",
    "print(mha.cache_k[:batch_size])\n",
    "print(\"prefill output\")\n",
    "print(output)\n",
    "with_cache_sequential.append(output)\n",
    "\n",
    "# now we can process the rest of the sequence sequentially and update the cache\n",
    "for i in range(3, seq_len):\n",
    "    print(\"i: \", i)\n",
    "    print(\"input\")\n",
    "    print(input[:, i : (i + 1)])\n",
    "    output = mha.forward(input[:, i : (i + 1)], use_kv_cache=True, start_pos=i)\n",
    "    with_cache_sequential.append(output)\n",
    "    print(\"current key-cache:\")\n",
    "    print(mha.cache_k[:batch_size])\n",
    "    print(\"output\")\n",
    "    print(output)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see in the outputs above, we just output the attention output for the new token. If we concatenate all of our outputs, we would get the same output as if we had not used key-value caching but on the whole sequence (not exactly though due to some numerical precision issues, but it should be very close):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
       "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
       "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076],\n",
       "         [-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309],\n",
       "         [-0.1518,  0.3811,  0.3394, -0.0629, -0.0441,  0.0369],\n",
       "         [-0.1075,  0.3083,  0.3525, -0.0363, -0.0332,  0.0558],\n",
       "         [-0.1208,  0.3341,  0.3497, -0.0366, -0.0378,  0.0537],\n",
       "         [-0.1245,  0.3554,  0.3567, -0.0293, -0.0346,  0.0549]],\n",
       "\n",
       "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
       "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
       "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795],\n",
       "         [-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177],\n",
       "         [-0.0364,  0.0500,  0.4691, -0.0181,  0.0731,  0.2813],\n",
       "         [-0.0682,  0.1153,  0.4586, -0.0224,  0.0610,  0.2626],\n",
       "         [-0.0420,  0.0875,  0.4597, -0.0096,  0.0579,  0.2689],\n",
       "         [-0.0447,  0.0996,  0.4243, -0.0101,  0.0275,  0.2254]]],\n",
       "       grad_fn=<CatBackward0>)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.concatenate(with_cache_sequential, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.equal(torch.concatenate(with_cache_sequential, dim=1), with_kv_cache)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [False, False,  True, False,  True, False],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True]],\n",
       "\n",
       "        [[ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True],\n",
       "         [False,  True, False, False, False,  True],\n",
       "         [False, False, False, False, False,  True],\n",
       "         [ True,  True,  True,  True,  True,  True]]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.concatenate(with_cache_sequential, dim=1) == with_kv_cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.allclose(torch.concatenate(with_cache_sequential, dim=1), with_kv_cache)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, we just have a reminder of what we did without key-value caching.\n",
    "\n",
    "Here, at each step, we'd input all the tokens up to that point to generate the next token. This is not efficient as we are re-computing the key-value tensors for the entire sequence each time. We also are saving computation since we only compute the attention output for the new token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input:\n",
      "tensor([[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009],\n",
      "         [0.2566, 0.7936, 0.9408, 0.1332, 0.9346, 0.5936],\n",
      "         [0.8694, 0.5677, 0.7411, 0.4294, 0.8854, 0.5739]],\n",
      "\n",
      "        [[0.7539, 0.1952, 0.0050, 0.3068, 0.1165, 0.9103],\n",
      "         [0.6440, 0.7071, 0.6581, 0.4913, 0.8913, 0.1447],\n",
      "         [0.5315, 0.1587, 0.6542, 0.3278, 0.6532, 0.3958]]])\n",
      "output:\n",
      "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
      "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
      "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076]],\n",
      "\n",
      "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
      "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
      "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  4\n",
      "input:\n",
      "tensor([[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009],\n",
      "         [0.2566, 0.7936, 0.9408, 0.1332, 0.9346, 0.5936],\n",
      "         [0.8694, 0.5677, 0.7411, 0.4294, 0.8854, 0.5739],\n",
      "         [0.2666, 0.6274, 0.2696, 0.4414, 0.2969, 0.8317]],\n",
      "\n",
      "        [[0.7539, 0.1952, 0.0050, 0.3068, 0.1165, 0.9103],\n",
      "         [0.6440, 0.7071, 0.6581, 0.4913, 0.8913, 0.1447],\n",
      "         [0.5315, 0.1587, 0.6542, 0.3278, 0.6532, 0.3958],\n",
      "         [0.9147, 0.2036, 0.2018, 0.2018, 0.9497, 0.6666]]])\n",
      "output:\n",
      "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
      "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
      "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076],\n",
      "         [-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309]],\n",
      "\n",
      "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
      "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
      "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795],\n",
      "         [-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  5\n",
      "input:\n",
      "tensor([[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009],\n",
      "         [0.2566, 0.7936, 0.9408, 0.1332, 0.9346, 0.5936],\n",
      "         [0.8694, 0.5677, 0.7411, 0.4294, 0.8854, 0.5739],\n",
      "         [0.2666, 0.6274, 0.2696, 0.4414, 0.2969, 0.8317],\n",
      "         [0.1053, 0.2695, 0.3588, 0.1994, 0.5472, 0.0062]],\n",
      "\n",
      "        [[0.7539, 0.1952, 0.0050, 0.3068, 0.1165, 0.9103],\n",
      "         [0.6440, 0.7071, 0.6581, 0.4913, 0.8913, 0.1447],\n",
      "         [0.5315, 0.1587, 0.6542, 0.3278, 0.6532, 0.3958],\n",
      "         [0.9147, 0.2036, 0.2018, 0.2018, 0.9497, 0.6666],\n",
      "         [0.9811, 0.0874, 0.0041, 0.1088, 0.1637, 0.7025]]])\n",
      "output:\n",
      "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
      "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
      "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076],\n",
      "         [-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309],\n",
      "         [-0.1518,  0.3811,  0.3394, -0.0629, -0.0441,  0.0369]],\n",
      "\n",
      "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
      "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
      "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795],\n",
      "         [-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177],\n",
      "         [-0.0364,  0.0500,  0.4691, -0.0181,  0.0731,  0.2813]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  6\n",
      "input:\n",
      "tensor([[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009],\n",
      "         [0.2566, 0.7936, 0.9408, 0.1332, 0.9346, 0.5936],\n",
      "         [0.8694, 0.5677, 0.7411, 0.4294, 0.8854, 0.5739],\n",
      "         [0.2666, 0.6274, 0.2696, 0.4414, 0.2969, 0.8317],\n",
      "         [0.1053, 0.2695, 0.3588, 0.1994, 0.5472, 0.0062],\n",
      "         [0.9516, 0.0753, 0.8860, 0.5832, 0.3376, 0.8090]],\n",
      "\n",
      "        [[0.7539, 0.1952, 0.0050, 0.3068, 0.1165, 0.9103],\n",
      "         [0.6440, 0.7071, 0.6581, 0.4913, 0.8913, 0.1447],\n",
      "         [0.5315, 0.1587, 0.6542, 0.3278, 0.6532, 0.3958],\n",
      "         [0.9147, 0.2036, 0.2018, 0.2018, 0.9497, 0.6666],\n",
      "         [0.9811, 0.0874, 0.0041, 0.1088, 0.1637, 0.7025],\n",
      "         [0.6790, 0.9155, 0.2418, 0.1591, 0.7653, 0.2979]]])\n",
      "output:\n",
      "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
      "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
      "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076],\n",
      "         [-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309],\n",
      "         [-0.1518,  0.3811,  0.3394, -0.0629, -0.0441,  0.0369],\n",
      "         [-0.1075,  0.3083,  0.3525, -0.0363, -0.0332,  0.0558]],\n",
      "\n",
      "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
      "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
      "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795],\n",
      "         [-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177],\n",
      "         [-0.0364,  0.0500,  0.4691, -0.0181,  0.0731,  0.2813],\n",
      "         [-0.0682,  0.1153,  0.4586, -0.0224,  0.0610,  0.2626]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  7\n",
      "input:\n",
      "tensor([[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009],\n",
      "         [0.2566, 0.7936, 0.9408, 0.1332, 0.9346, 0.5936],\n",
      "         [0.8694, 0.5677, 0.7411, 0.4294, 0.8854, 0.5739],\n",
      "         [0.2666, 0.6274, 0.2696, 0.4414, 0.2969, 0.8317],\n",
      "         [0.1053, 0.2695, 0.3588, 0.1994, 0.5472, 0.0062],\n",
      "         [0.9516, 0.0753, 0.8860, 0.5832, 0.3376, 0.8090],\n",
      "         [0.5779, 0.9040, 0.5547, 0.3423, 0.6343, 0.3644]],\n",
      "\n",
      "        [[0.7539, 0.1952, 0.0050, 0.3068, 0.1165, 0.9103],\n",
      "         [0.6440, 0.7071, 0.6581, 0.4913, 0.8913, 0.1447],\n",
      "         [0.5315, 0.1587, 0.6542, 0.3278, 0.6532, 0.3958],\n",
      "         [0.9147, 0.2036, 0.2018, 0.2018, 0.9497, 0.6666],\n",
      "         [0.9811, 0.0874, 0.0041, 0.1088, 0.1637, 0.7025],\n",
      "         [0.6790, 0.9155, 0.2418, 0.1591, 0.7653, 0.2979],\n",
      "         [0.8035, 0.3813, 0.7860, 0.1115, 0.2477, 0.6524]]])\n",
      "output:\n",
      "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
      "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
      "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076],\n",
      "         [-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309],\n",
      "         [-0.1518,  0.3811,  0.3394, -0.0629, -0.0441,  0.0369],\n",
      "         [-0.1075,  0.3083,  0.3525, -0.0363, -0.0332,  0.0558],\n",
      "         [-0.1208,  0.3341,  0.3497, -0.0366, -0.0378,  0.0537]],\n",
      "\n",
      "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
      "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
      "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795],\n",
      "         [-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177],\n",
      "         [-0.0364,  0.0500,  0.4691, -0.0181,  0.0731,  0.2813],\n",
      "         [-0.0682,  0.1153,  0.4586, -0.0224,  0.0610,  0.2626],\n",
      "         [-0.0420,  0.0875,  0.4597, -0.0096,  0.0579,  0.2689]]],\n",
      "       grad_fn=<ViewBackward0>)\n",
      "i:  8\n",
      "input:\n",
      "tensor([[[0.8823, 0.9150, 0.3829, 0.9593, 0.3904, 0.6009],\n",
      "         [0.2566, 0.7936, 0.9408, 0.1332, 0.9346, 0.5936],\n",
      "         [0.8694, 0.5677, 0.7411, 0.4294, 0.8854, 0.5739],\n",
      "         [0.2666, 0.6274, 0.2696, 0.4414, 0.2969, 0.8317],\n",
      "         [0.1053, 0.2695, 0.3588, 0.1994, 0.5472, 0.0062],\n",
      "         [0.9516, 0.0753, 0.8860, 0.5832, 0.3376, 0.8090],\n",
      "         [0.5779, 0.9040, 0.5547, 0.3423, 0.6343, 0.3644],\n",
      "         [0.7104, 0.9464, 0.7890, 0.2814, 0.7886, 0.5895]],\n",
      "\n",
      "        [[0.7539, 0.1952, 0.0050, 0.3068, 0.1165, 0.9103],\n",
      "         [0.6440, 0.7071, 0.6581, 0.4913, 0.8913, 0.1447],\n",
      "         [0.5315, 0.1587, 0.6542, 0.3278, 0.6532, 0.3958],\n",
      "         [0.9147, 0.2036, 0.2018, 0.2018, 0.9497, 0.6666],\n",
      "         [0.9811, 0.0874, 0.0041, 0.1088, 0.1637, 0.7025],\n",
      "         [0.6790, 0.9155, 0.2418, 0.1591, 0.7653, 0.2979],\n",
      "         [0.8035, 0.3813, 0.7860, 0.1115, 0.2477, 0.6524],\n",
      "         [0.6057, 0.3725, 0.7980, 0.8399, 0.1374, 0.2331]]])\n",
      "output:\n",
      "tensor([[[-0.2179,  0.4357,  0.3333,  0.0294, -0.0310, -0.0278],\n",
      "         [-0.1889,  0.5065,  0.3372, -0.0225, -0.0488, -0.0271],\n",
      "         [-0.1682,  0.4691,  0.3606, -0.0076, -0.0312,  0.0076],\n",
      "         [-0.1536,  0.4177,  0.3679, -0.0322, -0.0192,  0.0309],\n",
      "         [-0.1518,  0.3811,  0.3394, -0.0629, -0.0441,  0.0369],\n",
      "         [-0.1075,  0.3083,  0.3525, -0.0363, -0.0332,  0.0558],\n",
      "         [-0.1208,  0.3341,  0.3497, -0.0366, -0.0378,  0.0537],\n",
      "         [-0.1245,  0.3554,  0.3567, -0.0293, -0.0346,  0.0549]],\n",
      "\n",
      "        [[ 0.0404, -0.1563,  0.5403, -0.0059,  0.1507,  0.3894],\n",
      "         [-0.0704,  0.1210,  0.4321, -0.0240,  0.0440,  0.2168],\n",
      "         [-0.0642,  0.1390,  0.4008, -0.0369,  0.0128,  0.1795],\n",
      "         [-0.0706,  0.1430,  0.4408, -0.0333,  0.0486,  0.2177],\n",
      "         [-0.0364,  0.0500,  0.4691, -0.0181,  0.0731,  0.2813],\n",
      "         [-0.0682,  0.1153,  0.4586, -0.0224,  0.0610,  0.2626],\n",
      "         [-0.0420,  0.0875,  0.4597, -0.0096,  0.0579,  0.2689],\n",
      "         [-0.0447,  0.0996,  0.4243, -0.0101,  0.0275,  0.2254]]],\n",
      "       grad_fn=<ViewBackward0>)\n"
     ]
    }
   ],
   "source": [
    "set_seed(42)\n",
    "mha = CausalSelfAttention(config)\n",
    "mha.eval()\n",
    "\n",
    "print(\"input:\")\n",
    "print(input[:, :3])\n",
    "without_cache = mha.forward(input[:, :3])\n",
    "print(\"output:\")\n",
    "print(without_cache)\n",
    "\n",
    "# now we can process the rest of the sequence sequentially and update the cache\n",
    "for i in range(4, seq_len + 1):\n",
    "    print(\"i: \", i)\n",
    "    print(\"input:\")\n",
    "    print(input[:, :i])\n",
    "    without_cache = mha.forward(input[:, :i])\n",
    "    print(\"output:\")\n",
    "    print(without_cache)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now do a quick time comparison to see how much faster key-value caching is. We will use a larger batch size, sequence length, dimension and number of heads now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([8, 64, 128])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "set_seed(42)\n",
    "batch_size = 8\n",
    "seq_len = 64\n",
    "dim = 128\n",
    "\n",
    "input = torch.rand((batch_size, seq_len, dim))\n",
    "input.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = GPT.get_default_config()\n",
    "\n",
    "config.n_head = 8\n",
    "config.n_embd = dim\n",
    "config.block_size = seq_len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "19.6 ms ± 1.3 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 100\n",
    "\n",
    "set_seed(42)\n",
    "mha = CausalSelfAttention(config)\n",
    "mha.eval()\n",
    "\n",
    "# prefill the cache with the first 3 tokens\n",
    "mha.forward(input[:, :3], use_kv_cache=True, start_pos=0)\n",
    "\n",
    "# now we can process the rest of the sequence sequentially and update the cache\n",
    "for i in range(3, seq_len):\n",
    "    with_cache_sequential = mha.forward(\n",
    "        input[:, i : (i + 1)], use_kv_cache=True, start_pos=i\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "43.9 ms ± 3.06 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit -n 100\n",
    "\n",
    "set_seed(42)\n",
    "mha = CausalSelfAttention(config)\n",
    "mha.eval()\n",
    "\n",
    "mha.forward(input[:, :3])\n",
    "\n",
    "# now we can process the rest of the sequence sequentially and update the cache\n",
    "for i in range(4, seq_len + 1):\n",
    "    without_cache = mha.forward(input[:, :i])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mingpt",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}