diff --git a/README.md b/README.md index c24cf0b..e3dc845 100644 --- a/README.md +++ b/README.md @@ -20,26 +20,16 @@ git clone --recursive https://github.com/DD-DuDa/BitDecoding.git conda create -n bitdecode python=3.10 conda activate bitdecode pip install -r requirements.txt -python setup.py install +bash install.sh ``` ## Quick Start -1. See benchmark/bench_single_decode.ipynb -2. (Optional) Play with libtorch c++ - ``` - # download libtorch - - cd BitDecoding/csrc/bit_decode - mkdir build && cd build - cmake -DCMAKE_PREFIX_PATH= .. - make -j12 - ``` - -## Release Progress - -- [x] Page Implementation -- [ ] End-2-end LLMs Inference -- [ ] Hopper Implementation +1. Run the GSM8K example + ``` + cd evaluation + bash scripts/example.sh + ``` + ## Citation diff --git a/benchmark/bench_single_decode.ipynb b/benchmark/bench_single_decode.ipynb deleted file mode 100644 index 69a69c6..0000000 --- a/benchmark/bench_single_decode.ipynb +++ /dev/null @@ -1,414 +0,0 @@ -{ - "cells": [ - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [], - "source": [ - "import torch\n", - "import torch.nn as nn\n", - "import math\n", - "import triton\n", - "from einops import rearrange, repeat\n", - "import numpy as np\n", - "\n", - "from flash_attn import flash_attn_with_kvcache\n", - "from bit_decode import kvcache_pack_int, fwd_kvcache_int" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# 1. test" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [], - "source": [ - "def attention_ref(\n", - " q,\n", - " k,\n", - " v,\n", - "):\n", - " \"\"\"\n", - " Arguments:\n", - " q: (batch_size, seqlen_q, nheads, head_dim)\n", - " k: (batch_size, seqlen_k, nheads_k, head_dim)\n", - " v: (batch_size, seqlen_k, nheads_k, head_dim)\n", - " Output:\n", - " output: (batch_size, seqlen_q, nheads, head_dim)\n", - " attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout\n", - " \"\"\"\n", - " dtype_og = q.dtype\n", - "\n", - " d = q.shape[-1]\n", - "\n", - " scores = torch.einsum(\"bthd,bshd->bhts\", q / math.sqrt(d), k)\n", - " \n", - " attention = torch.softmax(scores, dim=-1).to(v.dtype)\n", - "\n", - " output = torch.einsum(\"bhts,bshd->bthd\", attention, v)\n", - "\n", - " return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [], - "source": [ - "# Define constants\n", - "batch_size = 1\n", - "nheads = 32\n", - "nheads_k = 32\n", - "d = 128\n", - "\n", - "# Sequence length\n", - "seqlen_q = 1\n", - "seqlen_kv = 4096\n", - "\n", - "# Quantization parameters\n", - "quant_mode = \"k-channel\"\n", - "num_bits = 4\n", - "pack_nums = 16 / num_bits\n", - "group_size = 128" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "seqlen_kv:4096 BitDecode vs Pytorch: 0.0028667449951171875\n", - "out_ref: \n", - "tensor([-0.0581, -0.0140, -0.0241, -0.0495, 0.0539, 0.0224, 0.0342, 0.0048],\n", - " device='cuda:0', dtype=torch.float16)\n", - "out_bitdecode: \n", - "tensor([-0.0574, -0.0154, -0.0268, -0.0524, 0.0555, 0.0147, 0.0376, 0.0105],\n", - " device='cuda:0', dtype=torch.float16)\n" - ] - } - ], - "source": [ - "# Set seed and parameters\n", - "device = \"cuda\"\n", - "dtype = torch.float16\n", - "torch.random.manual_seed(0)\n", - "\n", - "# Initialize tensors\n", - "q = torch.randn(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype)\n", - "k_cache = torch.randn(batch_size, seqlen_kv, nheads_k, d, device=device, dtype=dtype)\n", - "v_cache = torch.randn(batch_size, seqlen_kv, nheads_k, d, device=device, dtype=dtype)\n", - "\n", - "k_cache_rep = repeat(k_cache, \"b s h d -> b s (h g) d\", g=nheads // nheads_k)\n", - "v_cache_rep = repeat(v_cache, \"b s h d -> b s (h g) d\", g=nheads // nheads_k)\n", - "\n", - "# Reference attention computation\n", - "out_ref, _ = attention_ref(q, k_cache_rep, v_cache_rep)\n", - "\n", - "##################### BitDecoding Packing Kernel ##################### \n", - "\n", - "# Initialize quantization tensors\n", - "if quant_mode == \"k-channel\":\n", - " k_pack = torch.zeros((batch_size, int(seqlen_kv // pack_nums), nheads_k, d), dtype=torch.uint16, device=device)\n", - " k_params = torch.zeros((batch_size, int(seqlen_kv // group_size), nheads_k, d), dtype=torch.float32, device=device)\n", - "else:\n", - " k_pack = torch.zeros((batch_size, seqlen_kv, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device)\n", - " k_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seqlen_kv), dtype=torch.float32, device=device)\n", - "\n", - "v_pack = torch.zeros((batch_size, seqlen_kv, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device)\n", - "v_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seqlen_kv), dtype=torch.float32, device=device)\n", - "\n", - "cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_kv, seqlen_kv, dtype=torch.int32, device=device)\n", - "\n", - "kvcache_pack_int(\n", - " k_cache, k_pack, k_params,\n", - " v_cache, v_pack, v_params,\n", - " None, # opt_block_table\n", - " cu_seqlens_k, \n", - " seqlen_kv,\n", - " quant_mode,\n", - " group_size,\n", - " num_bits\n", - ")\n", - "\n", - "sm_scale = 1.0 / math.sqrt(d)\n", - "out_bitdecode = fwd_kvcache_int(\n", - " q,\n", - " k_pack, k_params, \n", - " v_pack, v_params,\n", - " None, # opt_block_table\n", - " sm_scale,\n", - " quant_mode, \n", - " group_size,\n", - " num_bits\n", - " )\n", - "\n", - "print(f\"seqlen_kv:{seqlen_kv} BitDecode vs Pytorch: {(out_bitdecode - out_ref).abs().mean().item()}\")\n", - "\n", - "print(f\"out_ref: \\n{out_ref[0,0,0,:8]}\")\n", - "print(f\"out_bitdecode: \\n{out_bitdecode[0,0,0,:8]}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## 2. Benchmark" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Single" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "data": { - "image/png": "iVBORw0KGgoAAAANSUhEUgAAAhYAAAGxCAYAAAA+tv8YAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjEsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvc2/+5QAAAAlwSFlzAAAPYQAAD2EBqD+naQAAbehJREFUeJzt3XlYlOX+BvB72PcBZN8XFUFlF2QR3EqtPO2ZWbmlWWqamWaeUjPTUsuyMsuT9jtZ2ilt39QEBtwVdyMXBBcUV/Z9nt8fj4yMoAICM8D9ua658p15Z+aZN2Run+2rEEIIEBERETUBA103gIiIiNoOBgsiIiJqMgwWRERE1GQYLIiIiKjJMFgQERFRk2GwICIioibDYEFERERNhsGCiIiImoxRS7+hWq3G2bNnYW1tDYVC0dJvT0RERI0ghEBBQQHc3NxgYHDzfokWDxZnz56Fp6dnS78tERERNYFTp07Bw8Pjpo+3eLCwtrYGIBtmY2PT0m9PREREjZCfnw9PT0/N9/jNtHiwqB7+sLGxYbAgIiJqZW43jYGTN4mIiKjJMFgQERFRk2GwICIioibT4nMs6kOtVqO8vFzXzSCCsbExDA0Ndd0MIqJWQ++CRXl5OTIzM6FWq3XdFCIAgK2tLVxcXLjvChFRPehVsBBCICcnB4aGhvD09LzlBhxEzU0IgeLiYuTm5gIAXF1dddwiIiL9p1fBorKyEsXFxXBzc4OFhYWum0MEc3NzAEBubi6cnJw4LEJEdBt61SVQVVUFADAxMdFxS4iuqw65FRUVOm4JEZH+06tgUY1j2aRP+PNIRFR/ehksiIiIqHVisGgCQgiMHTsW9vb2UCgUsLW1xeTJk5vktU+ePAmFQoG9e/c2yesRERE1JwaLJvD7779j1apV+Pnnn5GTk4Nu3brpukn14uPjgyVLlmjdt2rVKtja2rbI+yclJeH++++Hq6srLC0tERoaitWrV7fIexMRUfPQq1UhrdXx48fh6uqK2NhYAICRES9rfWzZsgXBwcGYPn06nJ2d8fPPP+Ppp5+GUqnEfffdp+vmERFRIzSox2L27NlQKBRaty5dujRX21qFESNGYOLEicjOzoZCoYCPj0+tc/773/8iMjIS1tbWcHFxwRNPPKHZGwEArly5gmHDhsHR0RHm5ubo1KkTVq5cqfUaJ06cQJ8+fWBhYYGQkBBs3br1lu06fvw47r//fjg7O8PKygo9evTAxo0bNY/37t0bWVlZePHFFzX/L5OSkjBy5Ejk5eVp7ps9ezYA2bvx1ltvYdSoUbC2toaXlxc+/fTTm77/n3/+CTMzM1y9elXr/kmTJqFv374AgFdffRVz585FbGws/P39MWnSJAwcOBDr1q275WcjIqLahADefx8YOVK37WjwUEjXrl2Rk5OjuaWmpjZHuwDIi1RUpJubEPVr4/vvv4833ngDHh4eyMnJwc6dO2udU1FRgblz52Lfvn34/vvvcfLkSYwYMULz+GuvvYbDhw/jt99+w5EjR7Bs2TI4ODhovcbMmTMxdepU7N27F507d8bQoUNRWVl503YVFhbinnvuwaZNm5Ceno6BAwdi8ODByM7OBgCsW7cOHh4eeOONNzT/L2NjY7FkyRLY2Nho7ps6darmNRcvXozIyEikp6fj+eefx3PPPYeMjIw6379fv36wtbXFd999p7mvqqoKa9euxbBhw27a7ry8PNjb29/0cSIiqu3yZeD++4HJk4FVq4BfftFhY0QDzJo1S4SEhDTkKbXk5eUJACIvL6/WYyUlJeLw4cOipKRECCFEYaEQ8iu+5W+FhfX/TO+9957w9vbWHCcmJopJkybd9PydO3cKAKKgoEAIIcTgwYPFyJEj6zw3MzNTABArVqzQ3Hfo0CEBQBw5cqT+jRRCdO3aVSxdulRz7O3tLd577z2tc1auXCmUSmWt53p7e4snn3xSc6xWq4WTk5NYtmzZTd9v0qRJom/fvprjP/74Q5iamoorV67Uef7atWuFiYmJOHjwYP0+UAu58eeSiEifpKYK4eEhv7uMjYV46SUhioub/n1u9f1dU4N7LI4ePQo3Nzf4+flh2LBhmn8B083t3r0bgwcPhpeXF6ytrZGYmAgAmmv33HPPYc2aNQgNDcW0adOwZcuWWq8RHBys+XP11tLVwylWVlaa27hx4wDIHoupU6ciMDAQtra2sLKywpEjR+7o/1fNNigUCri4uGjaMGjQIE0bunbtCgAYNmwYkpKScPbsWQDA6tWrce+999Y5OXTz5s0YOXIkPvvsM83ziYjo5tRqYN48IDEROH0acHMDVq4EZswArm0arBMNmmUYHR2NVatWISAgADk5OZgzZw569eqFgwcPwtraus7nlJWVoaysTHOcn59f7/ezsAAKCxvSwqbTVDuKFxUVYcCAARgwYABWr14NR0dHZGdnY8CAAZoKroMGDUJWVhZ+/fVXbNiwAf369cP48eOxaNEizesYGxtr/ly9YVN1obaaS1FtbGwAAFOnTsWGDRuwaNEidOzYEebm5njkkUfuqGpszTZUt6O6DStWrEBJSYnWeT169IC/vz/WrFmD5557DuvXr8eqVatqvW5ycjIGDx6M9957D08//XSj20dE1F6cOwcMGwb89Zc87t8fWLAACArSbagAGhgsBg0apPlzcHAwoqOj4e3tjW+++QajR4+u8znz58/HnDlzGtU4hQKwtGzUU/XG33//jUuXLmHBggXw9PQEAOzatavWeY6Ojhg+fDiGDx+OXr164eWXX9YKFrfSsWPHWvelpaVhxIgRePDBBwHIHoyTJ09qnWNiYqLZRv1W99WHu7t7nfcPGzYMq1evhoeHBwwMDHDvvfdqPZ6UlIT77rsPb7/9NsaOHdvg9yUiam82bJCh4sIFwMwMmDIFGDdO9ljoQzmjO9rHwtbWFp07d8axY8dues6MGTOQl5enuZ06depO3rLV8fLygomJCZYuXYoTJ07gxx9/xNy5c7XOef311/HDDz/g2LFjOHToEH7++WcEBgbe0ft26tQJ69atw969e7Fv3z488cQTtUrR+/j4ICUlBWfOnMHFixc19xUWFmLTpk24ePEiiouL76gdw4YNw549ezBv3jw88sgjMDU11Ty2efNm3HvvvXjhhRfw8MMP49y5czh37hwuX758R+9JRNQWVVTIYY4BA2So8PUFvvwSmDoV8PTUj1AB3GGwKCws1OzhcDOmpqawsbHRurUnjo6OWLVqFf73v/8hKCgICxYsqNUTYWJighkzZiA4OBgJCQkwNDTEmjVr7uh93333XdjZ2SE2NhaDBw/GgAEDEB4ernXOG2+8gZMnT8Lf3x+Ojo4AgNjYWIwbNw5DhgyBo6Mj3nnnnTtqR8eOHREVFYX9+/fXWg3yxRdfoLi4GPPnz4erq6vm9tBDD93RexIRtTVZWUBCghzuEAL417+AH34A7r0XsLPTdeu0KYSo78JKOW4/ePBgeHt74+zZs5g1axb27t2Lw4cPa76Ybic/Px9KpRJ5eXm1QkZpaSkyMzPh6+sLMzOzhn0SombCn0si0qXvvwdGjADy8uT8v5kz5V4Vzs6AQQvun32r7++aGjTH4vTp0xg6dCguXboER0dHxMfHY9u2bfUOFURERFQ/paVymOOjj+Rxly7A4sVAXBygVOq2bbfSoGBxp93zREREdHv//AM89hiwb588fvxxYNYswM8PMDHRbdtuh0UtiIiI9MiXX8pVHkVFsmdi1izgiScAJye5WlLfMVgQERHpgaIiYPx44Isv5HFIiBz6iIoCbrJVlF5isCAiItKx/fvl0EdGhpyQOXy4nKTp5QXcsDeh3mOwICIi0hEhgE8+AV58ESgrAzp0kNt0P/QQ4ODQOoY+bsRgQUREpANXrwLPPANUF4GOjgYWLQLCwlr3rtMMFkRERC1s+3ZgyBC58ZWRkZys+dJLgIeHPG7NWnBrjbard+/emDx58k0f9/HxwZIlSxr8urNnz0ZoaGij29VUFAoFvv/+ewDAyZMnoVAotAqfNZfbXVciotZGrQbeeQeIj5ehwsUF+PxzYM4cwMen9YcKgMGiRezcuVOrwFbNL+rWxtPTEzk5OejWrZuum3JHysrKEBoa2mIhiYjowgW5Bff06UBlJdC7N/Djj8AjjwD29rpuXdNhsGgBjo6OsGiqOuw6ZmhoCBcXFxi18lg9bdo0uLm56boZRNRObN4MBAcDv/8uN7iaPh3473+B8HDdlzlvagwWTaSyshITJkyAUqmEg4MDXnvtNVSXYak5FOLj4wMAePDBB6FQKDTHALBgwQI4OzvD2toao0ePRmlpaa33WbFiBQIDA2FmZoYuXbrg448/1nq8ett1e3t7WFpaIjIyEtu3b9c8vmzZMvj7+8PExAQBAQH473//q/X8o0ePIiEhAWZmZggKCsKGDRu0Hr9xKCQpKQkKhQKbNm1CZGQkLCwsEBsbi4yMDK3nvfnmm3BycoK1tTWeeeYZvPLKKw0e5vnll1+gVCqxevVqzX2ff/45unbtClNTU7i6umLChAm3fZ3ffvsNf/75Z73L0hMRNVZlJfD660C/fsC5c3L56JdfAq+8IudT6EtF0iYlWlheXp4AIPLy8mo9VlJSIg4fPixKSkpaull3JDExUVhZWYlJkyaJv//+W3z55ZfCwsJCfPrpp0IIIby9vcV7770nhBAiNzdXABArV64UOTk5Ijc3VwghxNq1a4WpqalYsWKF+Pvvv8XMmTOFtbW1CAkJ0bzPl19+KVxdXcV3330nTpw4Ib777jthb28vVq1aJYQQoqCgQPj5+YlevXoJlUoljh49KtauXSu2bNkihBBi3bp1wtjYWHz00UciIyNDLF68WBgaGoq//vpLCCFEVVWV6Natm+jXr5/Yu3evSE5OFmFhYQKAWL9+vRBCiMzMTAFApKenCyGE2Lx5swAgoqOjRVJSkjh06JDo1auXiI2N1Wq3mZmZ+Pzzz0VGRoaYM2eOsLGx0fpsN7uukyZNEkIIsXr1amFtbS1++uknzeMff/yxMDMzE0uWLBEZGRlix44dmut8M+fOnRPu7u5i586dtT7LzbTWn0si0q3Tp4Xo1UsIuahUiHvuEWLvXiFKS3Xdssa51fd3TfodLNRqIQoLdXNTq+v9mRITE0VgYKBQ13jO9OnTRWBgoBBCO1gIIbS+qKvFxMSI559/Xuu+6OhorS9ff39/8dVXX2mdM3fuXBETEyOEEGL58uXC2tpaXLp0qc52xsbGijFjxmjd9+ijj4p77rlHCCHEH3/8IYyMjMSZM2c0j//222/1ChYbN27UPOeXX34RADT/H6Ojo8X48eO13jcuLq7eweLDDz8USqVSJCUlaT3u5uYmZs6cecvXqEmtVouBAweKuXPn1vlZbobBgoga6uefhbC3l4HC3FyIOXOEyMkRoqpK1y1rvPoGC/0eCikuBqysdHMrLm5QU3v27AlFjZ1MYmJicPToUVRVVdXr+UeOHEF0dLTWfTExMZo/FxUV4fjx4xg9ejSsrKw0tzfffBPHjx8HAOzduxdhYWGwv8ksoCNHjiAuLk7rvri4OBw5ckTzuKenp9bcg5ptuJXg4GDNn11dXQEAubm5AICMjAxERUVpnV/zWKVSaX2mmkMd3377LV588UVs2LABiYmJmvtzc3Nx9uxZ9OvXr872jBs3Tus1AWDp0qUoKCjAjBkz6vWZiIgaqrwcmDIFuO8+4PJloFMnYM0aYPJkuQKkJcuc60rrnoHXjhQWFgIAPvvss1oBxPDaIJ25DmcAGdfYc7Y6YKnV6no9NzIyUmtlhrOzs+bPYWFh2LNnDz7//HNERkZqXvt2n/WNN97A1KlTte7766+/sHXrVpiamtZ6/2HDhuGL6g36iYga4cQJuTfFrl3y+OGHgblzAX9//a9I2pT0O1hYWADXvlB18t4NUHOCJABs27YNnTp10nzp12RsbFyrJyMwMBDbt2/H008/rfUa1ZydneHm5oYTJ05g2LBhdbYhODgYK1aswOXLl+vstQgMDERaWhqGDx+uuS8tLQ1BQUGax0+dOoWcnBxNr0PNNjRWQEAAdu7cqfXZdu7cqfmzubk5OnbsWOdz/f39sXjxYvTu3RuGhob48MMPAQDW1tbw8fHBpk2b0KdPn1rPc3JygpOTk9Z9H3zwAd58803N8dmzZzFgwACsXbu2VlgjImqItWuBMWOAggJZMGzWLOCppwBHx9a5Lfed0O9goVC0mn1Ns7OzMWXKFDz77LPYs2cPli5disWLF9d5bvUXYlxcHExNTWFnZ4dJkyZhxIgRiIyMRFxcHFavXo1Dhw7Bz89P87w5c+bghRdegFKpxMCBA1FWVoZdu3bhypUrmDJlCoYOHYq33noLDzzwAObPnw9XV1ekp6fDzc0NMTExePnll/HYY48hLCwM/fv3x08//YR169Zh48aNAID+/fujc+fOGD58OBYuXIj8/HzMnDnzjq/NxIkTMWbMGERGRiI2NhZr167F/v37tT7brXTu3BmbN29G7969YWRkpFlhM3v2bIwbNw5OTk4YNGgQCgoKkJaWhokTJ9b5Ol5eXlrH1UMk/v7+8PDwaPwHJKJ2q7gYmDQJWLFCHnftCixZAvTsKUfV26UWmvOh0VZXhTz//PNi3LhxwsbGRtjZ2YlXX31VM5nzxsmbP/74o+jYsaMwMjIS3t7emvvnzZsnHBwchJWVlRg+fLiYNm1arQmOq1evFqGhocLExETY2dmJhIQEsW7dOs3jJ0+eFA8//LCwsbERFhYWIjIyUmzfvl3z+Mcffyz8/PyEsbGx6Ny5s/i///s/rdfPyMgQ8fHxwsTERHTu3Fn8/vvv9Zq8eeXKFc1rpKenCwAiMzNTc98bb7yh+WyjRo0SL7zwgujZs+dtr2v1qhAhhDh8+LBwcnISU6ZM0dz3ySefiICAAGFsbCxcXV3FxIkTb/maNXHyJhHdiUOHhAgKkhM0FQohnn5aiKNHhSgv13XLmkd9J28qhLi22UILyc/Ph1KpRF5eHmxsbLQeKy0tRWZmJnx9fWFmZtaSzaIWdtddd8HFxaXWPhr6iD+XRFSTEHIb7okTgZISwM4OePNNWfa8Q4e2O/Rxq+/vmvR7KITahOLiYnzyyScYMGAADA0N8fXXX2Pjxo21Nt8iItJ3+fnAs8/KlR4AEBEBvPsuEBnZ4Kl5bRaDBTU7hUKBX3/9FfPmzUNpaSkCAgLw3XffoX///rpuGhFRve3eLXslTpyQy0bHjpU7aLq7t43iYU2Fl4Kanbm5uWaCKBFRayME8P77wLRpQEUF4OQELFgA3H9/2yoe1lQYLIiIiG7i0iVg5Ejgp5/kcXw8sGiRLCjW1oqHNRUGCyIiojqkpgKPPw6cOQMYG8vJmi++CLi6ttHiYU2EwYKIiKiGqipg/nxg9mz5Z3d32UsxcCBga6vr1uk/BgsiIqJrcnKAJ58E/vpLHt99N/D220BgIHBDNQC6CQYLIiIiAH/8IbfhvnABMDMDXn4ZeP55OVmzPRQPayoMFkRE1K5VVAD//jfwzjvy2M8PWLwY6NsXuMU+UHQTzGBNoHfv3pg8efJNH/fx8dHUt2iI2bNnIzQ0tNHtaioKhQLff/89AODkyZNQKBRa1Uiby+2uKxHRncrKAhISroeKBx6QK0DuuYehorEYLFrAzp07MXbsWM1xzS/q1sbT0xM5OTno1q2brpvSYCdPnsTo0aPh6+sLc3Nz+Pv7Y9asWSgvL9d104hIB9atA0JCgG3bZL3L+fOB5cvlfIr2VOa8qXEopAU4OjrquglNxtDQEC4uLrpuRqP8/fffUKvVWL58OTp27IiDBw9izJgxKCoqwqJFi3TdPCJqIaWlwEsvAR9/LI+7dAHee0/uUdFuK5I2IfZYNJHKykpMmDABSqUSDg4OeO2111Bd363mUIiPjw8A4MEHH4RCodAcA8CCBQvg7OwMa2trjB49GqWlpbXeZ8WKFQgMDISZmRm6dOmCj6v/Zlxz+vRpDB06FPb29rC0tERkZCS2b9+ueXzZsmXw9/eHiYkJAgICahUBO3r0KBISEmBmZoagoKBa9TxuHApJSkqCQqHApk2bEBkZCQsLC8TGxiIjI0PreW+++SacnJxgbW2NZ555Bq+88kqDh3l++eUXKJVKrF69WnPf559/jq5du8LU1BSurq6YMGHCTZ8/cOBArFy5EnfffTf8/Pzwr3/9C1OnTsW6desa1A4iar3++QeIjr4eKoYOlUMf/foxVDQVBosm8sUXX8DIyAg7duzA+++/j3fffRcrVqyodd7OnTsBACtXrkROTo7m+JtvvsHs2bPx1ltvYdeuXXB1da0VGlavXo3XX38d8+bNw5EjR/DWW2/htddewxdffAEAKCwsRGJiIs6cOYMff/wR+/btw7Rp06BWqwEA69evx6RJk/DSSy/h4MGDePbZZzFy5Ehs3rwZAKBWq/HQQw/BxMQE27dvxyeffILp06fX6/PPnDkTixcvxq5du2BkZIRRo0ZptXvevHl4++23sXv3bnh5eWHZsmUNur5fffUVhg4ditWrV2PYsGEAZEgaP348xo4diwMHDuDHH39Ex44dG/S6eXl5sOeevETtwn//C4SHA/v3A0ql3Kb7gw8Af3+5ARY1kZao4V7Treq5l5SUiMOHD4uSkhIhhBBqtVoUlhXq5KZWq+v9mRITE0VgYKDWc6ZPny4CAwOFEEJ4e3uL9957T/MYALF+/Xqt14iJiRHPP/+81n3R0dEiJCREc+zv7y+++uorrXPmzp0rYmJihBBCLF++XFhbW4tLly7V2c7Y2FgxZswYrfseffRRcc899wghhPjjjz+EkZGROHPmjObx3377Tau9mZmZAoBIT08XQgixefNmAUBs3LhR85xffvlFAND8f4yOjhbjx4/Xet+4uDitz1aXxMREMWnSJPHhhx8KpVIpkpKStB53c3MTM2fOvOVr3MrRo0eFjY2N+PTTT2953o0/l0TUuhQUCPH000LIqh9ChIYKkZQkRFGRrlvWutzq+7smvZ5jUVxRDKv5uumbKpxRCEsTy3qf37NnTygUCs1xTEwMFi9ejKqqqno9/8iRIxg3bpzWfTExMZrehKKiIhw/fhyjR4/GmDFjNOdUVlZCqVQCAPbu3YuwsLCb/gv8yJEjWpNIASAuLg7vv/++5nFPT0+4ublptaE+goODNX92dXUFAOTm5sLLywsZGRl4/vnntc6PiorCX9d2oFGpVBg0aJDmseXLl2t6Jb799lvk5uYiLS0NPXr00JyTm5uLs2fPol+/fnW2Z9y4cfjyyy81x4WFhVqPnzlzBgMHDsSjjz6qdT2JqG3Zt09WJP3nH7kXxahRwMyZgKcnt+VuLnodLOi66i/Gzz77DNHR0VqPGV7722Guw4o4xjX6EasDVvUQzO1ERkZqLV91dnbW/DksLAx79uzB559/jsjISM1r3+6zvvHGG5g6dWqdj509exZ9+vRBbGwsPv3003q1kYhaFyGAZcuAKVOAsjLAwUGu+njoIVYkbW56HSwsjC1QOKPw9ic203s3RM0JkgCwbds2dOrUSfOlX5OxsXGtnozAwEBs374dTz/9tNZrVHN2doabmxtOnDih+df8jYKDg7FixQpcvny5zl6LwMBApKWlYfjw4Zr70tLSEBQUpHn81KlTyMnJ0fQ61GxDYwUEBGDnzp1an616bgkgQ8LN5kb4+/tj8eLF6N27NwwNDfHhhx8CAKytreHj44NNmzahT58+tZ7n5OQEJyenWvefOXMGffr0QUREBFauXAkDbqdH1OZcvQqMHi2XkwJATAzw7rtAaKjcUZOal14HC4VC0aDhCF3Kzs7GlClT8Oyzz2LPnj1YunQpFi9eXOe51V+IcXFxMDU1hZ2dHSZNmoQRI0YgMjIScXFxWL16NQ4dOgQ/Pz/N8+bMmYMXXngBSqUSAwcORFlZGXbt2oUrV65gypQpGDp0KN566y088MADmD9/PlxdXZGeng43NzfExMTg5ZdfxmOPPYawsDD0798fP/30E9atW4eNGzcCAPr374/OnTtj+PDhWLhwIfLz8zFz5sw7vjYTJ07EmDFjEBkZidjYWKxduxb79+/X+my30rlzZ2zevBm9e/eGkZGRZoXN7NmzMW7cODg5OWHQoEEoKChAWloaJk6cWOfrnDlzBr1794a3tzcWLVqECxcuaB5rrUtoiUjbtm3AkCFAdjZgZCS35H75ZVYkbVEtNOdDoyGTN1uLxMRE8fzzz4tx48YJGxsbYWdnJ1599VXNZM4bJ2/++OOPomPHjsLIyEh4e3tr7p83b55wcHAQVlZWYvjw4WLatGm1JjiuXr1ahIaGChMTE2FnZycSEhLEunXrNI+fPHlSPPzww8LGxkZYWFiIyMhIsX37ds3jH3/8sfDz8xPGxsaic+fO4v/+7/+0Xj8jI0PEx8cLExMT0blzZ/H777/Xa/LmlStXNK+Rnp4uAIjMzEzNfW+88Ybms40aNUq88MILomfPnre9rpMmTdIcHz58WDg5OYkpU6Zo7vvkk09EQECAMDY2Fq6urmLixIk3fb2VK1cKAHXebqW1/lwStSdVVUIsWCCEkZGcoOnqKsSXXwpR41cT3aH6Tt5UCHFts4UWkp+fD6VSiby8PNjcsF9qaWkpMjMz4evrCzP2V7Vpd911F1xcXGrto6GP+HNJpN9yc2XxsD//lMd9+8oy50FBrEjalG71/V2TXg+FUNtQXFyMTz75BAMGDIChoSG+/vprbNy4sdbmW0REDfXXX8ATTwDnz8ttuKdMAV54AXB2ZkVSXWGwoGanUCjw66+/Yt68eSgtLUVAQAC+++479O/fX9dNI6JWqrISmDMHmDdPrgDx9pYVSe+6i8XDdI3Bgpqdubm5ZoIoEdGdOn1absWdmiqP770XePttoFMnFg/TBwwWRETUavz8MzB8OHD5MmBhAcyYAYwdCzg6AjX2KCQdYrAgIiK9V14OTJsm63sAQOfOsiJpQgKLh+kbvQwWLbxQheiW+PNIpFvHj8ttuffskcePPgq8+Sbg5yf3qiD9oldzZqt3qSwvL9dxS4iuKy4uBqC9bTkRtYw1a4CwMBkqrK3lDpoffyznUzBU6Ce9+t9iZGQECwsLXLhwAcbGxtxumXRKCIHi4mLk5ubC1ta2zu3Ziah5FBfLZaP/+Y887t5dDn3ExMi5FaS/9CpYKBQKuLq6IjMzE1lZWbpuDhEAwNbWllt+E7WgQ4fk0Mfhw3JC5tNPA7NnsyJpa6FXwQIATExM0KlTJw6HkF4wNjZmTwVRCxECWLECmDQJKCmRVUjnzZO1P+zsdN06qi+9CxYAYGBgwK2TiYjakfx8uWx07Vp5HBkJLFkCRESwImlro5fBgoiI2o9du2SvxIkTchvucePk/hRubtyWuzVisCAiIp0QQu5LMW0aUFEh63u8/TbwwAOAUqnr1lFjMVgQEVGLu3RJ7qD5yy/yOCFBLiXt1o0VSVs7BgsiImpRKhXw+OPA2bOAsTEwebKsSurkxKGPtoDBgoiIWkRVlVzlMWcOoFYDHh6yIumgQXLzK2obGCyIiKjZ5eQAw4YBmzfL4wEDgIULgS5dZK8FtR131Om0YMECKBQKTJ48uYmaQ0REbc3vvwPBwTJUmJnJza7++185n4Khou1pdI/Fzp07sXz5cgQHBzdle4iIqI2oqABmzpQ9E4AsGrZkCdC3L2BpqdOmUTNqVI9FYWEhhg0bhs8++wx23A6NiIhucPIk0KvX9VDx4IPAr7/K+RQMFW1bo4LF+PHjce+996J///5N3R4iImrl1q0DQkOB7dtliHjnHeCzz4DOnVmRtD1o8P/iNWvWYM+ePdi5c2e9zi8rK0NZWZnmOD8/v6FvSURErUBpqVw2umyZPA4KkkMfcXGsSNqeNKjH4tSpU5g0aRJWr15d71oe8+fPh1Kp1Nw8PT0b1VAiItJfGRlAdPT1UPHkk8DPP8v5FAwV7YtCCCHqe/L333+PBx98UKvaY1VVFRQKBQwMDFBWVlarEmRdPRaenp7Iy8uDjY1NE3wEIiLSpS++AMaPB4qKAFtb4M03gSeeYEXStiY/Px9KpfK2398NGgrp168fDhw4oHXfyJEj0aVLF0yfPr3O8tKmpqYw5f6sRERtTmEh8NxzwJdfyuPwcDn00aMHK5K2Zw0KFtbW1ujWrZvWfZaWlujQoUOt+4mIqO3auxd47DHg6FG5DfczzwD//jfg7s5tuds7zs8lIqJ6EwL4+GPgpZeAsjLA0VFWJH3oIVYkJemOg0VSUlITNIOIiPTdlSvAqFHA99/L45gYOfQREsKKpHQdeyyIiOi2tm4FhgwBTp2Se1FMmABMnw44OwMKha5bR/qEwYKIiG5KrZa7Z86cKauTurkBixYB993HiqRUNwYLIiKqU24u8NRTwJ9/yuN+/WSZ86AgFg+jm2OwICKiWjZtkmXOz5+X8yemTgVeeEFO1uTQB90KgwUREWlUVgKzZgHz58sVID4+wHvvAXfdxeJhVD8MFkREBEBOzBw6FEhLk8eDB8sCYh07sngY1R9/VIiICD/+CIwYIZeUWlgAr70GjBkDdOig65ZRa8NgQUTUjpWVAdOmAR98II8DAuTQR+/egLm5TptGrRSDBRFRO3XsmNyWOz1dHg8ZArz1FuDtDdRR+omoXrijOxFRO/T117JoWHo6YGMjd9Bctgzw82OooDvDHgsionakuBiYOBH4/HN5HBwsh0Gio1mRlJoGgwURUTtx8KAc+jhyRO5FMXKkXFrq4cGKpNR0GCyIiNo4IYAVK+QGV6WlgL09sGCBDBmsSEpNjcGCiKgNy8sDxo4FvvlGHkdFyaGPsDDAxES3baO2icGCiKiN2rlTrvTIzJQTMp97Dnj1VcDFhdtyU/NhsCAiamOEkHtRvPIKUFEhg8TChcD997MiKTU/Bgsiojbk4kVg+HDg11/lcWIi8P77rEhKLYfBgoiojUhJkbU+zp6VIWLKFHljRVJqSQwWREStXFUVMG8eMGcOoFYDnp5yKGTgQFYkpZbHYEFE1IqdPQs88QSQnCyPBw0CFi0COndmRVLSDf7YERG1Ur//Djz5JHDpktw189VXgeefZ0VS0i0GCyKiVqaiQoaIRYvksb+/rPXRrx8rkpLuMVgQEbUimZnA448DO3bI40ceAebPB3x9WTyM9AODBRFRK/Htt8Azz8jdNK2sgNmzgVGjADs7XbeM6DoGCyIiPVdSArz4IrB8uTzu2lVuyx0by4qkpH8YLIiI9Njff8tiYQcOyOOnnwbmzmVFUtJf/LEkItJDQgCrVgERETJU2NkBH38MLF0KeHkxVJD+Yo8FEZGeKSiQBcNWr5bHEREyUEREsCIp6T8GCyIiPZKeLoc+jh2TvRJjxwKvv86KpNR6MFgQEekBIYCPPgJeegkoL5f1PRYuBB56iBVJqXVhsCAi0rErV+Sy0e+/l8fx8bIiaffurEhKrQ+DBRGRDm3ZIje8OnVK1vaYNAmYPh1wcODQB7VODBZERDqgVgPvvAP8+9+yOqm7O7B4MXDffaxISq0bgwURUQs7f14WD9u4UR7fdZcscx4QwIqk1PrxR5iIqAVt3AgMGwbk5gKmpnLY44UXWJGU2g4GCyKiFlBZCcyaJQuGCSGLhi1ZInsrWJGU2hIGCyKiZpadDQwdKidqAsD998ulpH5+rEhKbQ+DBRFRM/rhB2DkSLmk1NJSbnY1dixga6vrlhE1DwYLIqJmUFYGTJ0KfPihPA4IkH/u1UvOrSBqqxgsiIia2NGjwJAhcntuQA6DLFjAiqTUPvBHnIioCX31FRAeLkOFjQ3wwQfAJ5+wIim1H+yxICJqAkVFwMSJwMqV8jgkRFYkjY5mRVJqXxgsiIju0IEDsiLp33/LbbhHjwbeeIMVSal9YsccEVEjCQEsXw5ERclQ0aEDsGKF3EXT1ZWhgton9lgQETVCXh7wzDPAt9/K45495dBHSAgrklL7xmBBRNRAO3bIVR8nT8oNriZMAF59FXB0ZC8FEYMFEVE9qdVymOOVV+QW3S4usiLpAw8AFha6bh2RfmCwICKqhwsXgOHDgd9+k8d9+8paH4GBrEhKVBP/OhAR3UZSEvDEE0BOjlw6OnUqMGUKK5IS1YXBgojoJqqqgDfflEtH1Wq5ydV77wGDBrEiKdHNMFgQEdXh7FnZS5GcLI/vvRd4913A358VSYluhcGCiOgGv/0GPP00cPGi7JmYORMYP54VSYnqg8GCiOia8nJgxgzZMwEAnToB778vJ2qyIilR/TBYEBEBOHECePxxYOdOefzoo8DChYCnJ4uHETUE/7oQUbv3v/8BYWEyVFhbywmaK1YA3t4MFUQNxR4LImq3SkqAyZOBTz+Vx926AR9+CMTEsCIpUWM1KIsvW7YMwcHBsLGxgY2NDWJiYvBb9W4xREStyJEjsnjYp5/KbbhHjAB+/x1ISGCoILoTDQoWHh4eWLBgAXbv3o1du3ahb9++uP/++3Ho0KHmah8RUZMSAvj8cyAyEjh4ELC3Bz75RBYQc3dnrQ+iO6UQQog7eQF7e3ssXLgQo0ePrtf5+fn5UCqVyMvLg42NzZ28NRFRgxQUAM89B6xeLY8jI4GPPgLCw7ktN9Ht1Pf7u9F/laqqqvC///0PRUVFiImJuel5ZWVlKCsr02oYEVFLS08HHnsMOHZMbnD17LPArFmsSErU1Bo83/nAgQOwsrKCqakpxo0bh/Xr1yMoKOim58+fPx9KpVJz8/T0vKMGExE1hBBymKNnTxkqnJyAlSvlUlInJ4YKoqbW4KGQ8vJyZGdnIy8vD99++y1WrFiB5OTkm4aLunosPD09ORRCRM3u8mVg5Ejgxx/lca9ectVH167clpuooeo7FHLHcyz69+8Pf39/LF++vEkbRkR0J9LS5IZXp08DxsZyWen06axIStRYzT7HoppardbqkSAi0iW1GliwAHj9dVmd1N1dbng1eDBgZqbr1hG1fQ0KFjNmzMCgQYPg5eWFgoICfPXVV0hKSsIff/zRXO0jIqq3c+eAJ58ENm2SxwMGAEuWAJ07cwdNopbSoGCRm5uLp59+Gjk5OVAqlQgODsYff/yBu+66q7naR0RULxs2yFCRmyt7Jl55BZg0iRVJiVpag4LFf/7zn+ZqBxFRo1RUyGWjCxbIFSB+frIi6V13sSIpkS5wSxgiarWysoChQ4GtW+XxQw8BixaxeBiRLvGvHhG1St9/D4SGylBhaSn3pVi5EvD1Zagg0iX2WBBRq1JWBkydKvejAIDAQLkBVq9eLB5GpA8YLIio1fjnH2DIEGDvXnn85JPA/PksHkakT9hhSEStwpdfAhERMlQolbJ42LJlgIcHQwWRPmGPBRHptaIiYMIEYNUqeRwWJodBoqJYkZRIH/GvJRHprf37ZUXSjAw5IXPMGGDOHBYPI9JnHAohIr0jBPDJJ7JXIiMDcHAAVqwA3n0XcHZmqCDSZ+yxICK9cvWq7Jn49lt5HBMj51MEB7MiKVFrwGBBRHpj+3a56iMrS86fmDgRmDmTFUmJWhMOhRCRzqnVcsfM+HgZKlxdgf/+F3jrLYYKotaGPRZEpFMXLgBPPw38/rs87tcP+OADoEsX7qBJ1BoxWBCRziQlAcOGAWfPyl0zp00DXnqJFUmJWjMGCyJqcVVVwBtvAHPnyhUg3t6yIunAgaxIStTaMVgQUYs6c0ZWJFWp5PHgwcCSJYCPD4c+iNoC/jUmohbzyy9ASIgMFebmss7H6tWAnx9DBVFbwR4LImp25eXAjBlygysA6NxZViTt0wcwNtZt24ioaTFYEFGzOnFC7k2xa5c8HjIEWLiQxcOI2ip2PhJRs1m7VhYN27ULsLaWy0j/8x/A05OhgqitYo8FETW54mJg0iRZ3wOQ23F/9BHQsycrkhK1dfwrTkRN6vBhWZH00CHZKzFqFDBvniweRkRtH4dCiKhJCAF8/jkQGSlDhb09sHy5HP5gqCBqP9hjQUR3rKAAGDcO+OoreRwVJYc+wsJYkZSovWGwIKI7kp4uhz6OHZMh4vnngVmzWDyMqL3iUAgRNYoQwMcfywmZx44BTk7AqlXAO+8wVBC1Z+yxIKIGu3oVeOYZ4Lvv5HGvXnLoo2tX7qBJ1N4xWBBRg+zYIYc+srLk0tHJk4FXXwXs7HTdMiLSBwwWRFQvQshiYdOnAxUVgKsr8N57wAMPsCIpEV3HYEFEt3XpEjBiBPDzz/K4b1859NG5M4c+iEgbgwUR3VJamixzfuoUYGICTJ0KvPwyYGur65YRkT5isCCiOqnVcoXHv/8NVFXJomFLlgCDB8uAQURUFwYLIqolNxd46ingzz/l8YABssy5vz+HPojo1hgsiEhLUhLwxBNATo6clPnKK8CLLwJKpa5bRkStAYMFEQGQwx1vvgm88YYcBvHxAd5/Hxg4kEMfRFR/DBZEhJwcYNgwYPNmeTx4sJxP4esrK5QSEdUXgwVRO/fnn8CTTwIXLgDm5sBrrwETJgDW1rpuGRG1RpyGRdROVVYCM2fKoY4LF+TEzHXr5HJShgoiaiz2WBC1Q6dOAY8/DmzZIo8fegh4913Ay4tDH0R0Z9hjQdTO/PILEBoqQ4WlJbBwIfDFF4C3N0MFEd059lgQtRPl5bJY2OLF8jggQG7LnZAAGBvrtm1E1HYwWBC1A5mZwJAhwM6d8vjxx2VPhbs7eymIqGlxKISojVu3DggLk6HCykouI12xQm7RzVBBRE2NPRZEbVRpqSwW9uGH8rhrV+Djj4HYWMCIf/OJqJnw1wtRG3TsGPDoo8DevfL46aeBBQsAFxf2UhBR8+JQCFEbs2YNEB4uQ4VSKXspli0DXF0ZKoio+bHHgqiNKCkBJk0CPvtMHoeEAJ98AkRGcuiDiFoOf90QtQFHjsihj0OHZK/E6NGyoJizs65bRkTtDYdCiFq5L76QvRKHDgH29sDy5cAHHzBUEJFusMeCqJUqLJTFwr74Qh5HRsr5FOHhgKGhbttGRO0XgwVRK3TggBz6yMgADAyAceOA2bMBR0ddt4yI2jsOhRC1IkLIyZlRUTJUODgAn38OLFrEUEFE+oE9FkStRH4+8OyzcjkpAMTEyFofwcEc+iAi/cFgQdQK7Nkja30cOyZDxIQJwGuvAR066LplRETaGCyI9JgQslfipZdkdVJnZ1md9OGHATMzXbeOiKg2BgsiPXXlityPYv16eZyYCCxdKmt+GHB2FBHpqQb9epo/fz569OgBa2trODk54YEHHkBGRkZztY2o3dq+XVYkXb9e7pr58svyz927M1QQkX5r0K+o5ORkjB8/Htu2bcOGDRtQUVGBu+++G0VFRc3VPqJ2Ra2WQx3x8UBWFuDmBqxeDcydC9jZ6bp1RES3pxBCiMY++cKFC3ByckJycjISEhLq9Zz8/HwolUrk5eXBxsamsW9N1OZcuiSrkP76qzzu10+WPO/cmb0URKR79f3+vqM5Fnl5eQAAe3v7O3kZonYvNRV4/HHgzBnAxASYNk1O2LS11XXLiIgaptHBQq1WY/LkyYiLi0O3bt1uel5ZWRnKyso0x/n5+Y19S6I2R60G3n5bLh2tqgI8PYElS4B77wVMTXXdOiKihmt0sBg/fjwOHjyI1NTUW543f/58zJkzp7FvQ9RmnT8PPPkksHGjPB44UK768PPj0AcRtV6NmmMxYcIE/PDDD0hJSYGvr+8tz62rx8LT05NzLKhd++svYNgw4Nw52TPx6qvApEmAUqnrlhER1a1Z5lgIITBx4kSsX78eSUlJtw0VAGBqagpT9ukSAZDDHXPnAm+8ITe/8vWVJc7vvlvOrSAiau0aFCzGjx+Pr776Cj/88AOsra1x7tw5AIBSqYS5uXmzNJCorTh7VvZSJCXJ4/vvB959V4YLhUKnTSMiajINGgpR3OS338qVKzFixIh6vQaXm1J79McfwFNPARcuAObmwOuvA+PHA9bWum4ZEVH9NNtQCBHVX0WFDBELFsjjTp3kBM2+fQFjY922jYioObBWCFEzOXVKViTdulUeP/KI3FXT05NDH0TUdnFRG1Ez+OknIDRUhgpLS2DRImDlSsDLi6GCiNo29lgQNaHycuCVV4D33pPHgYFyW+5evTj0QUTtA4MFURPJzAQeewzYtUseP/EE8M47spAYeymIqL3gUAhRE/j2W1nmfNcuudLjgw+ATz8F3N0ZKoiofWGPBdEdKC2VxcI+/lged+8u/9yzJ2DEv11E1A7xVx9RIx09Cjz6KLBvnzweMQJ46y3AxYW9FETUfnEohKgRvvoKCA+XocLWFli2DPjoI8DVlaGCiNo39lgQNUBxMTBxIvD55/I4PFwOfUREcOiDiAhgsCCqt8OH5aqPQ4dkr8QzzwBvvgk4Oem6ZURE+oNDIUS3IYTc3CoyUoYKe3vgs8+A999nqCAiuhGDBdEtFBYCw4cDo0YBJSVAjx6yoNiIEbKYGBGRXikoANLSdNoEDoUQ3cS+fXLo459/AENDYNw4YPZswMFB1y0jIromPx9ITgaSkoCUFGD/fkCtBs6fl92rOsBgQXQDIeTmVpMmAWVlgKOjLB72yCPspSAiHcvLk0Fi82ZApZJBoqJC+xwHB7kVMIMFke7l5wNjxgDffCOPY2PlMtLu3WWvBRFRi7py5XqPRHWQqKzUPsfRUY7T9uoF9O0rf2Hp8F9BDBZE1+zeLYc+TpyQIeKFF4CZM4EOHXTdMiJqNy5f1u6ROHAAqKrSPsfJSQaJhAQZJIKCZJDQk010GCyo3RMCWLoUmDpV9ii6uADvvgs8+CBgZqbr1hFRm3bpkgwRSUlAaipw8GDtIOHsLINEYqIMEl263DJI5BblwslSd0vWGCyoXbtyRa74+P57eZyYKMucBwUBBlwzRURN7cIFGSSSk2WPxOHDtYOEiwsQFXW9R6JLF/mvnDqChFqoceD8AWw+uRmqLBW2nt6KyyWXceHlC7A2tW6hD6WNwYLarW3bgMcfB7KyAGNjYMoUYPp0wM5O1y0jojYjN1cGic2bZY/EkSNy1UZNbm7XeyT69JFBwtS0ziBRXlWO7We2IykzCamnUrH99HbkleVpnWOoMMThC4cR7RHdnJ/sphgsqN1Rq+Uqj1dflXOg3N2BJUuAwYPl32UiokY7fx746y85tJGWVneQ8PDQHtro1OmmQaKgrACqLBWSs5KReioVe3L2oLSyVOscMyMzdHPshgi3CPTy6oU4zzh42Xo144e8NQYLalcuXgSefhr47Td53L+/HPro1IlDH0TUQEIA585d75FISwP+/lveX5OHBxAdLYc2+vS5ZZDILcrF5szNSM5KRtqpNBzKPYQqoT1UojRVItQlFBFuEUj0TkSUWxRszGxgZmQGA4Xuf5ExWFC7kZICPPEEcOYMYGIihz2mTJHVSYmIbksIICcH2LRJzpFISwMyMmoHCS8v7R4Jf/86Z4ILIXDiyglsPrkZySeTseX0Fpy4cqLWeS5WLghzCUOUexQSvRMR4hwCCxMLmBqaQqEnK0FqYrCgNq+qCpg/H5g1S/ZIenvLoY9Bgzj0QUS3IIT8l8imTfJfJqmpciveG3l7y8mWvXvLm79/nb9cqtRVOHD+AP46+RdSslKw7fQ2nC86X+s8P1s/hLuGo6dHT/T26Y2O9h1hbmwOE0OTpv+MzYDBgtq0c+eAJ5+UvxcA4N57Zajw8+PQBxHdQAjg1Ck5RyI5WQaJY8dqn+fjI4c2qidb+vrWGSTKKsuw/cx2bM7crJloWVBeoHWOkYERunToggi3CMR6xiLBKwFuNm4wNzKHsaFxM33Q5sVgQW3Wpk3AsGFyLpWZmdzs6oUXABsbXbeMiPSCEEB2tvbQxvHjtc/z89PukfDxqTNI5JXmQZV9baJltpxoWV5VrnWOuZE5gp2DEeEagXiveMR6xqKDRQeYG5nD0KBtbO/LYEFtTmUl8MYbwJtvyt8bfn5yA6z+/eXcCiJqp4SQNTSqeyS2bJFb7d7I31+7R8LLq84gkVOQg+SsZCSdTEJadhoOXzwMtdBeAWJnZodQl1BEukUi0TsRkW6RsDa11puJls2BwYLalDNngKFD5b4zAPDAA3IXTR8fvdntlohaihAyOFQv/9y6VQaLG3XsKINE794ySHh41AoSQggcu3xMa6Llyasna72Um7UbwlzCEO0ejd4+vdHVsateT7RsDgwW1Gb89ptcSnrxotztdtYs4PnnAWvdbD5HRC1NrZZBonpoY8sWuQNeTQrF9SDRp48MEx4etbozq9RV2Hd+HzZnbkZKVgq2nt6KC8UXtF8KCvjb+yPcJRwxHjHo7dsbvra+rWqiZXNgsKBWr6IC+Pe/gXfekcedO8uhjz595I6aRNRGqdVycmXNIHHqlPY5CoX8pVDdI9G7t9wV74YgUVJRgu2ntyMpKwmqbBV2nNmBwvJCrXOMDYwR5BiEcNdwxHvGI8EnAc6WzjA3NoeRAb9Oq/FKUKuWlSW35d62TR4/9hiwaJH8B0g76XUkaj/Uarncs+bQxunT2ucYGFwPEtU9Eq6utYLElZIrSM1ORVJWElKzUpF+Lh0V6gqtcyyNLTUTLRO9E9HTsyfszOxgZmTWZiZaNgcGC2q1fvgBGDlSFhKzspITNseMkX8mojZArZYbUG3cKPeR2LpVTqSqycAACAgAevaUQSIxURbxuiFInM4/jZSsFM1EyyMXj0BAe2OrDuYdEOYSpploGeEWASsTK5gZmbWb+RFNgcGCWp3ycmDaNOD99+VxYCDw0UdAfDyHPohatcpKuSV2dRnxrVvlTpc1GRrKIl3VPRIJCbWChBACGRf/RvLJZLk1dnYasvOza72dp42n3NHSIwp9fPqgS4cusDSxhImhCYPEHWCwoFbl+HE59LFrlzweNgx4+21ZHJC/B4hamcpKWaSremhj2za5q11NhoZAUJB2kHB21vpXRKW6EulndiIpKwkpJ1Ow5fQWXC65rPUyBgoDdLTviAjXCDnR0qc3vJRe7X6iZXNgsKBW43//A555BsjPl5tcvfUWMGIEYGmp65YRUb1UVgKHDl3fInvbNrmDXU3VQSImRs6PSEgAnJy0gkRxRTG2ZaqQdDIJKVkp2Hl2J4orirVexsTQBF0du8qJll7xSPBOgKOFIydatgBeXdJ7paXAiy8Cn3wij4OD5dBHz56AEX+CifRXZSWwf78c2khOBrZvB3Jztc8xMgK6dr0+RyIhAXBw0AoSl4ovIfV4KpJOJiE1OxV7z+9FpbpS62WsTKwQ4hyCSNdIJHgnINojGrZmtjA3Nm+zG1HpK/5aJr2WkSFXeuzfL4c6RoyQPRXOzhz6INI7FRXAgQPaPRIXL2qfY2QEdOsmeyT69AF69QI6dNAKEllXs+TW2CeTocpWIeNSRq23crRwRJhrGCJdI9HbpzfCXMNgaWzJiZZ6gMGC9NaXXwLjxgFFRbK0+dtvy4JiFha6bhkRAZBBYt++60Fi+3bg0iXtc4yNge7dr/dIxMdrBQm1UOPIhSNyxUaWXLFxpuBMrbfyVnojzFXuaNnXpy86degEC2MLmBqxRLG+YbAgvVNUBEycCKxcKY/Dw4Fly+R/OfRBpEMVFUB6+vVaGzt2AJe1J0nCxOR6kOjbF4iLk0Hi2l/e8qpy7MnZjZSsFCSfTMbW01txpfSK1ksYKAwQ0CEA4a7hiPWIRaJPIjyVnq264md7wl/TpFcOHZJDH4cPy+XpY8bI/SmcnHTdMqJ2qKzsepBISZFB4op2CICJiZz4VD20ERcH2NtrgkRheSG2Zm1GclYyUrJSsOvsLpRUlmi9hJmRGbo6dtVU/Iz3ioejpWObqvjZnjBYkF4QQvZQTJgAlJTIf+AsXAgMGcKhD6IWU1YG7N59fbLljh1AXp72OaamQEgIEBsrV23ExgJ2dpogkVuUi9SjPyIlKwUpWSnYf34/qkSV1ksoTZVyoqXbtYmW7tGwMbNp0xU/2xMGC9K5ggLgueeA1avlcXS0XPURGipXnhFRMykrA3buvB4kdu2qHSTMzK4HiT595H+VSsDICEIIZF7NROqh3zQTLY9ePlrrbVysXBDmEoYebj2Q6J2IYJdgWJlYtauKn+0JgwXpVHq6HPo4dkyGiOeek1VJHRx03TKiNqik5HqQSEmRQSI/X/scc3PtIBETowkSaqHGwdyDUO1Zg+SsZKRmpyKnMKfW2/jZ+iHMNQxR7lHo69sXHe07wtzInBMt2wkGC9IJIYAPPwSmTpVbdDs5yeJhjzwif68RURMoKZHDGTWDREGB9jnm5kBYmAwQffvKSZc2NoCREcoqy7Dr7C6kHE5BcpacaJlfph1EDBWG6OLQRe5o6RmD3t694WbjxomW7RiDBbW4y5eB0aOB77+Xx4mJssx5UBCHPojuSHGxXPJZHSR27wYKtUt/w8JCBonYWBkkoqNlkDA0RH5ZPrac2oKU9BSoslTYeXYnyqrKtJ5ubmSO7k7dEeEWgXjPeMR5xaGDRQdOtCQNBgtqUVu2AEOHAtnZchn7iy8Cr7wi534RUQMVFclNqDZvBlQqGSSKirTPsbSUQSI+Xk62jIrSBIlzheegytqIlK0pUGWrcCD3ANRCrfV0WzNbhLqEajaiinSLhLWpNSda0k0xWFCLUKvlBlevvQZUVQEeHsB77wGDB8tJ5kRUD4WFsuJnUpLskdizR/ZS1GRlJTd9qe6RiIoCrKwgDAxw7PIxqE6uR0qWDBInrpyo9RZu1m6a0uF9ffqim1M3WJhYcKIl1RuDBTW78+fljpkbN8rjAQOADz4AOnaUe1UQ0U0UFgJpaTJIqFQySJRo7wEBa2sZJOLi5GTLHj0AKytUKYD95/dD9fcXmtLh54u0C34poIC/vb8sHe4ehX6+/eBn58eKn3RHGCyoWW3cKEPF+fNy1dorrwCTJ8tJ5kR0g/x87SCRni6r8NVkbQ1ERMgg0bev/LOVFUpFBXac2QHVvg+RnJWMbae3oaBce6KmsYExAh0CEe4Wjhj3GPTx7QMXKxdW/KQmxZ8kahaVlcDs2bJgmBCAn5/spbjrLrlRH1G7J4QMElu2yJ0tVSpg7165t0RNSqXskYiPl0EiPBywtMTVigKkZachZcdbSMlOwZ6cPSivKtd6qqWxJbo7d0ekayTiveRESzszO5gZmXGiJTUbBgtqcqdOyQmaaWny+MEHgXffBby9WZGU2rHqIKFSyR6J1NS6g4StreyFiI+XQxthYYClJc4UnYMqW4UU1StIzU7FwdyDEBBaT7U3t9fMj0jwSkAP9x6wNLGEuZE550dQi2GwoCb144/AyJFySamlpey1GDdOzicjaleEAK5e1Q4S+/fXHSQiI6/3SISGQlhYIOPKUaiyVFBt+g9U2SqcvHqy1lt42Hgg3CUckW6R6OPbB0EOQZxoSTrHYEFNoqwMmD4deP99edyli9wAKyFBUx2ZqG0TQhboUqnk8s+0NBkkyrWHJ2BvL3skevWSQSIkBJVmJtibux+qLBVSfnkXqadScbH4otbTDBQG6GTfSW6N7d4D/f36w1vpzYmWpHcYLOiOHTsmi4Xt2SOPhw2TS0vd3Dj0QW2YELJrLiXleo/EgQOytHhNHTrIHonqING9O4qNge1nd0KVvQkp62Zh2+ltKKrQ3n/CxNAEQY5BiHCNQLR7NPr69oWTpRMnWpLe408n3ZGvvwaefVbuEqxUysmaI0awIim1QUIAFy9qB4mDB+VM5ZocHOSSz/h4oF8/oGtXXEIJ0s5shSprHVLWTMaenD2oVGs/z8rECiHOIYhwjUCcVxziPONga2YLc2NzbkRFrQqDBTVKcTHwwgvAf/4jj0NDZUXSqChN9WSi1k0I4MIFWfUzKUkObRw8KHd4q8nRUQaJ6h6JoCBkV1yE6lQaVNkrofo/FQ5fOFzr5R0tHBHqEooebj3Qy6sXIt0jYWlsCTMjM86PoFaNXwHUYAcPyqGPw4flUMczzwBz5wLOzrpuGdEdUKtlkEhKkmEiLQ04dKh2kHB2vh4k+vSBOrALjhSeROrpLVBlf4CUtBScyj9V6+W9ld6y4qdbFBJ9EhHoEAgLYwtW/KQ2p8HBIiUlBQsXLsTu3buRk5OD9evX44EHHmiGppG+EQJYsUL2VJSWyqHjhQuBxx9nRVJqhdRquXNbcrKcbLllC3DkSO0g4eKi1SNR0ckfu68elkHi2JtI3ZyKyyWXtZ5iqDBE5w6dEe4ajh5uPdDPtx+8bL1Y8ZPahQYHi6KiIoSEhGDUqFF46KGHmqNNpIfy84GxY4G1a+Vxz55y1UdoKCuSUiuhVgPnzskeiaSk60FCrV10C25ucrJlQgLQpw8K/Tyw7dI+qLJToTr4Mrb9vg0lldrbapsamqKbUzeEu4ZrJlo6Wjqy4ie1Sw0OFoMGDcKgQYOaoy2kp3btkkMfJ07IEDFxIvDvf8seCyK9pVYDOTnXg0RaGpCRUTtIuLtr9Uhc8LBD6oXdMkjsehbpP6ejSmj3YtiY2iDEOQSRbpGI8YhBvFc8lGZKVvwkAudY0C0IIfelmDZNrqBzcZE7aD70ECuSkh6qqgLOntXukcjIkD/INXl4yCCRkADRpw9OOhpDlbsTquxUpKZ+hr8v/V3rpV0sXeRES/ceiPeMR7hbOKxMrLgRFVEdmj1YlJWVoazGTnP5+fnN/ZbUBC5dkstGf/5ZHvfuLYc+AgNZkZT0RFUVcOaMrLORnCzLif/zT+0g4ekpg0RiItR9euOgTSlU53cg9VQaVJsW4UzBmVov7WfrhzDXMES6RqK3T290cewCcyNzTrQkqodmDxbz58/HnDlzmvttqAmpVMATTwCnT8uCYS+9JHfVZEVS0qmqKlmIpmaPxLFjtYOEt7emR6IsIQ67LK4i9fwOqLL/RNqvs3C19KrW6YYKQwQ6BGpWbPT27Q0vJSdaEjWWQogb/1Y24MkKxW1XhdTVY+Hp6Ym8vDzY2Ng09q2pGVRVAfPnA7NmyWFoLy/gvfeA++5jRVLSgaoqIDtbrthISpI9EseO1T7Px0fTI5Ef3wNbjc9BlbMdquxU7Di7A6WV2mXHzY3M0d2pu5xo6RGNRO9ETrQkqof8/Hwolcrbfn83e4+FqakpTDkgr/dycoAnn5S9ygBw773AkiWAvz+35aYWUlkpg8SmTXJoY9s24Pjx2uf5+sqd2BIScD42GCpkQ5WzDars/2DfDy9ALbQnZ9qZ2ckdLd0i0NOjJ+K94mFjasOJlkTNpMHBorCwEMdq/KshMzMTe/fuhb29Pby8vJq0cdQy/vgDeOopuTeQuTkwcyYwaRIrklIzq6wETp7UniORmVn7PH9/oEcPiIQEHO/hD5U6E6qcbUjNXoKjPxytdbqbtZumdHisRywi3CJgaWLJiZZELaTBQyFJSUno06dPrfuHDx+OVatW3fb59e1KoeZXUQG89posGAYAHTsCH3wA9O/PiqTUDCor5ZrlzZvlbft2GSxu5O8PREejKjEB+4Odoao8jtSc7VBlq3Cu8JzWqQoo4G/vLyt+uvVAgncCujh0YcVPombQbEMhvXv3xh1MyyA9kZUFDB0q/5EIAI8+CixaJCfQ8x911CQqKuRQRnWPxLZtcqijJoVCJtroaJT2isGOrrZILT8G1Zmt2HJ6GvL/0F5FZmxgjCDHIM1Ey17evTSlw1nxk0g/8G9iO7R+PTBqFHD1qhzueOMNWaGUFUnpjlRUyMmV1UFi+/a6g0SnTkBUFK4mRGFLFwuoiv+G6sxW7Dz7DcpzyrVOtzS2RHfn7ohwiUC0RzTiPOPgbOUMMyMzTrQk0lMMFu1IaSkwdaqsQgoAXbvKvSni41mRlBqhokJuiV1dtGvHDrlGuSYDAxkkoqNxNj4EKn9jqIoOQXVmKw6cXQ1xVrv3s4N5B4S6hCLCVU607OnRE0ozJcyNzDk/gqiV4NdJO/HPP3Jb7r175fHw4XJpqaurTptFrUlFBbB/v3aQOH9e+xwDAyAgACI6Cv/EBkDlrYCq4CBUp1XIPPt/wFnt0z1tPDUTLWM8YhDqEgorU+5oSdSaMVi0A19+CYwbBxQVAba2MlAMH86KpHQb5eXAnj1yoqVKJYPEpUva5xgaAkFBqIyKxN6e3lC5VSI17yBSz/yG3LNfaAUJA4UBOtl3QpirnGgZ5xmHAIcAWBhbcKIlURvCYNGGFRYCEyYAX3whj8PD5TBIjx6sSEp1KC+XEyyTkoCUFFl9Li9P+xwTE6BbN5REhWNHlDtUjsVQXU7HljP/Q+GZQqDG7tgmhibo5tQNoc6hiPKIQrxnPDyVnrAwtuBES6I2jH+726j9++XQx99/y97psWOBuXMBBwddt4z0RmmprPiZlCR7JHbtkt1aNZmZASEhyOsZii3hjkixzYPqwm7szPk/lGeXAzXmZlqZWCHEOQThruFyfoR7TzhbOcPc2JwbURG1IwwWbYwQwPLlwOTJQFkZ4OgILFwIPP44K5K2e0VFMkBUB4k9e2S4qMnSEggLw/me3aAKsUOKxQWocndif+5nUGdq72jpYO6AUNdQRLpGoqdHT/Rw6wE7czuYGZlxfgRRO8Zg0YZcvQqMGQN8+608jo8Hli4FgoNZkbRdKii4XrBLpZIzdysqtM9RKiHCw5DZswtU3ayRYnwWqnPbcfRyKnDDbtoeNh6ajah6evREqEsobExtWPGTiLQwWLQRO3bIoY+TJ+XS0UmT5Nbcdna6bhm1mCtX5B4SSUlAaipw4IAs5FVThw5QR4TjUIw/VJ3NkKLIhursNpwtTAJq7I6tgAId7Tsi3CUcPdx7INYzFl0cusDKxIoVP4nolhgsWjm1Gnj3XWDGDLljspubPH7wQVYkbfNyc68X7EpNlXtKqLWHK+DsjIoe4dgd7QWVvxFSKo4jLWc7rpRu0AoSRgZGCHIMQrhLOKLdoxHrGQsfOx9YGltyIyoiahAGi1bswgW5bPS33+Rx//6y1keXLtyWu006exbYuFH2SKSlyc1JbuThgaIeIdgW7Q6VF6AqzcDWs0koqSwBapxubmSOYOdgWTrcXe5o6WrtyomWRHTHGCxaqaQkYNgw+V1jagpMnw689BLAum5tSFYWsGGD7JFIS6u78qevLy5HByM1whEqtwqoCg9j9/k/UFlaqRUkbE1t5Y6WbhGIdo9GjEcMHCwduBEVETU5BotWpqpKLhudO1f2evv4AO+/DwwaxIqkrZoQss7Gxo3XhzbOnNE+51rBrtMxXaEK6wCVQxFUeftx8OIPQBG0hjacLZ0R7hqOSLdIRLtHo4d7D9ia2XIjKiJqdgwWrciZM7KXIjlZHv/rX8B77wG+vhz6aHWEAA4flkEiJUX2SNSxPbboEoCjMQFICbGFyjYPqkvpyMz7HsiDvF3jY+ujNdGyu1N3WJtacyMqImpx/K3TSvz2G/D008DFi7IK6WuvARMnym0HqBWoqpK7llVPttyyBbh8WfscY2NUBQVif5w/UoIsobK8BNWFXcgtPgJchrxBbo0d0CEAEa4R6OEut8bu3KEzLIwtONGSiHSOwULPVVQA//438M478jggAFiyRE7UZEVSPVZRITeg2rRJ9khs21Z7e2xTU5SFdMXOnl5ICTCDyvQctuTuQX75fuAi5A3XtsZ27IYItwhEucutsb1svVjxk4j0Er+a9NipU3LHzC1b5PGQIXIXTU9P3baL6lBWBuzceb1HYvt2oLhY+xxzcxREdMeWaDeoOhojxeAUduSmo6xqD3Dh+mmWxpaa0uHRHnLFhrOVM8yMzFr2MxERNQKDhZ76+We5lPTyZcDKSk7WfPZZViTVG8XFsheiukdi504ZLmqytkZuVFekRjpB5WOAFPUJ7L24C2qhBmpMp7A3t0e4azgiXCPQ06MnYjxiYG9uz42oiKhVYrDQMxUVcsfMhQvlcWCg3Ja7d29WJNWp/HzZdVQdJNLTa22PLWyVyIrtClV4B6jcq5BS9g8yrm6TD9YIEu7W7gh3DUcPNznRMtw1HDamNpwfQURtAoOFHsnOlkMfW7fK46FD5dwKDw/dtqtdunJF1tfYvFkObezbV2tXS7WjA47EBUAVYguVSxlURUdwqvDauFWNINHRviMiXK5PtOzq1BWWJpbciIqI2iQGCz3BoQ8dy82VPRHVQeLwYbkktIYKdxekx3eEqqs1VA5FSM0/iEulafLBa0HCyMAIQQ5BCHcNR5R7FHp59YKfvR8nWhJRu8FgoWMVFcCrrwKLFsnjoCC5LTeHPprZmTMyQCQlyf/WsT12sa8Htsf7QtXFHCplHrbmHURRRSqgBpArzzEzMkOIc4hma+xe3r3gbu3Oip9E1G4xWOjQjUMfTzwhhz7c3XXbrjZHCFn2tWaQOHmy1mlXAn2QFusJVUcTqCwvYtfVw6hQnwYqoFn6aW1ijQjXiOtLP73i4WjhyImWRETXMFjoyE8/yaGPK1euD32MGweYcUXhnRNC9kBUB4mUlNrbYxsY4GywL1TRLlD5GUFleg4Hrv4DgZNAGeQNgJOlEyJcIxDpFokYjxhEe0RDaarkREsioptgsGhhdQ19LF0KJCZy6KPR1Grg0CEZJJKTZZDIzdU6RRga4HikP1J6OMmqn4ancLzgOIDjQAnkDYC30lszPyLOMw4hziGwNrXm/AgionpisGhBWVly6GPbtRWITzwhl5W6uem2Xa1OZaVcpVEdJFQq2fVTQ5WpMQ729IMqrANS3CuhUmfiXMlRAEeBa/tWKaBAgEPA9fkRXr0Q4BDAiZZERHeAwaKF3Dj08eabctUHhz7qobwc2L37epBISwMKCrRPsTTDrjhfqIJtkeJcirSK48grz5APFsn/GBsYo5tTN0S4RiDKIwqJ3onwVnpzoiURURNisGhmFRXAjBnA4sXyuGvX66s+DLiNQd1KS+WW2NVBYutWoKRE65RCO0tsjfeBqqsVUjoUYnvpMZRWHZEPXgsSFsYWCHUJRaRrJHq490CidyJcrV1Z8ZOIqBnxN2wzunHoY9gwueqDQx83KC6Wu1pWr9jYsUP2UtRw0dkaqfFemqWfe0pOoEockg9eCxJ2ZnaaFRs93Xsi3jse9ub23IiKiKgFMVg0k5pDH9bWwLx5wJgxHPoAIIPE1q0ySGzeLIPEDdtjZ3vbQhXrLpd+Wl3C4ZJsANeCxLU5Em7WbpoVG3Gecejh1oMTLYmIdIzBoomVl8uhj3fflccc+oAcxqgOEklJcpijRo+EAPB3gD1U0a5Q+RpCZXYeWWXnAVy99nz5Hz87P82wRoJXAro7d4e5MbcmJSLSJwwWTSgrS5Y2375dHj/5JPD22+1w6KO0VI7/VPdIbNumFSQqDYB9QbZIiXaByhtINT6HCxWXAVyWJ5QBhgpDBDoEItLt+vyITh06wcTQRCcfiYiI6ofBoon8+KMc+rh6tR0OfZSVXQ8SSUmyd6JGCfFSI2BHsBIpkY5QeQpsMTqLwqqr0PRIVACmhqYIdg6WQcKtB/r49oGHjQcnWhIRtTL8rX2H6hr6qN7wqs0OfZSVyXkR1T0SW7fKXopr8kyBLeHWSAl3gMq9EjsNzqFc5AHIkydUya2xw1zCEOkWiZ4ePdHLqxecrJw40ZKIqJVjsLgDNw59PPWUHPpwddVtu5pcefn1IJGUJFdw1Fj+ed4SUEVZQRVqhxTXcuxXXIAaBQCu7TUhAAcLB0S4RqCHWw/EesYixiMGSjMlJ1oSEbUxDBaN9MMPwIgR2kMfY8cCpm1hr6XycmDXLtkbkZQkN6S6FiQEgEw7QBVmodmM6qjiMoDCazfJ08ZTFupyi0K8dzzCXcJhaWKpi09DREQtiMGigcrLgenTgSVL5HG3bsCHHwK9erXioY+KChkkqoc20tLkklAAagVwyBFQ9TCHqpsNUhyLcFZRCLnmU56jgAKdO3RGhJvskUj0TkSQYxB3tCQiaocYLBrg+HE59LF7tzxutUMflZXyQ1T3SKSmAkVyl6kKA2C3G6DqaQZVVyuk2hfiiqIUNSt1GRkYoZtjN02Q6OPbB762viwdTkREDBb1tWaNHOooKACUSjn08cwzrWToo7IS2LPneo9EaipQKIctioyBbR6AqrMpVF3Msc22CMWKCgCl126AuZG53Bq7xooNVytXlg4nIqJaGCxuo7gYmDwZ+OwzeRwaKld9xMUBejvvsLIS2Lv3eo+ESqUp2nXZHEjzBFI6m0DV2RS7rYtQqVADKLt2A5SmSkS4RSDS9fqKjQ4WHTjRkoiIbovB4hYOHwYeeww4dEiGiGeekT0Vjo66btkNqqpkkKjukVCpgPx8AMAZa0DlDaR0NIaqozEOWl3bDxvl126Ai6WLHNZw74E4zzhEuUfBxtRGF5+EiIhaOQaLOggBrFwJTJggF0N06CDnUjz1FGCiDxs/VlUB+/df75FISQHy8iAAHO0AqPyBlI5GUPkZIdO8en+Jims3wNfWF5FukYh0i0SCVwKCnYNhYWKhow9DRERtCYPFDQoKgOeeA1avlsfR0XLoIzJSh0MfarUMEtU9EikpwNWrqFIA+50BVRcgxc8QqT4GOG9aXcyrEkAlDBQG6OLQRTM/ordPb3Tu0JlbYxMRUbNgsKghPV2u+jh6FDA0BMaPB2bPBuzsWrgh1T0S1WXEU1KAK1dQZgjsdAdU3QCVnyHSPIF846rqJwGogomhCbo7ddcEib6+feGp9OTW2ERE1CL4bQM59PHRR8BLL8l9KpydgcWL5fwK45ZYQVlVBezbpx0krl7FVTNgqweQGgak+hpgu5tAmaGofhIAwNLYEmGuYZqJlgneCXC2cubW2EREpBPtPlhcuQKMHg2sXy+PExOB998HgoObceij5mTL6lUbeXk4ZQOkegGpsUCqjwEOOKohNG1QAwA6mHdAhGsEIt0iEesZi1jPWNia2XLFBhER6YV2HSy2bQMef1zW/DAykj0WM2bIfSqaVPXyzxpBQl2Qj8OOMkio+gGp3gpkK0WNJ8kg4aX0QrhrOCJdIxHvFY8I1whYmVo1cQOJiIiaRrsMFkIA770nt+aurAQ8POTxAw/IgHHHKivlhI3qIJGaitLifOxyu9YjcR+Q5gVc1SqpLmCoMESgQyDC3WSQSPRJRECHAG6NTURErUa7CxZXrgAjR8oiYgBw990yVAQF3cGL1tzZ8lqQuFxZgC2e14LEI8BON6D8hqtdvaNl9dBGb5/e8LDx4I6WRETUarWrYLFzp5yQefKknJQ5bZq82TR0L6jqWhvXgoRIVSHLqEiGCC8g9SngkFPtp3Uw76DZ0TLaIxpxnnGwN7fn/AgiImoz2kWwEELuRTF1qizk6e4uJ2jef389hz7Uarlq46+/gM2bUaVKxkHzwutBYgxwuo55GX62vgi/FiTivOIQ5hLG0uFERNSmtflgkZcnV31895087tdPhoquXW/xJCGAI0dkkPjrL1SokrDb7ApSvIEUbyD1eSDPTPspRgpDdHXqhgjXCES4RSDROxGdOnTiRlRERNSutOlgsWePHPo4flz2TLz8MvDKK3UMfQghT7rWI1Gs+gvbjXOR4i3rbGwdCxTfkA8sjS0Q5hKGiGsbUSX6JLLiJxERtXttMlgIASxfDkyaJDe8cnMDliwBHnywxtBHVpbcHnvzZuSlbsIWgzOaHomdo4CKG/KBrYkSke6R6OEehRiPGMR5xcHOzI7zI4iIiGpoc8GipEQOfXz9tTzu00cOfXRXZgNfJQGbN+PC1k1QGZyCyksGib1PAuobNqp0Nne8FiSiEesZq6n4ySBBRER0c20qWOTlAYMHy40sjYyA+U8fwUSL/yB36Df4yuCUpkfiyNDaz/WydEekRxR6eEQh3jMeIS4hsDa1bvkPQURE1Io1Klh89NFHWLhwIc6dO4eQkBAsXboUUVFRTd22BikokBMzM3YX4BmblYgP/wg7zv2Dbv7AsUdrn9/J2geRntHo4RGNeK94BDkGccUGERHRHWpwsFi7di2mTJmCTz75BNHR0ViyZAkGDBiAjIwMODnVsXlDC6isEJj44NdwsF+A8nGH8LmTGitqDG0YQoEgZSdEevVEpGc0Er0T4W/vDzMjs5u/KBERETWYQgghbn/addHR0ejRowc+/PBDAIBarYanpycmTpyIV1555bbPz8/Ph1KpRF5eHmwavDNV3TL2Z6H7tz5aEy79K23Q0yMGscH3YGDAffCy9WLpcCIiokaq7/d3g75py8vLsXv3bsyYMUNzn4GBAfr374+tW7fW+ZyysjKUlZVpNaypBQR74/7/88OVkmI8GNAfsYPGIsg3ijU2iIiIWliDgsXFixdRVVUFZ2dnrfudnZ3x999/1/mc+fPnY86cOY1vYT2tXXgUBgqD259IREREzabZv4lnzJiBvLw8ze3UqVPN8j4MFURERLrXoB4LBwcHGBoa4vz581r3nz9/Hi4uLnU+x9TUFKamHJIgIiJqDxr0z3wTExNERERg06ZNmvvUajU2bdqEmJiYJm8cERERtS4NXiYxZcoUDB8+HJGRkYiKisKSJUtQVFSEkSNHNkf7iIiIqBVpcLAYMmQILly4gNdffx3nzp1DaGgofv/991oTOomIiKj9afA+FneqOfaxICIiouZV3+9vLqUgIiKiJsNgQURERE2GwYKIiIiaDIMFERERNRkGCyIiImoyDBZERETUZBgsiIiIqMkwWBAREVGTafDOm3eqej+u/Pz8ln5rIiIiaqTq7+3b7avZ4sGioKAAAODp6dnSb01ERER3qKCgAEql8qaPt/iW3mq1GmfPnoW1tTUUCkWTvGZ+fj48PT1x6tQpbhPehHhdmw+vbfPgdW0evK7No7VdVyEECgoK4ObmBgODm8+kaPEeCwMDA3h4eDTLa9vY2LSK/zmtDa9r8+G1bR68rs2D17V5tKbrequeimqcvElERERNhsGCiIiImkybCBampqaYNWsWTE1Ndd2UNoXXtfnw2jYPXtfmwevaPNrqdW3xyZtERETUdrWJHgsiIiLSDwwWRERE1GQYLIiIiKjJtIlg8dFHH8HHxwdmZmaIjo7Gjh07dN2kFpOSkoLBgwfDzc0NCoUC33//vdbjQgi8/vrrcHV1hbm5Ofr374+jR49qnXP58mUMGzYMNjY2sLW1xejRo1FYWKh1zv79+9GrVy+YmZnB09MT77zzTq22/O9//0OXLl1gZmaG7t2749dff21wW/TB/Pnz0aNHD1hbW8PJyQkPPPAAMjIytM4pLS3F+PHj0aFDB1hZWeHhhx/G+fPntc7Jzs7GvffeCwsLCzg5OeHll19GZWWl1jlJSUkIDw+HqakpOnbsiFWrVtVqz+1+vuvTFn2xbNkyBAcHa9btx8TE4LffftM8zut65xYsWACFQoHJkydr7uN1bZzZs2dDoVBo3bp06aJ5nNf1JkQrt2bNGmFiYiI+//xzcejQITFmzBhha2srzp8/r+umtYhff/1VzJw5U6xbt04AEOvXr9d6fMGCBUKpVIrvv/9e7Nu3T/zrX/8Svr6+oqSkRHPOwIEDRUhIiNi2bZtQqVSiY8eOYujQoZrH8/LyhLOzsxg2bJg4ePCg+Prrr4W5ublYvny55py0tDRhaGgo3nnnHXH48GHx73//WxgbG4sDBw40qC36YMCAAWLlypXi4MGDYu/eveKee+4RXl5eorCwUHPOuHHjhKenp9i0aZPYtWuX6Nmzp4iNjdU8XllZKbp16yb69+8v0tPTxa+//iocHBzEjBkzNOecOHFCWFhYiClTpojDhw+LpUuXCkNDQ/H7779rzqnPz/ft2qJPfvzxR/HLL7+If/75R2RkZIhXX31VGBsbi4MHDwoheF3v1I4dO4SPj48IDg4WkyZN0tzP69o4s2bNEl27dhU5OTma24ULFzSP87rWrdUHi6ioKDF+/HjNcVVVlXBzcxPz58/XYat048ZgoVarhYuLi1i4cKHmvqtXrwpTU1Px9ddfCyGEOHz4sAAgdu7cqTnnt99+EwqFQpw5c0YIIcTHH38s7OzsRFlZmeac6dOni4CAAM3xY489Ju69916t9kRHR4tnn3223m3RV7m5uQKASE5OFkLIdhsbG4v//e9/mnOOHDkiAIitW7cKIWTgMzAwEOfOndOcs2zZMmFjY6O5jtOmTRNdu3bVeq8hQ4aIAQMGaI5v9/Ndn7boOzs7O7FixQpe1ztUUFAgOnXqJDZs2CASExM1wYLXtfFmzZolQkJC6nyM1/XmWvVQSHl5OXbv3o3+/ftr7jMwMED//v2xdetWHbZMP2RmZuLcuXNa10epVCI6OlpzfbZu3QpbW1tERkZqzunfvz8MDAywfft2zTkJCQkwMTHRnDNgwABkZGTgypUrmnNqvk/1OdXvU5+26Ku8vDwAgL29PQBg9+7dqKio0PosXbp0gZeXl9Z17d69O5ydnTXnDBgwAPn5+Th06JDmnFtds/r8fNenLfqqqqoKa9asQVFREWJiYnhd79D48eNx77331vrsvK535ujRo3Bzc4Ofnx+GDRuG7OxsALyut9Kqg8XFixdRVVWl9T8NAJydnXHu3DkdtUp/VF+DW12fc+fOwcnJSetxIyMj2Nvba51T12vUfI+bnVPz8du1RR+p1WpMnjwZcXFx6NatGwD5WUxMTGBra6t17o2ft7HXLD8/HyUlJfX6+a5PW/TNgQMHYGVlBVNTU4wbNw7r169HUFAQr+sdWLNmDfbs2YP58+fXeozXtfGio6OxatUq/P7771i2bBkyMzPRq1cvFBQU8LreQosXISNqTcaPH4+DBw8iNTVV101pMwICArB3717k5eXh22+/xfDhw5GcnKzrZrVap06dwqRJk7BhwwaYmZnpujltyqBBgzR/Dg4ORnR0NLy9vfHNN9/A3Nxchy3Tb626x8LBwQGGhoa1Zr6eP38eLi4uOmqV/qi+Bre6Pi4uLsjNzdV6vLKyEpcvX9Y6p67XqPkeNzun5uO3a4u+mTBhAn7++Wds3rxZqyKvi4sLysvLcfXqVa3zb/y8jb1mNjY2MDc3r9fPd33aom9MTEzQsWNHREREYP78+QgJCcH777/P69pIu3fvRm5uLsLDw2FkZAQjIyMkJyfjgw8+gJGREZydnXldm4itrS06d+6MY8eO8ef1Flp1sDAxMUFERAQ2bdqkuU+tVmPTpk2IiYnRYcv0g6+vL1xcXLSuT35+PrZv3665PjExMbh69Sp2796tOeevv/6CWq1GdHS05pyUlBRUVFRoztmwYQMCAgJgZ2enOafm+1SfU/0+9WmLvhBCYMKECVi/fj3++usv+Pr6aj0eEREBY2Njrc+SkZGB7Oxsret64MABrdC2YcMG2NjYICgoSHPOra5ZfX6+69MWfadWq1FWVsbr2kj9+vXDgQMHsHfvXs0tMjISw4YN0/yZ17VpFBYW4vjx43B1deXP6620+HTRJrZmzRphamoqVq1aJQ4fPizGjh0rbG1ttWbhtmUFBQUiPT1dpKenCwDi3XffFenp6SIrK0sIIZd42traih9++EHs379f3H///XUuNw0LCxPbt28XqampolOnTlrLTa9evSqcnZ3FU089JQ4ePCjWrFkjLCwsai03NTIyEosWLRJHjhwRs2bNqnO56e3aog+ee+45oVQqRVJSktYys+LiYs0548aNE15eXuKvv/4Su3btEjExMSImJkbzePUys7vvvlvs3btX/P7778LR0bHOZWYvv/yyOHLkiPjoo4/qXGZ2u5/v27VFn7zyyisiOTlZZGZmiv3794tXXnlFKBQK8eeffwoheF2bSs1VIULwujbWSy+9JJKSkkRmZqZIS0sT/fv3Fw4ODiI3N1cIwet6M60+WAghxNKlS4WXl5cwMTERUVFRYtu2bbpuUovZvHmzAFDrNnz4cCGEXOb52muvCWdnZ2Fqair69esnMjIytF7j0qVLYujQocLKykrY2NiIkSNHioKCAq1z9u3bJ+Lj44Wpqalwd3cXCxYsqNWWb775RnTu3FmYmJiIrl27il9++UXr8fq0RR/UdT0BiJUrV2rOKSkpEc8//7yws7MTFhYW4sEHHxQ5OTlar3Py5EkxaNAgYW5uLhwcHMRLL70kKioqtM7ZvHmzCA0NFSYmJsLPz0/rPard7ue7Pm3RF6NGjRLe3t7CxMREODo6in79+mlChRC8rk3lxmDB69o4Q4YMEa6ursLExES4u7uLIUOGiGPHjmke53WtG6ubEhERUZNp1XMsiIiISL8wWBAREVGTYbAgIiKiJsNgQURERE2GwYKIiIiaDIMFERERNRkGCyIiImoyDBZERETUZBgsiKjZJSUlQaFQ1CqSRERtD4MFERERNRkGCyIiImoyDBZE7ci3336L7t27w9zcHB06dED//v1RVFQEAFixYgUCAwNhZmaGLl264OOPP9Z67o4dOxAWFgYzMzNERkZi/fr1UCgU2Lt3b6Pakpqail69esHc3Byenp544YUXNG0BAB8fH7z11lsYNWoUrK2t4eXlhU8//bTRn52IWgaDBVE7kZOTg6FDh2LUqFE4cuQIkpKS8NBDD0EIgdWrV+P111/HvHnzcOTIEbz11lt47bXX8MUXXwAACgsLcd999yEoKAi7d+/G7NmzMXXq1Ea35fjx4xg4cCAefvhh7N+/H2vXrkVqaiomTJigdd7ixYsRGRmJ9PR0PP/883juueeQkZFxR9eBiJqZTmqqElGL2717twAgTp48Wesxf39/8dVXX2ndN3fuXBETEyOEEGL58uWiQ4cOoqSkRPP4smXLBACRnp5+2/fevHmzACCuXLkihBBi9OjRYuzYsVrnqFQqYWBgoHkPb29v8eSTT2oeV6vVwsnJSSxbtqxen5eIdMNIx7mGiFpISEgI+vXrh+7du2PAgAG4++678cgjj8DExATHjx/H6NGjMWbMGM35lZWVUCqVAIAjR44gODgYZmZmmsdjYmIa3ZZ9+/Zh//79WL16teY+IQTUajUyMzMRGBgIAAgODtY8rlAo4OLigtzc3Ea/LxE1PwYLonbC0NAQGzZswJYtW/Dnn39i6dKlmDlzJn766ScAwGeffYbo6Ohaz2kOhYWFePbZZ/HCCy/UeszLy0vzZ2NjY63HFAoF1Gp1s7SJiJoGgwVRO6JQKBAXF4e4uDi8/vrr8Pb2RlpaGtzc3HDixAkMGzaszucFBgbiv//9L0pLSzW9Ftu2bWt0O8LDw3H48GF07Nix0a9BRPqJkzeJ2ont27fjrbfewq5du5CdnY1169bhwoULCAwMxJw5czB//nx88MEH+Oeff3DgwAGsXLkS7777LgDgiSeegEKhwJgxY3D48GH8+uuvWLRoUaPbMn36dGzZsgUTJkzA3r17cfToUfzwww+1Jm8SUevDHguidsLGxgYpKSlYsmQJ8vPz4e3tjcWLF2PQoEEAAAsLCyxcuBAvv/wyLC0t0b17d0yePBkAYGVlhZ9++gnjxo1DWFgYgoKC8Pbbb+Phhx9uVFuCg4ORnJyMmTNnolevXhBCwN/fH0OGDGmqj0tEOqIQQghdN4KIWp+TJ0/C19cX6enpCA0N1XVziEhPcCiEiIiImgyDBRHdsXHjxsHKyqrO27hx43TdPCJqQRwKIaI7lpubi/z8/Dofs7GxgZOTUwu3iIh0hcGCiIiImgyHQoiIiKjJMFgQERFRk2GwICIioibDYEFERERNhsGCiIiImgyDBRERETUZBgsiIiJqMgwWRERE1GT+H63R3tFRNrEkAAAAAElFTkSuQmCC", - "text/plain": [ - "
" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "decoding benchmark:\n", - " seq_len flash-attn-v2 bitdecoding-kc-4 bitdecoding-kc-2\n", - "0 1024.0 0.032768 0.024576 0.023552\n", - "1 2048.0 0.048128 0.029696 0.028672\n", - "2 4096.0 0.072704 0.039936 0.034816\n", - "3 8192.0 0.112640 0.060416 0.049152\n", - "4 16384.0 0.192512 0.087040 0.074752\n", - "5 32768.0 0.349184 0.141312 0.125952\n", - "6 65536.0 0.666624 0.246784 0.225280\n", - "7 131072.0 1.292288 0.459776 0.428032\n", - "8 262144.0 2.551808 0.873472 0.833536\n", - "9 524288.0 5.052416 1.715200 1.638400\n" - ] - } - ], - "source": [ - "@triton.testing.perf_report(\n", - " triton.testing.Benchmark(\n", - " x_names=[\"seq_len\"],\n", - " x_vals=[2**i for i in range(10, 20, 1)],\n", - " line_arg='provider', # Argument name whose value corresponds to a different line in the plot.\n", - " line_vals=['flash-attn-v2', 'bitdecoding-kc-4', 'bitdecoding-kc-2'], # Possible values for `line_arg`.\n", - " line_names=['flash-attn-v2', 'bitdecoding-kc-4', 'bitdecoding-kc-2'], # Label name for the lines.\n", - " styles=[('blue', '-'), ('red', '-'), ('green', '-')], # Line color and style.\n", - " plot_name=\"decoding benchmark\",\n", - " args={}, # Values for function arguments not in `x_names` and `y_name`.\n", - " )\n", - ")\n", - "def benchmark(seq_len, provider):\n", - " torch.random.manual_seed(0)\n", - " device = \"cuda\"\n", - " dtype = torch.float16\n", - "\n", - " batch_size = 1\n", - " nheads = 32\n", - " nheads_k = 32\n", - " d = 128\n", - " sm_scale = 1.0 / math.sqrt(d)\n", - "\n", - " # Quantization parameters\n", - " quant_mode = \"k-channel\"\n", - " group_size = 128\n", - "\n", - " q = torch.randn(batch_size, 1, nheads, d, device=device, dtype=dtype)\n", - " k_cache = torch.randn(batch_size, seq_len, nheads_k, d, device=device, dtype=dtype)\n", - " v_cache = torch.randn(batch_size, seq_len, nheads_k, d, device=device, dtype=dtype)\n", - "\n", - " quantiles = [0.5, 0.2, 0.8]\n", - "\n", - " if provider == 'flash-attn-v2':\n", - " ms, min_ms, max_ms = triton.testing.do_bench(lambda: flash_attn_with_kvcache(q, k_cache, v_cache), quantiles=quantiles)\n", - " elif provider == 'bitdecoding-kc-4':\n", - " num_bits = 4\n", - " pack_nums = 16 / num_bits\n", - " k_pack = torch.zeros((batch_size, int(seq_len // pack_nums), nheads_k, d), dtype=torch.uint16, device=device)\n", - " k_params = torch.zeros((batch_size, int(seq_len // group_size), nheads_k, d), dtype=torch.float32, device=device)\n", - " v_pack = torch.zeros((batch_size, seq_len, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device)\n", - " v_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seq_len), dtype=torch.float32, device=device)\n", - " ms, min_ms, max_ms = triton.testing.do_bench(lambda: fwd_kvcache_int(q,\n", - " k_pack, k_params, \n", - " v_pack, v_params,\n", - " None, # opt_block_table\n", - " sm_scale,\n", - " quant_mode, \n", - " group_size,\n", - " num_bits), quantiles=quantiles)\n", - " elif provider == 'bitdecoding-kc-2':\n", - " num_bits = 2\n", - " pack_nums = 16 / num_bits\n", - " k_pack = torch.zeros((batch_size, int(seq_len // pack_nums), nheads_k, d), dtype=torch.uint16, device=device)\n", - " k_params = torch.zeros((batch_size, int(seq_len // group_size), nheads_k, d), dtype=torch.float32, device=device)\n", - " v_pack = torch.zeros((batch_size, seq_len, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device)\n", - " v_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seq_len), dtype=torch.float32, device=device)\n", - " ms, min_ms, max_ms = triton.testing.do_bench(lambda: fwd_kvcache_int(q,\n", - " k_pack, k_params, \n", - " v_pack, v_params,\n", - " None, # opt_block_table\n", - " sm_scale,\n", - " quant_mode, \n", - " group_size,\n", - " num_bits), quantiles=quantiles)\n", - " \n", - " return ms, min_ms, max_ms\n", - "\n", - "benchmark.run(show_plots=True, print_data=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Batches" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "metadata": {}, - "outputs": [ - { - "data": { - "image/png": "iVBORw0KGgoAAAANSUhEUgAAAh8AAAGxCAYAAADCo9TSAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjEsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvc2/+5QAAAAlwSFlzAAAPYQAAD2EBqD+naQAAbtZJREFUeJzt3Xd409X+B/B3OtLddC+6J2WPMgqyBBkqVwW9XuQqKIJI2SJDBBkyHCiKCioKXEXl5xUceFERGYrsPdqyWspoaUtpuldyfn8ckja0hRbaJmnfr+fJY/PNN+lphObN+X7O5yiEEAJEREREDcTC2AMgIiKipoXhg4iIiBoUwwcRERE1KIYPIiIialAMH0RERNSgGD6IiIioQTF8EBERUYNi+CAiIqIGZWXsAdxKq9Xi6tWrcHJygkKhMPZwiIiIqAaEEMjNzYWfnx8sLG4/t2Fy4ePq1asICAgw9jCIiIjoLly6dAn+/v63PcfkwoeTkxMAOXhnZ2cjj4aIiIhqIicnBwEBAfrP8dsxufChu9Ti7OzM8EFERGRmalIywYJTIiIialAMH0RERNSgGD6IiIioQZlczUdNaTQalJaWGnsYRLC2toalpaWxh0FEZDbMLnwIIZCWlobs7GxjD4VIz8XFBT4+PuxNQ0RUA2YXPnTBw8vLC/b29vxlT0YlhEBBQQHS09MBAL6+vkYeERGR6TOr8KHRaPTBw93d3djDIQIA2NnZAQDS09Ph5eXFSzBERHdgVgWnuhoPe3t7I4+EyJDuzyTrkIiI7syswocOL7WQqeGfSSKimjPL8EFERETmi+GjgQghMGbMGLi5uUGhUMDFxQWTJ0+uk9dOTk6GQqHA0aNH6+T1iIiI6hPDRwP55ZdfsHbtWmzevBmpqalo1aqVsYdUI8HBwVi+fLnBsbVr18LFxaVBvv+OHTvwyCOPwNfXFw4ODmjXrh3Wr1/fIN+biIjqh1mtdjFn58+fh6+vL7p16wYAsLLiW18Tf//9N9q0aYMZM2bA29sbmzdvxjPPPAOVSoWHH37Y2MMjIqK7wJmPBjBy5EhMmDABKSkpUCgUCA4OrnTOF198gZiYGDg5OcHHxwdPPfWUvncEANy4cQPDhw+Hp6cn7OzsEBERgTVr1hi8xoULF9CnTx/Y29ujbdu22LNnz23Hdf78eTzyyCPw9vaGo6MjOnXqhN9//13/eO/evXHx4kVMmTIFCoUCCoUCO3bswLPPPgu1Wq0/Nm/ePABylmTx4sV47rnn4OTkhMDAQHzyySfVfv/ffvsNtra2lRrGTZo0Cffffz8A4JVXXsHChQvRrVs3hIWFYdKkSRg4cCA2btx425+NiIiqptEAhYXGHYPZhw8hgPx849yEqNkY33vvPSxYsAD+/v5ITU3FgQMHKp1TWlqKhQsX4tixY/j++++RnJyMkSNH6h+fM2cOTp8+jS1btiA+Ph4rV66Eh4eHwWvMnj0b06ZNw9GjRxEZGYlhw4ahrKys2nHl5eXhwQcfxLZt23DkyBEMHDgQgwcPRkpKCgBg48aN8Pf3x4IFC5CamorU1FR069YNy5cvh7Ozs/7YtGnT9K+5bNkyxMTE4MiRIxg3bhxefPFFJCYmVvn9+/btCxcXF3z33Xf6YxqNBhs2bMDw4cOrHbdarYabm1u1jxMRUdXy84E33wT++18jD0SYGLVaLQAItVpd6bHCwkJx+vRpUVhYqD+WlyeEjAENf8vLq/nP9e6774qgoCD9/V69eolJkyZVe/6BAwcEAJGbmyuEEGLw4MHi2WefrfLcpKQkAUCsXr1af+zUqVMCgIiPj6/5IIUQLVu2FCtWrNDfDwoKEu+++67BOWvWrBEqlarSc4OCgsS///1v/X2tViu8vLzEypUrq/1+kyZNEvfff7/+/q+//ipsbGzEjRs3qjx/w4YNQqlUipMnT9bsB2ogVf3ZJCIyFVqtEOfPCzFokPz88vERoppfs3ftdp/ftzL7mY/G4tChQxg8eDACAwPh5OSEXr16AYB+FuLFF1/EN998g3bt2mH69On4+++/K71GmzZt9F/r2nzrLt04Ojrqb2PHjgUgZz6mTZuG6OhouLi4wNHREfHx8frveTcqjkGhUMDHx0c/hkGDBunH0LJlSwDA8OHDsWPHDly9ehUAsH79ejz00ENVFrRu374dzz77LD799FP984mI6PZKS4EffgB69gS2bAEsLIDHHgOcnIw3JrOverS3B/LyjPe960J+fj4GDBiAAQMGYP369fD09ERKSgoGDBiAkpISAPKD++LFi/jf//6HrVu3om/fvoiLi8Pbb7+tfx1ra2v917qmV1qtFgAMluE6OzsDAKZNm4atW7fi7bffRnh4OOzs7PD444/rv+fdqDgG3Th0Y1i9ejUKb15o1J3XqVMnhIWF4ZtvvsGLL76ITZs2Ye3atZVed+fOnRg8eDDeffddPPPMM3c9PiKipiQ7G3j9dWD5clnr4esLzJ4NPPAAYMydIMw+fCgUgIODsUdxbxISEnD9+nUsXboUAQEBAICDBw9WOs/T0xMjRozAiBEj0KNHD7z88ssG4eN2wsPDKx3bvXs3Ro4cicceewyAnAlJTk42OEepVEKj0dzxWE00a9asyuPDhw/H+vXr4e/vDwsLCzz00EMGj+/YsQMPP/ww3njjDYwZM6bW35eIqKnRaoFTp4AxY4C9e+WxPn2Al18GWrQAqvl13GB42cUEBAYGQqlUYsWKFbhw4QJ+/PFHLFy40OCcuXPn4ocffsC5c+dw6tQpbN68GdHR0ff0fSMiIrBx40YcPXoUx44dw1NPPaWfpdAJDg7Grl27cOXKFWRmZuqP5eXlYdu2bcjMzERBQcE9jWP48OE4fPgwFi1ahMcffxw2Njb6x7Zv346HHnoIEydOxNChQ5GWloa0tDRkZWXd0/ckImqsioqAtWuBXr1k8LCxAaZNA956C+jcGQgKAozd7YHhwwR4enpi7dq1+Pbbb9GiRQssXbq00oyGUqnErFmz0KZNG/Ts2ROWlpb45ptv7un7vvPOO3B1dUW3bt0wePBgDBgwAB06dDA4Z8GCBUhOTkZYWBg8PT0BAN26dcPYsWPx5JNPwtPTE2+++eY9jSM8PBydO3fG8ePHK61yWbduHQoKCrBkyRL4+vrqb0OGDLmn70lE1Bhduwa88AIwahRw4wYQHg58/DHw/PNyxsNUNoRXCFHTBaMNIycnByqVCmq1Wl+boFNUVISkpCSEhITA1tbWSCMkqox/NonImDQa4O+/5WWWhAR5bMgQYMIEICIC8PGp/xqP231+38rsaz6IiIiasvx84MMPgXnzZPMwZ2dZ2/GPf8hLLCqVsUdYGcMHERGRGRICSE6Wsxs//yyPtW8PvPIK0KED4O8PKJVGHWK1GD6IiIjMTGmp7NkRFwdcvix7d4wcKWs7wsIAT0+5GtRUMXwQERGZkexsYNEi2bujrAzw9pa9O/r2BQIDAUdHY4/wzhg+iIiIzIBWC5w+LYtKdfuG9uoFzJhR3rvD2Etoa8pMhklERNR0FRUB33wDvPQSkJUlaznGjweeegoIDgbc3Ez7MsutGD6IiIhMWHq6nN3Q7TwRGgrMmQPExsrLLHZ2Rh3eXWH4ICIiMkEajexQOmaMvNwCAI8+Ckya1HC9O+oLO5w2kN69e2Py5MnVPh4cHIzly5fX+nXnzZuHdu3a3fW46opCocD3338PAEhOToZCoTDYzK6+3Ol9JSIyR/n5wDvvAP37y+Dh5AQsWCB7ebRrJ+s7zDV4AAwfJuPAgQMGm6ZV/DA3NwEBAUhNTUWrVq2MPZR7UlxcjHbt2jVYkCIiEgJISgKGDQOmTwcKCoC2bYHPPpPHmjcHXFyMPcp7x8suJkK3b0pjYGlpCR8fH2MP455Nnz4dfn5+OHbsmLGHQkRNQGkp8OuvwLhxwKVLsoB0xAhg9GjZu8PLy7yKSm+HMx8NqKysDOPHj4dKpYKHhwfmzJkD3dY6FS+7BAcHAwAee+wxKBQK/X0AWLp0Kby9veHk5IRRo0ahqKio0vdZvXo1oqOjYWtri+bNm+Ojjz4yePzy5csYNmwY3Nzc4ODggJiYGOzbt0//+MqVKxEWFgalUomoqCh88cUXBs8/e/YsevbsCVtbW7Ro0QJbt241ePzWyy47duyAQqHAtm3bEBMTA3t7e3Tr1g2JiYkGz3v99dfh5eUFJycnPP/885g5c2atLyn9/PPPUKlUWL9+vf7Y559/jpYtW8LGxga+vr4YP378HV9ny5Yt+O233ypt8EdEVB/UallEOmSIDB5eXvKyy7RpQOvWspdHYwkeAABhYtRqtQAg1Gp1pccKCwvF6dOnRWFhoRFGdm969eolHB0dxaRJk0RCQoL48ssvhb29vfjkk0+EEEIEBQWJd999VwghRHp6ugAg1qxZI1JTU0V6eroQQogNGzYIGxsbsXr1apGQkCBmz54tnJycRNu2bfXf58svvxS+vr7iu+++ExcuXBDfffedcHNzE2vXrhVCCJGbmytCQ0NFjx49xJ9//inOnj0rNmzYIP7++28hhBAbN24U1tbW4sMPPxSJiYli2bJlwtLSUvzxxx9CCCE0Go1o1aqV6Nu3rzh69KjYuXOnaN++vQAgNm3aJIQQIikpSQAQR44cEUIIsX37dgFAdOnSRezYsUOcOnVK9OjRQ3Tr1s1g3La2tuLzzz8XiYmJYv78+cLZ2dngZ6vufZ00aZIQQoj169cLJycn8dNPP+kf/+ijj4Stra1Yvny5SExMFPv379e/z9VJS0sTzZo1EwcOHKj0s1THnP9sEpHxaDRCnDghRPfuQsiLLkL06CHE5s1CJCcLUVJi7BHW3O0+v29Vq/CxePFiERMTIxwdHYWnp6d45JFHREJCgsE5vXr1EgAMbi+88EKdDL7KX/BarRB5eca5abU1/rl69eoloqOjhbbCc2bMmCGio6OFEIbhQwhh8GGuExsbK8aNG2dwrEuXLgYf0GFhYeKrr74yOGfhwoUiNjZWCCHExx9/LJycnMT169erHGe3bt3E6NGjDY498cQT4sEHHxRCCPHrr78KKysrceXKFf3jW7ZsqVH4+P333/XP+fnnnwUA/f/LLl26iLi4OIPv27179xqHjw8++ECoVCqxY8cOg8f9/PzE7Nmzb/saFWm1WjFw4ECxcOHCKn+W6jB8EFFtFRYKsW6dEO7uMnQolUJMnizEwYNCZGTU6iPGJNQmfNTqssvOnTsRFxeHvXv3YuvWrSgtLUX//v2Rn59vcN7o0aORmpqqv7355pt3PzVzJwUFspesMW4FBbUaateuXaGoMG8WGxuLs2fPQqPR1Oj58fHx6NKli8Gx2NhY/df5+fk4f/48Ro0aBUdHR/3t9ddfx/nz5wEAR48eRfv27eHm5lbt9+jevbvBse7duyM+Pl7/eEBAAPz8/Kocw+20adNG/7Wvry8AID09HQCQmJiIzp07G5xf8f6ff/5p8DNVvKzy3//+F1OmTMHWrVvRq1cv/fH09HRcvXoVffv2rXI8Y8eONXhNAFixYgVyc3Mxa9asGv1MRER3Iz1d7ssyciRw/ToQEgKsWgW88AIQHQ14eDSyyyy3qFXB6S+//GJwf+3atfDy8sKhQ4fQs2dP/XF7e/tGUXBobvLy8gAAn376aaWQYnlzTZadEbvRWFtb67/WhTCtVluj58bExBisOPH29tZ/3b59exw+fBiff/45YmJi9K99p591wYIFmDZtmsGxP/74A3v27IGNjU2l7z98+HCsW7euRuMlIqqKRgPs2yd7d5w6JY/94x/A5Mmyd4evr3kvoa2pe1rtolarAaDSv6LXr1+PL7/8Ej4+Phg8eDDmzJkDe3v7Kl+juLgYxcXF+vs5OTm1G4S9PXDzQ7fBVfMzVadiUScA7N27FxEREfpgUJG1tXWlGZHo6Gjs27cPzzzzjMFr6Hh7e8PPzw8XLlzA8OHDqxxDmzZtsHr1amRlZVU5+xEdHY3du3djxIgR+mO7d+9GixYt9I9funQJqamp+tmLimO4W1FRUThw4IDBz3bgwAH913Z2dggPD6/yuWFhYVi2bBl69+4NS0tLfPDBBwAAJycnBAcHY9u2bejTp0+l53l5ecHLy8vg2Pvvv4/XX39df//q1asYMGAANmzYUCnQERHVRkGBnN2YO1f28XByAqZOBR57DAgKahxLaGvqrsOHVqvF5MmT0b17d4N+Dk899RSCgoLg5+eH48ePY8aMGUhMTMTGjRurfJ0lS5Zg/vz5dzsMOS/l4HD3z29AKSkpmDp1Kl544QUcPnwYK1aswLJly6o8V/eh2b17d9jY2MDV1RWTJk3CyJEjERMTg+7du2P9+vU4deoUQkND9c+bP38+Jk6cCJVKhYEDB6K4uBgHDx7EjRs3MHXqVAwbNgyLFy/Go48+iiVLlsDX1xdHjhyBn58fYmNj8fLLL+Of//wn2rdvj379+uGnn37Cxo0b8fvvvwMA+vXrh8jISIwYMQJvvfUWcnJyMHv27Ht+byZMmIDRo0cjJiYG3bp1w4YNG3D8+HGDn+12IiMjsX37dvTu3RtWVlb6lUPz5s3D2LFj4eXlhUGDBiE3Nxe7d+/GhAkTqnydwMBAg/u6yzFhYWHw9/e/+x+QiJosIYCLF+Xsxg8/yGNt2sidaDt0AAICgFsmWxu/uy0sGTt2rAgKChKXLl267Xnbtm0TAMS5c+eqfLyoqEio1Wr97dKlS412tcu4cePE2LFjhbOzs3B1dRWvvPKKvgD11oLTH3/8UYSHhwsrKysRFBSkP75o0SLh4eEhHB0dxYgRI8T06dMrFWWuX79etGvXTiiVSuHq6ip69uwpNm7cqH88OTlZDB06VDg7Owt7e3sRExMj9u3bp3/8o48+EqGhocLa2lpERkaK//znPwavn5iYKO677z6hVCpFZGSk+OWXX2pUcHrjxg39axw5ckQAEElJSfpjCxYs0P9szz33nJg4caLo2rXrHd9X3WoXIYQ4ffq08PLyElOnTtUfW7VqlYiKihLW1tbC19dXTJgw4bavWRELTonoXpSUCPHzz0IEBcmiUoVCiGeeEWL3biHS0uRql8aiNgWnCiFuNpqohfHjx+OHH37Arl27EBIScttz8/Pz4ejoiF9++QUDBgy442vn5ORApVJBrVbD2dnZ4LGioiIkJSUhJCQEtra2tR02mZkHHngAPj4+lfqMmCL+2SSiW6nVwNKlwLJlsoGYpyfwyivAAw/IDeGcnIw9wrp1u8/vW9XqsosQAhMmTMCmTZuwY8eOOwYPAPoiQV19AFFVCgoKsGrVKgwYMACWlpb4+uuv8fvvv1dqYEZEZOq0WiA+Hhg7FvjrL3nsvvvkzrStWsl9WSrU3zdJtQofcXFx+Oqrr/DDDz/AyckJaWlpAACVSgU7OzucP38eX331FR588EG4u7vj+PHjmDJlCnr27GmwzJLoVgqFAv/73/+waNEiFBUVISoqCt999x369etn7KEREdVYURHw7beykDQzU4aMceOAf/8bCA4G3N0b9xLamqpV+Fi5ciUAuZNoRWvWrMHIkSOhVCrx+++/Y/ny5cjPz0dAQACGDh2KV199tc4GTI2TnZ2dvqiViMgcZWTIyyqffSaLTIODgVdfBbp3l5dZarlAslGr9WWX2wkICMDOnTvvaUBERETmRKMBDhyQvTtOnJDHHn4YmDIFiIxsOr07aoO72hIREd2lggLgk0/kpnB5ebL5dcXeHa6uxh6haWL4ICIiqiUhgJQUObuxaZM81qqVvMzSsWMT7d1RCwwfREREtVBaCmzdKgtJL16UBaRPPQW8+CIQFgZ4eQEWtdo5relh+CAiIqohtRp4803grbdkCPHwAGbNAvr3l5dZGlvvjvrC8EFERHQHWi2QkCBnN3btkse6dQNmzgRat2bvjtrixFAD6d27NyZPnlzt48HBwfr9SGpj3rx5aNeu3V2Pq64oFAp8//33AIDk5GQoFAqDXWjry53eVyKie1VUBHzzDdC7twweVlbA+PHAe+8BsbFyxoPBo3YYPkzEgQMHMGbMGP39ih/m5iYgIACpqakGGw6ai+TkZIwaNQohISGws7NDWFgYXnvtNZSUlBh7aERkBBkZwKRJwNNPy68DA4GVK2W9R4sW8rILm4bVHi+7mAhPT09jD6HOWFpawsfHx9jDuCsJCQnQarX4+OOPER4ejpMnT2L06NHIz8/H22+/bezhEVED0WiAgwdl747jx+WxBx+Uy2ijoti7415x5qMBlZWVYfz48VCpVPDw8MCcOXP0jdsqXnYJDg4GADz22GNQKBT6+wCwdOlSeHt7w8nJCaNGjUJRUVGl77N69WpER0fD1tYWzZs3x0cffWTw+OXLlzFs2DC4ubnBwcEBMTEx2Ldvn/7xlStXIiwsDEqlElFRUZU2djt79ix69uwJW1tbtGjRotL+K7dedtmxYwcUCgW2bduGmJgY2Nvbo1u3bkhMTDR43uuvvw4vLy84OTnh+eefx8yZM2t9Sennn3+GSqXC+vXr9cc+//xztGzZEjY2NvD19cX48eOrff7AgQOxZs0a9O/fH6GhofjHP/6BadOmYePGjbUaBxGZr4IC4MMP5QZwx48DDg5yCe2iRUCHDoC/P4PHvWL4aEDr1q2DlZUV9u/fj/feew/vvPMOVq9eXem8AwcOAJBt61NTU/X3/+///g/z5s3D4sWLcfDgQfj6+lYKFuvXr8fcuXOxaNEixMfHY/HixZgzZw7WrVsHAMjLy0OvXr1w5coV/Pjjjzh27BimT58OrVYLANi0aRMmTZqEl156CSdPnsQLL7yAZ599Ftu3bwcAaLVaDBkyBEqlEvv27cOqVaswY8aMGv38s2fPxrJly3Dw4EFYWVnhueeeMxj3okWL8MYbb+DQoUMIDAzUt/Ovqa+++grDhg3D+vXrMXz4cAAySMXFxWHMmDE4ceIEfvzxR4SHh9fqddVqNdzc3Gr1HCIyP7reHc88Iy+15OYCLVsCq1fLY82bs2lYnREmRq1WCwBCrVZXeqywsFCcPn1aFBYW6o9ptVqRV5xnlJtWq63xz9WrVy8RHR1t8JwZM2aI6OhoIYQQQUFB4t1339U/BkBs2rTJ4DViY2PFuHHjDI516dJFtG3bVn8/LCxMfPXVVwbnLFy4UMTGxgohhPj444+Fk5OTuH79epXj7Natmxg9erTBsSeeeEI8+OCDQgghfv31V2FlZSWuXLmif3zLli0G401KShIAxJEjR4QQQmzfvl0AEL///rv+OT///LMAoP9/2aVLFxEXF2fwfbt3727ws1WlV69eYtKkSeKDDz4QKpVK7Nixw+BxPz8/MXv27Nu+xu2cPXtWODs7i08++eS251X1Z5OIzEdJiRBbtggREiKEjCFCPPWUEH/+KURqqhAajbFHaPpu9/l9K7Ov+SgoLYDjEkejfO+8WXlwUDrU+PyuXbtCUaEyKTY2FsuWLYNGo6nR8+Pj4zF27FiDY7GxsfpZifz8fJw/fx6jRo3C6NGj9eeUlZVBpVIBAI4ePYr27dtX+y/5+Ph4g8JXAOjevTvee+89/eMBAQHw8/MzGENNVNzZ2NfXFwCQnp6OwMBAJCYmYty4cQbnd+7cGX/88QcA4M8//8SgQYP0j3388cf62Y3//ve/SE9Px+7du9GpUyf9Oenp6bh69Sr69u1b5XjGjh2LL7/8Un8/Ly/P4PErV65g4MCBeOKJJwzeTyJqXNRqYNky4I03gJISufPszJnAwIGywNTZ2dgjbHzMPnxQOd2H56effoouXboYPGZ58wKlnZ1dg49Lx7rCWjRdCNNd7rmTmJgYg6W73t7e+q/bt2+Pw4cP4/PPP0dMTIz+te/0sy5YsADTpk2r8rGrV6+iT58+6NatGz755JMajZGIzItWCyQmypUrO3bIY127yqZhut4dSqVRh9homX34sLe2R96svDufWE/fuzYqFnUCwN69exEREaEPBhVZW1tXmhGJjo7Gvn378Mwzzxi8ho63tzf8/Pxw4cIF/azArdq0aYPVq1cjKyurytmP6Oho7N69GyNGjNAf2717N1q0aKF//NKlS0hNTdXPXlQcw92KiorCgQMHDH42Xa0LIINEdbUaYWFhWLZsGXr37g1LS0t88MEHAAAnJycEBwdj27Zt6NOnT6XneXl5wcvLq9LxK1euoE+fPujYsSPWrFkDC/ZJJmp0iorkniyTJwPp6bJ3x5gxwIgRQEgIl9DWN7MPHwqFolaXPowpJSUFU6dOxQsvvIDDhw9jxYoVWLZsWZXn6j40u3fvDhsbG7i6umLSpEkYOXIkYmJi0L17d6xfvx6nTp1CaGio/nnz58/HxIkToVKpMHDgQBQXF+PgwYO4ceMGpk6dimHDhmHx4sV49NFHsWTJEvj6+uLIkSPw8/NDbGwsXn75Zfzzn/9E+/bt0a9fP/z000/YuHEjfv/9dwBAv379EBkZiREjRuCtt95CTk4OZs+efc/vzYQJEzB69GjExMSgW7du2LBhA44fP27ws91OZGQktm/fjt69e8PKykq/cmjevHkYO3YsvLy8MGjQIOTm5mL37t2YMGFCla9z5coV9O7dG0FBQXj77beRkZGhf8xclw8TkaHMTLkL7SefyNmPgABg9mygRw/ZMMzBPD5SzFsD1KDUSm0LTs1Fr169xLhx48TYsWOFs7OzcHV1Fa+88oq+APXWgtMff/xRhIeHCysrKxEUFKQ/vmjRIuHh4SEcHR3FiBEjxPTp0ysVZa5fv160a9dOKJVK4erqKnr27Ck2btyofzw5OVkMHTpUODs7C3t7exETEyP27dunf/yjjz4SoaGhwtraWkRGRor//Oc/Bq+fmJgo7rvvPqFUKkVkZKT45ZdfalRweuPGDf1rHDlyRAAQSUlJ+mMLFizQ/2zPPfecmDhxoujatesd39dJkybp758+fVp4eXmJqVOn6o+tWrVKREVFCWtra+Hr6ysmTJhQ7eutWbNGAKjydjvm/GeTqKkoKxNi/34h2rUrLyodOFCI334TIiVFiNJSY4/QvNWm4FQhxM1GEyYiJycHKpUKarUazrdU+RQVFSEpKQkhISGwtbU10gipoTzwwAPw8fGp1GfEFPHPJpFpKygAPvtM9uvIyQHs7eUll8cfB4KDuYS2Ltzu8/tWZn/ZhRqHgoICrFq1CgMGDIClpSW+/vpr/P7775UamBER1YYQwOXLwEsvAd9+K49FR8sQ0qmTvOTCfy80PIYPMgkKhQL/+9//sGjRIhQVFSEqKgrfffcd+vXrZ+yhEZGZKi0Ftm+XO9FeuCCPDRsmV7eEhQHe3gDryY2D4YNMgp2dnb6olYjoXqnVwDvvyN4dxcWAm1t5746gIPbuMDaGDyIiajS0WuDMGSAuDrjZoxBdupT37vD3Z+8OU8DwQUREjUJREfDjj8DEicC1a7J3x+jRwMiR7N1haswyfJjYAh0i/pkkMrLMTGDePGDVKkCjkTMcr77K3h2myqzCh649d0FBgVHbhBPdqqCgAIBhC3kiqn8aDXD4MDB2rPwvAPTvD0ybJneh9fWVMyBkWszqf4mlpSVcXFyQnp4OALC3tzfYqI2ooQkhUFBQgPT0dLi4uFTZKp+I6kdBAbBmDfDKK7J3h52d7N0xdKi8zFLN/plkAswqfADlLa51AYTIFLi4uLD9OlEDEQK4ckXObmzYII9FRcmW6Z06yZ1o2bvDtJld+FAoFPD19YWXlxdKS0uNPRwiWFtbc8aDqIGUlgI7d8reHefOyWNPPilXt4SFAT4+7N1hDswufOhYWlryFz4RUROiVgPvvQcsWSJXtri5ATNmlPfuUKmMPUKqKbMNH0RE1DRotXKWIy4O0PUi7NRJ9u5o25a9O8wRwwcREZmsoiLgp59k7460NMDSEhg1CnjuOVlU6unJ3h3miOGDiIhMUmYmMH8+sHKlXFLr5yd7d/Tqxd4d5o7hg4iITIpGAxw5Int3HDokj/XrJ1e3REfLEMLeHeaN//uIiMhkFBQA//mP3AROrZa9OyZNAh5/nL07GhOGDyIiMjohgKtXgZdfBr7+Wh6LjJSXWTp3lr072Ni68WD4ICIioyotBXbtkr07zp6Vx554Ahg/vrx3BzsrNC4MH0REZDQ5OcD77wOLFsmVLS4usnfHgw+yd0djxvBBREQNTqsFzp+Xsxu//SaPxcTI3h3t2rF3R2PH8EFERA2quLi8d0dqqryk8txz8hYayt4dTQHDBxERNZjr14EFC4APPyzv3TF7NtC7tywqdXQ09gipITB8EBFRvdNogKNHZVHpgQPyWN++5b07mjVj746mhP+riYioXhUUAF98IXt3ZGfL7e4nTgT++U8gOBhwdzf2CKmhMXwQEVG90PXumDEDWL9eHouIAObMYe+Opo7hg4iI6lxpKfDnn/Iyy5kz8tjjj8vVLeHh7N3R1DF8EBFRncrNBVasABYulL07VCpg+nTgoYfkZRb27iCGDyIiqhO63h0TJwK//CKPdegAvPJKee8OGxujDpFMBMMHERHds+Ji4OefZfC4cgWwsACefRZ4/nm5IZyXF3t3UDmGDyIiuifXr8v26O+/L5fU+vrK2Y4+fWSLdPbuoFsxfBAR0V3RaIDjx2VR6b598lifPnJn2hYtZAMxa2vjjpFME8MHERHVWkGBXD47YwZw44as5ZgwAfjXv2RRqZsbL7NQ9Rg+iIioxoSQ+7HMnCkbhwFy6eyrrwJdusjeHfb2xh0jmT6GDyIiqpHSUuCvv4Bx44CEBHlsyBA54xERwd4dVHMWtTl5yZIl6NSpE5ycnODl5YVHH30UiYmJBucUFRUhLi4O7u7ucHR0xNChQ3Ht2rU6HTQRETWs3Fxg2TLZqyMhAXB2ln085s6Vy2ibNWPwoJqrVfjYuXMn4uLisHfvXmzduhWlpaXo378/8vPz9edMmTIFP/30E7799lvs3LkTV69exZAhQ+p84EREVP+0WuDcOVnLMWsWUFgItG8PfPaZPNa8OeDiYuxRkrlRCCHE3T45IyMDXl5e2LlzJ3r27Am1Wg1PT0989dVXePzxxwEACQkJiI6Oxp49e9C1a9c7vmZOTg5UKhXUajWcnZ3vdmhERHSPiouBLVvkZZXLl2XvjpEjgVGjgLAwwNNTHiMCavf5fU9/bNRqNQDAzc0NAHDo0CGUlpaiX79++nOaN2+OwMBA7Nmz516+FRERNaDr12WvjieekMHD2xtYvhyYOhVo3VreZ/Cgu3XXBadarRaTJ09G9+7d0apVKwBAWloalEolXG6Zg/P29kZaWlqVr1NcXIzi4mL9/ZycnLsdEhER3SONBjhxQhaV6v7N2Lt3ee+OZs3Yu4Pu3V2Hj7i4OJw8eRJ//fXXPQ1gyZIlmD9//j29BhER3buCAuCbb2TQyMoClEq5C+2wYbJ3h7s7e3dQ3birSbPx48dj8+bN2L59O/z9/fXHfXx8UFJSguzsbIPzr127Bh8fnypfa9asWVCr1frbpUuX7mZIRER0l3S9O8aNk3uxZGUBoaHAqlXA6NFyxsPDg8GD6k6tZj6EEJgwYQI2bdqEHTt2ICQkxODxjh07wtraGtu2bcPQoUMBAImJiUhJSUFsbGyVr2ljYwMbbnNIRGQUpaXy8srYsUB8vDz26KNyg7iICLlPC5fQUl2rVfiIi4vDV199hR9++AFOTk76Og6VSgU7OzuoVCqMGjUKU6dOhZubG5ydnTFhwgTExsbWaKULERE1nNxcYOVKYP58ecnFyUlecvnHP+SGcFxCS/WlVkttFdXMua1ZswYjR44EIJuMvfTSS/j6669RXFyMAQMG4KOPPqr2ssutuNSWiKh+abVAUhIweTKwebM81rYtMHu27OERECD3aiGqjdp8ft9Tn4/6wPBBRFR/iouBX36RvTsuXZJ1HCNGyNqO0FDAy4tLaOnu1Obzm3u7EBE1EVlZwNKlsl9HaakMGq+8AvTtKy+zODkZe4TUVDB8EBE1choNcPKkXM3y99/yWI8ecmda9u4gY2D4ICJqxAoKgP/7P2DaNNm1VKmUIWT4cPbuIONh+CAiaoSEANLSZBHp2rXyfkgIMGcOEBsLBAYC9vbGHiU1VQwfRESNTGkpsHcv8OKLwKlT8tgjj8jeHZGR7N1BxsfwQUTUiOTmAh9/DMybB+TnyyLSl16S4SMoCHB1NfYIiRg+iIgaBa0WSE6WvTt++kkea9PGsHeHra0xR0hUjuGDiMjMFRcDv/0mN4FLSZEFpE8/DYwZA4SFsXcHmR6GDyIiM5aVBbzxBvDuu7LWw9MTmDULeOABWVTKXo1kihg+iIjMkEYji0nj4oC//pLH7rsPmDEDaNkS8Pdn7w4yXQwfRERmpqAA+O9/ZSFpZqYMGRV7d3h4sHcHmTaGDyIiMyEEcO0a8OqrwOefy/vBwbKotFs3uZrFwcHYoyS6M4YPIiIzUFoK7NsnZzhOnJDHHn5Yrm7R9e6w4m90MhP8o0pEZOJyc4FPPwXmzpW9Oxwd5SWXRx9l7w4yTwwfREQmSte746WXgO+/l8datZKXXTp0YO8OMl8MH0REJqi4GPj9d7ma5eJFWUA6fDgwdix7d5D5Y/ggIjIxWVnA228Dy5YBJSVy9cqsWUD//uzdQY0DwwcRkYnQaIDTp2Wn0l275LHu3Q17dyiVxh0jUV1g+CAiMgEFBcB338n6jowMuXJl7FjgmWfYu4MaH4YPIiIj0vXumDsX+OwzWWQaGAjMmcPeHdR4MXwQERlJaSmwf7/s3XH8uDz24IPAlClAVBR7d1DjxT/WRERGkJsrZzrmzpVfOzjI0DFkiJztcHMz9giJ6g/DBxFRA9Jq5dLZl14CNm2Sx1q2lC3SO3aUl1zYu4MaO4YPIqIGUlwMbNsmV7MkJcljut4doaGAjw97d1DTwPBBRNQAsrJk346335a9O9zdZe+OBx6Qq1nYu4OaEoYPIqJ6pOvdMXEisGOHPNa1KzBzJtC6NXt3UNPE8EFEVE8KCoCNG2V9R3q6XLkyZozs3RESAnh6sncHNU0MH0REdUwIGTZee03uRqvVyk3gXn0VuO8+9u4gYvggIqpDpaXAwYOyd8fRo/LYwIHA1Kmyd4efH3t3EPGvABFRHcnNBdaskd1Jc3IAe3tg8mRg6FBZVMreHUQSwwcR0T3S9e6YPh3473/lsehoeZklJkZecrGzM+4YiUwJwwcR0T0oLgb++EP27rhwQR4bNkxedmHvDqKqMXwQEd2lrCxg+XLgzTdlCHFzk0toBwyQRaUqlbFHSGSaGD6IiGpJowHi42Xvju3b5bEuXWTwaNOGvTuI7oThg4ioFgoKgB9+kJvAXbsmV66MHg2MGMHeHUQ1xfBBRFQDut4d8+cDn3wiZz/8/WVRaY8eckM4R0djj5LIPDB8EBHdQWkpcOiQLCI9ckQe699fdi6NigKaNWPvDqLa4F8XIqLbyM0F1q2TMxxqtVwyW7F3h7u7sUdIZH4YPoiIqqDVAikpsnfHt9/KY82bywZiHTvKyyzs3UF0dxg+iIhuUVwsd6AdPx44d04ee/JJedklLEz27rC0NOoQicwawwcRUQVZWcB778neHUVFsnfH9Olyf5bgYPbuIKoLDB9ERJCrVxITZe+Obdvksc6dDXt32NgYd4xEjQXDBxE1eQUFwI8/yt4daWnyksrzzwPPPitnO7y82LuDqC4xfBBRkyUEkJEBLFgArFolZz/8/IDZs4FevWSLdPbuIKp7DB9E1CSVlgKHDwNxcbKHBwD06wdMmyZXtbB3B1H94V8tImpycnOBL74AXnmlvHfHpEnA44/LyyxubrzMQlSfGD6IqMnQaoFLl2QR6TffyGORkbKBWKdO8jILe3cQ1T+GDyJqEoqLgZ07Ze+Os2flsX/+U/buCA9n7w6ihsTwQUSNXlYW8MEHwNKlQGEh4OICzJgBDBokZztcXIw9QqKmheGDiBotjQY4c0bWc2zdKo/FxACzZsneHQEB7N1BZAwMH0TUKBUUAJs3y03gUlPlJZVRo4CRI4HQUPbuIDImhg8ialR0vTtefx346KPKvTsCAwEnJ2OPkqhps6jtE3bt2oXBgwfDz88PCoUC33//vcHjI0eOhEKhMLgNHDiwrsZLRFSt0lLg4EHgoYeAFStk8OjbF1i9Wh6LjGTwIDIFtZ75yM/PR9u2bfHcc89hyJAhVZ4zcOBArFmzRn/fhhdViaie5eYC69fLeo7sbMDWVu7T8sQTsneHuzsvsxCZilqHj0GDBmHQoEG3PcfGxgY+Pj53PSgioprSaoHLl2Xo+OoreSwiQvbu6NxZXmaxtzfuGInIUK0vu9TEjh074OXlhaioKLz44ou4fv16fXwbImriiouBP/4AHnigPHg8/jjwySfA/ffLEMLgQWR66rzgdODAgRgyZAhCQkJw/vx5vPLKKxg0aBD27NkDyyo6+BQXF6O4uFh/Pycnp66HRESNUHa2rOtYskT27lCpgOnTZW0He3cQmbY6Dx//+te/9F+3bt0abdq0QVhYGHbs2IG+fftWOn/JkiWYP39+XQ+DiBopjUZ2KJ08Gfj1V3msY0d52aVtW/buIDIH9XLZpaLQ0FB4eHjg3LlzVT4+a9YsqNVq/e3SpUv1PSQiMlMFBcCmTXL32V9/BSwsgOeflzMg990n+3cweBCZvnrv83H58mVcv34dvr6+VT5uY2PD1TBEdFu63h2LF8s26RoN4Osre3f07s3eHUTmptbhIy8vz2AWIykpCUePHoWbmxvc3Nwwf/58DB06FD4+Pjh//jymT5+O8PBwDBgwoE4HTkRNQ2kpcPSo3BBu/355rE8f4OWXgehooFkzwNraqEMkolqqdfg4ePAg+vTpo78/depUAMCIESOwcuVKHD9+HOvWrUN2djb8/PzQv39/LFy4kLMbRFRrubnA118DM2cCN27ISyoTJgBPPsneHUTmTCGEEMYeREU5OTlQqVRQq9VwdnY29nCIyAi0WuDKFVlEun69PBYeLnt3dOnC3h1Epqg2n9/c24WITEpxMfDXX/IyS0KCPDZkiJzxCA+XtR5VrNonIjPC8EFEJiM7W24G9/rrsneHs7Os7Xj4Ydm7w9XV2CMkorrA8EFERqfr3TF1KrBlizzWvj3wyitAu3bs3UHU2DB8EJFRFRTIwDF5styjxcICGDkSeO45ICwM8PKSx4io8WD4ICKj0PXuWLpUNgkrKwO8vWXvjj59ZFEpa86JGieGDyJqcKWlwLFjsoh07155rHfv8t4d/v7s3UHUmDF8EFGDys0FNmwAZswAsrIApVKubBk2TBaVeniwdwdRY8fwQUQNQqsFrl6Vl1W++EJedgkNLe/dERzM3h1ETQXDBxHVu+JiYPduOcMRHy+PPfqovOwSGcneHURNDcMHEdWr7Gxg1Spg4UK5ssXZGZg2DRg8mL07iJoqhg8iqhcaDXDuHPDSS8DPP8tjbdvKyy663h22tkYdIhEZCcMHEdW5ggLgl1+AKVOAlBRZQDpiBPD88+zdQUQMH0RUh4QAMjOBN94A3n9fLqn18pKdSvv2Ze8OIpIYPoioTpSWAsePyyLSPXvksZ49genTy3t3KJXGHSMRmQaGDyK6Z3l5wPr1coZD17tj3DjgqafkElr27iCiihg+iOiuCSFrOmbMkI3DAFnTMXu27N0RFAQ4OBh3jERkehg+iOiulJQA27bJ3h0XLshjjz8OxMXJAOLrC1jxNwwRVYG/Goio1rKzZVHpO+/IEOLmJms7Bg6Usx0uLsYeIRGZMoYPIqoxrRY4eVLWc+zeLY916yYvu7RqxaJSIqoZhg8iqpGiIuDLLw03hBs7Fnj6aW4IR0S1w/BBRHd09apsif711/K+rqg0Nlb27uCGcERUGwwfRFStsjJgxw45w3H+vDw2dKjs5RERAXh7c0M4Iqo9hg8iqlJeHrBkCfDWW7KBmK6o9OGH5b4s7FRKRHeL4YOIDAght70fM6a8qDQ2Fpg1C2jTBvDzA6ytjTtGIjJvDB9EpFdaKotKX3oJuHFDFpW+8ALw7LOyqNTNzdgjJKLGgOGDiAAAGRlyF9r16+X90FBgzhzgvvvkElpbW+OOj4gaD4YPoiZOqwV27QJGjwbOnZPHhg4FJk0CoqIAT08uoSWiusXwQdSEFRXJotKlS8s7lc6YIYtKuS8LEdUXhg+iJursWeC554C//pL3u3aVl1natJH7snAJLRHVF4YPoiZGowG++kpeVtEVlY4dK4NIcDCgUhl7hETU2DF8EDUhN27I0PHFF/J+WBgwdy7QowfQrBn3ZSGihsHwQdRE7N4NjBxpWFQ6dSoQGQm4u7OolIgaDsMHUSNXWiqLShctMiwqffRR2anUzs7YIySipobhg6gRS0qSsx27dsn7XbvKyyzt2wNeXoCFhVGHR0RNFMMHUSMkBLBhAxAXB2RllReVjhkjd6F1cjL2CImoKWP4IGpkcnKAiROBdevk/dBQYN48oHdvuYTWin/ricjI+GuIqBHZuxd4+mnDotKXX5ZFpa6uxh0bEZEOwwdRI6DRyC6lCxbIolJXV7kL7dChcgmtjY2xR0hEVI7hg8jMpaQAzzwD7Nwp73ftKi+zdOgAeHhwCS0RmR6GDyIz9n//B7z4YnlR6YsvAi+8IPdlsbc39uiIiKrG8EFkhvLzgQkTgDVr5P3QUGDhQqBXL8DHh/uyEJFpY/ggMjMHDgDDh8uN4QBZ1zFrFhARATg7G3dsREQ1wfBBZCa0WuDNN4HXXisvKn3lFeCJJwA/P8Da2tgjJCKqGYYPIjNw5Qrw738DO3bI+127ypUtHTvKEMKiUiIyJwwfRCbuu+9kEen16+VFpXFxcl8WW1tjj46IqPYYPohMVGGhLCr97DN5PyQEWLwY6NMH8PTkvixEZL4YPohM0OHDwLBhwJkz8v7QocDs2bKo1NHRuGMjIrpXDB9EJkQIWVQ6d255UemrrwL//KdcQst9WYioMeCvMiITkZoqi0r/+EPe79IFWLIEaN8ecHEx6tCIiOoUwweRCfj+e+D552VRqbW1LCgdP14WlSqVxh4dEVHdYvggMqLCQmDyZOCTT+T9kBA529G3L+DuziW0RNQ4MXwQGcnRo8C//gUkJsr7Q4fKWo+ICMDOzqhDIyKqV7VerLdr1y4MHjwYfn5+UCgU+P777w0eF0Jg7ty58PX1hZ2dHfr164ezuj7QRAQhgLffljUdiYmyqPSdd4APPgBatWLwIKLGr9bhIz8/H23btsWHH35Y5eNvvvkm3n//faxatQr79u2Dg4MDBgwYgKKionseLJG5S0sDHngAePlluZqlS5fyeg8fH/buIKKmodaXXQYNGoRBgwZV+ZgQAsuXL8err76KRx55BADwn//8B97e3vj+++/xr3/9695GS2TGfvgBGDWqvFPp+PGyiZi/P5fQElHTUqf/zkpKSkJaWhr69eunP6ZSqdClSxfs2bOnyucUFxcjJyfH4EbUmBQWyvbojz4qg0dICPDVV7JpWHAwgwcRNT11Gj7S0tIAAN7e3gbHvb299Y/dasmSJVCpVPpbQEBAXQ6JyKiOHwc6dChfzfL448BPPwEPPwy4uRl3bERExmL0K8yzZs2CWq3W3y5dumTsIRHdMyFkEWmnTkBCgiwqfe894MMPgRYtABsbY4+QiMh46nTC18fHBwBw7do1+Pr66o9fu3YN7dq1q/I5NjY2sOFvYmpErl0Dnn4a2LpV3u/aVa5uadcOcHAw6tCIiExCnc58hISEwMfHB9u2bdMfy8nJwb59+xAbG1uX34rIJP30E9CypQweSiXw0kvAt9/KAMLgQUQk1XrmIy8vD+fOndPfT0pKwtGjR+Hm5obAwEBMnjwZr7/+OiIiIhASEoI5c+bAz88Pjz76aF2Om8ikFBYCU6cCq1bJ+yEhwLJlwP33AyqVccdGRGRqah0+Dh48iD59+ujvT506FQAwYsQIrF27FtOnT0d+fj7GjBmD7Oxs3Hffffjll19ga2tbd6MmMiEnTgBPPgnEx8v7TzwBLFwoAwj3ZSEiqkwhhBDGHkRFOTk5UKlUUKvVcHZ2NvZwiKolBLB8OTBrFlBcLItKFywAhg2TK1m4LwsRNSW1+fxmhwGiu1BVUek77wDt2wOc5CMiuj2GD6Ja2rwZePZZIDNTXlaZOFEWlnp5sT06EVFNMHwQ1VBhITBtGvDRR/J+SIi87HL//YCjo1GHRkRkVhg+iGrg1qLSJ58sLyple3QiotrhJDHRbeiKSjt1ksHD1RX44AO5pDYigsGDiOhu8FcnUTWuXQOeeQb47Td5v1s34N13gbZt2R6diOheMHwQVeHnn4GRI8uLSidPlkWlnp5cQktEdK8YPogquLWoNDRUznb06wfY2xt3bEREjQXDB9FNx4/LBmGnT8v7//oXsGgREBQEWFoad2xERI0JC06pyRNCbnffubMMHq6ucubj44/lzAeDBxFR3eLMBzVpaWmytuPXX+X9bt2A998H2rQBrK2NOjQiokaL4YOarFuLSqdMAV5+GXB3N/bIiIgaN4YPanJuLSoNC5O9PPr1474sREQNgeGDmpRbi0qHDQOWLAECArgvCxFRQ+GvW2oStFrDolI3N9ml9NNP5WoWBg8ioobDmQ9q9NLSgBEjyjuVdu8uW6S3bs2VLERExsB/71GjtnmzXLny22+yqHTWLHmsXTsGDyIiY+HMBzVKhYVy5cqHH8r74eGyU2n//jKEEBE1aUIYda8Ihg9qdG4tKh0+HFi6FGjWjPuyEFETlZ0N7N4N7Nolb76+wMaNRhsOwwc1GlqtbBA2YwZQUiL7dSxeDDz9NGBnZ+zRERE1oIwM4M8/y8PG0aNytkPH2Vn+0jRStT3DBzUKt3Yqve++8qJSrmQhokbvypXyoLFrV/nUb0WBgUCHDkDXrsADDxj1lyPDB5m9zZuB556TQd/GRjYQmz5dBnsiokZHCCApyTBsnD9f+bzwcBk2unQB+vQBoqLkNLAJXH9m+CCzdWtRaUSELCodMACw4p9sImoshAASEgzDxuXLhudYWMhw0bGjnNno0wcICZFtm00gbNyKv6LJLFVVVPrGG7KolIjIrGk0wIkThmEjI8PwHCsroGVLICZGzmzcf7/8BWhjY5Jh41YMH2RWqioqXbpUFpXa2Bh7dEREd6G0FDh8uDxo/PWXXJ1SkY2NLGLr1Kn8Moq3t9n+4mP4ILNxa1Fpz54yiLRpYxZBn4hIKioC9u8vDxt//w3k5xue4+AAtG0rw0a3bkCPHvJfW42kURHDB5mFW4tKX35ZFpU6ORl7ZEREd5CfLwOGLmzs2wcUFxueo1IB7dvLsNG9u7y5uDTaArbG+VNRo3FrUWlkJPDOO8DAgWyPTkQmKjtbXjrRhY1Dh4CyMsNz3N1lcWinTrI3QNeu8l9TTeQXG8MHmaxbi0qfflrWd/j5GXdcREQG0tMNG3odO2bY0AuQHUV1y1579JBfOzg0mbBxK4YPMjm3FpV6eJQXlTaSy51EZM6uXAF27iwPG/Hxlc8JCioPG716yWJROzt2PbyJ4YNMyq1Fpb17y94dbduyqJSIjEDX0Kti2LhwofJ5ERHl3UN795bXiE2koZcpYvggk7F5M/Dss0BmpiwqnTFDditlUSkRNRhdQ6+KYePKFcNzLCyA5s3Le2z07g2EhppNjw1TwPBBRldYKEPGRx/J+1FR5Z1KOUNJRPVKo5EFZhUbemVmGp5jZQW0alUeNvr0Afz9zbbHhilg+CCj2r9fznboikpHjgRef52dSomonpSWytUnFRt6qdWG59jaygZCFWc2fHxYdFaHGD7IKPLygDlzgPfek7OcLColonpRVCT7alRs6FVQYHiOgwPQrp0MG926yQ6G7u6AtbVRhtwUMHxQg/vlF2DsWODiRXn/H/8A5s+Xf/eJiO5JXl7lhl4lJYbnqFSyOFTXPbSRN/QyRXynqcFkZABTpgDr18v7fn4ydPz733KWk4io1m7cAHbvLi8QPXRI1nFU5OEhG3p17iyDRpcuTaqhlyli+KB6J4QMHJMnA9evyyLSp5+Wl13Cwow9OiIyK+nphsWhx49X3dCrY8fKDb1YwW4yGD6oXiUny0ssur4dkZHAokXAI4/wcioR1cDly4bLXhMSKp8THFw+s9GzJxt6mQGGD6oXGo3sUvrqq7K2S6kExo0DZs6Uu0ATEVUihGzgVTFsJCVVPi8ionwlSq9ecn2+rS17bJgRhg+qc8ePA88/Dxw4IO/HxMjls/368RIrEVUghGxNXjFsXL1qeE7Fhl667qEhIWzoZeYYPqjOFBUBCxcCb74pN3B0dJTNwyZNkoXkRNTEVWzotXOn3IytuoZenTqVN/Rq1owNve6REAIlmhL9TaFQwM3OzWjjYfigOrFzJzBmDHDmjLzfr5+s7YiJ4WVXoiZL19BLN7Px119ATo7hORUbeulmNry92fDnLgkhUKotRXFZsT5o5JXkoaC0AMWaYpRqSlGmLYPKRsXwQeYrOxuYPh349FN539NTrmIZNQqwtzfq0IiooRUWyrbFurCxZ0/lhl6OjrKpT6dOQGysXI3Chl61pgsZJZoSfdDIL81Hfkm+PnRohRa5Jbm4knMFl3Iu4VLOJVzMvojzN87Dx9EH+57fZ7TxM3zQXdu4ERg/HkhNlfefeEL27WjenJdiiZqE3FwZMHRhY//+yg29XFyA9u3lShRdQy+Vig29aqiqkFFQWoC8kjx9yCjTluFa3jVcyb0ZMtSXcFF9EUnZScgsyKzydfNL8hv4JzHE//tUa1euyNDx/ffyflCQLCh94gleliVq1G7ckJdOdGHj8OHKDb08PWVfjS5dZNDo3JkNvWqoYk1GcVkxCkoLkF+aj+IyebkkvzQfl3Mu41LOJVzOuYwUdQqSs5NxUX0RRWVF1b6ut4M3gl2CEagKRIBzAPyd/dHKq1UD/mSVMXxQjWm18vLK9Onysq2lJfDcc/IyS0CAsUdHRHXu2jVZFKoLGydOVG7o5edn2NCrfXs29LqDUk0pijXlNRn5Jfn6kFFSVoLMwkykqFMqhYyruVchIKp8TWsLawS5BOkDRoBzAIJUQQh0CYTKRgVrS2vYWdnBUekIpaUS9tbGvS7O8EE1kpgIjB4tfw8Bshh90SJg0CBeqiVqNC5dKl+JsmuX/It/K11Dr65dZdho1UoWePFaayWlmpuXS24GDV3IKNGUoKCkAJdzLiNZnSxnM9SXZMhQJyOnOKfa11TZqBDsEowgVRD8nf1lwLgZOGytbfXBwt7aHkpLJZSWSlhbWkNpqYSFwnQCIcMH3VZJCfDWW8CCBfJrOztg4kS5hNbDw9ijI6K7JgRw/rxh2EhONjxHoZANvXQzG717yzbFbOhl4NaQUVhaiNySXJRoSpBVkIWk7CT9TIYuZKTkpKBMW1bl61koLODn5IcgVRACnAMQqApEoCoQwS7B8HTwhLWFNeyt7eFg7QCllVIfMpSWSlhZmMfHunmMkoxi3z7ZLOzkSXm/Wzdg8WJ5GZe1YkRmRquVDb0qhg1dtbiOpaVhQ69evWRDL+78CKA8ZOiCRmFpIfJK8lBUVoQrOVdwIfsCLmZfREpOir7os7qCTwCwtbKVtRjOgQhQlV8mCVGFwNHGEUoLJRyUDrCztiufxbCQsxgKMw9//AihSvLyZFv099+X/zhycZFt0ceOlUXqRGQGNBrg2DHDhl7XrxueY20tL5tU7LHh79/ke2yUacv0K0t0ISO/JB/qYjWSbiThQvYFOXuhLg8ZhWWF1b6ep72nnMVQlc9ihKhC4OfsB1srWzmLoXTQBwvdzdKi8Rbp1nn4mDdvHubPn29wLCoqCglVbQZEJmfLFhkyUlLk/QcflLUdrVuzWJ3IpJWUGDb02r276oZebdsazmz4+DTZwq2KIUO3hDW/JB/X8q/h7PWzSMpOQrI6WX+p5HYFn1YWVrLQUxcwnAMR5BKEcNdwuNi5wNbKFg7WDrC1sjWoxbC2sDb7WYy7US8zHy1btsTvv/9e/k04R2/yMjKAKVOA9evlfV9fYN48YPhwWbhORCamsFBeG9XNbOzZI49V5OgoV5/ExJQ39PLwaHLXTcu0ZQZ9MgrLCpFdmI3k7GScyzpnMJORok6Bulhd7Ws5KZ3kqhLnQP0sRqhrKIJcgvR1GPbW9voiT93NlIo9TUG9/Am0srKCj49Pfbw01TEhgC+/lMHj+nW5Om7YMNksLCSEq+WITEZuLvD33+VhY/9+2b68IhcXWRzaqZMs0urWTR5rItOWupChCxpFZUVIz09HYmZipZBxOecySrWlVb6OAgr4OProw4Wu2DPMNQzeDt6ws5ZLVm2sbAxqMawtm+YM0t2ol/Bx9uxZ+Pn5wdbWFrGxsViyZAkCAwPr41vRPUhKkpdYfvtN3g8Pl5dYBg+Wq1qIyIiysmRDL13YOHKk6oZeHTvKRl733SdDh7Nzo/9Xg0ar0a8sKS6TNRlJ2UlIyEzA+azz+tUlKeoUpBekV/s6NpY2BgEjUBWIMJcwhLqFQmWj0hd73lqL0RQvk9S1Og8fXbp0wdq1axEVFYXU1FTMnz8fPXr0wMmTJ+Hk5FTp/OLiYhQXF+vv59x6jZLqnEYji0lffVVuu6BUyk3hZs2Sl1v494rICNLSZFGoLmycOFH5nGbNysNGz55yj5RG3NCrYsgo0ZRAXaRGfGY8zlw/g3PXz+lDxp0KPt3s3PR1GIGqQIS4hCDMLQwBzgFwUDrAwdrBcBbD0tpslqyaK4UQt7arq1vZ2dkICgrCO++8g1GjRlV6vKoCVQBQq9Vwdnauz6E1SceOyeWzBw/K++3by+WzffqwNTpRg0pJMVz2qtsSuqKQkPIeGz17ypUpdnaN7l8IGq3GoE9GWl4aTlw7gTPXz+D8jfKZjNTcVGiEpsrXsFRYoplTM/0MRpBLEMJcwxDmFgYvBy99sWfFWoymWuxZX3JycqBSqWr0+V3v0c7FxQWRkZE4d+5clY/PmjULU6dO1d/PyclBAHt117nCQmDhQuDNN+XMh6MjMHWqbBjm5tbofpcRmRYhgHPnZMjQBY6LFw3PUShkAy9d2OjVq9E19NKFDN3Kkgs3LuBk+slKISO7KLva13CwdigPGKog/SxGqGsonG2c4aB0MLhEYm1h3aiXrJqreg8feXl5OH/+PJ5++ukqH7exsYEN/8ldr3bskK3RdfmvTx9Z29GxY5Nfzk9UP7Ra4PTp8rBRXUOv6Gi5EkXXPTQkpFFMQWqFVr+yJKswCwmZCTiVcQpnrp/BhRsXkJydjEs5l1CiKan2NbwdvPWNt4JdghHqEoow9zAEOAXAycbJIGDoLpWQ+ajz8DFt2jQMHjwYQUFBuHr1Kl577TVYWlpi2LBhdf2t6A5u3ABefhn47DN538MDmD0bGDlSFsATUR0pKytv6LVrV/UNvVq3Lt8XpU8fWcNhxv8C0IWM4rJiXM69jNPpp3Eq4xTOZp3Vh4xr+deqfb7SUmmwP0mISwjC3cIR7hYOD3sPg/1JdAGDS1YbhzoPH5cvX8awYcNw/fp1eHp64r777sPevXvh6elZ19+KqiEEsHEjEBcnN6UEgMcek/uzREU12X5CRHWnpEQWTunCxl9/yaWwFdnayoLQijMb3t5m+RdQK7Qo0ZQgrzgPZ66fwemM0zideRpns84i6UYSLqovIq8kr9rnu9i4GCxZDXULRbhrOMJcw+Bs4wwbKxuDWgwWezZ+df5/+Jtvvqnrl6RauHJFho4ffpD3AwJk6BgyRK7AI6K7UFBg2NBr797KDb2cnMobenXtKgtE3d3NqqGXLmSk56fjdMZpxGfE43SGDBnJ2XL31eoKPi0UFvB19NXvTxLqEoowtzBEuUfBz8mvUe5PQnfPfP5W0G1ptcAnnwAzZsiOypaWwIgRwCuvAEFBZvX7j8j4cnIMG3odOFC5oZerK9Chg1z2GhtrVg29tEKLorIiXLhxAafSTyEhMwHxmfE4l3UOydnJuF54vdrn2lnZlc9iqIIR5haGcLdwNPdoDlc71ya1PwndPX4kNQIJCXL57O7d8n50NPD660D//nJVCxHdwfXrlRt6abWG53h5lffY6N5d/tfJyaR7bAghkF2UjfhMOYORkJGA+OvxOJ91HsnZySjWFFf7XE97z/IVJa4hCHMNQ5RHFEJdQg1nMZrw/iR09xg+zFhJCfDGGzJolJTIS8xxcXIJrbe3WfwDjMg40tIMV6JU1dDL3788bPToIes3HB1NctmrVqvFldwrOJl+EvGZ8YjPiNcvX73TZmj+zv6VZjFaeLaAt4M39yehesPwYab27QNGjQJOnZL3u3aVy2e7dgXs7Y07NiKTc/GiYdioqqFXaKhh91ATbOhVqinFmetn9CEjMTMRCdcTcOHGBeQUV98dWrcZmsEshnsUmns0N+iLwf1JqKEwfJiZvDxZx/HBB3JVi4uLXE77/PNyKa0JzwATNQwhgLNnDRt6paQYnqNQyKVfVTX0MgHqIjVOpp/EqYxTsp145hmcyTqD5OxklGnLqnxOxc3QQlxCEOoaigi3CLT0agl/Z/9KfTF4mYSMieHDjGzZArzwAnDpkrw/cKBcydK6tcn8ziRqeFqtnAKsOLORlmZ4jqUl0LKlYdgwckMvIQQu51zGifQTcmXJzT1Lzlw/g/T822+GFqAKQIhLiOzu6RqG5h7NEe0ZDRdbF+5PQmaBfzLNQEaGbIOuW8Xs4wPMnQv8859sjU5NUFkZcPSoYUOvrCzDc5RKmcor9tgwUkOv4rJinLl+BqcyTulXlpzJOoNzWedQUFpQ7fPc7NzKu3u6hiLSPRItPVsi1DW00lbunMUgc8PwYcKEAL74ApgyRf5utbCQgePVV4Hw8EbRhZnozoqLDRt67d5duaGXnZ1hQ69evWRKb8A15tcLriM+Ix6nMk7pZzLOXj+LlJwUaIW2yudYKCzg7+SPIJcghLqGyiWr7s3R0qslPB08uT8JNVoMHyYqNRV4+mlg2zZ5PyxMXmIZNEjWefAfOtRoFRTIJl66sLFnD1BUZHiOk5PssaFr6NWjhyx6quclXhqtBsnZyTJkZJ5CfIYs+jyTdQZZhVnVPs/B2gFBLhVmMdwi0cqrFcLdwuGodOT+JNTkMHyYoIwM4P77Zf8OpVKuapk6VTYLM8POzES3l5MjZzN0YaOqhl5ubrJeIyZGNvOKjZVNvuqpwjqvRLYRT8hMwKl0WfSZkClXldyuN4a3gzeCXYIR4hKCCLcIRHlEoZVXKzRzaqa/VML9SYgYPkyOWi2bgyUkyH/IffaZnEFWqYw9MqI6cv26rNPQhY3qGnrFxJQ39OrUqc4begkhkJqXioTMBP2uqwkZCUi8nogruVeqfZ61hbW+FiPENQRR7lFo6dkSzT2aw8XOhfuTENUA/3aYkIIC4KGHZC2di4sMHgMHmvWml0TyGmLFlSgnT1Y+x9+/PGz06AG0bVtnDb1KNCU4n3VeHzJOZ5xGQqYMGbkludU+z8XWBcGqYH1fjBaeLdDSqyVCXEJgZ23H/UmI7gHDh4koLpY7z+7eDTg4AJ9+CgwYwOBBZkarBc6fl3UaurBx9mzl80JDy8NGz55yGew9NvS6UXhDHzASMhNwOvM0EjMTceHGhdtuhtbMqZmsxXAJRbh7OFp7t0Zz9+bwcfTh/iRE9YThwwSUlQHDhwO//Sb7dXz0EfDww1zNQiauoEC2JT96FDh2TP73+HEgP9/wPF1DL91KlJ4977qhl1ZokaJOMQgZ8RnxSLiecNveGPZW9ghyCdL3xYjyiEJr79aIdIuEk40T9ychamAMH0am1crupN99J4tJ330XeOIJNg0jEyKEbNp19Khh0Dh7tnKtBiBTc8Ww0asXEBxcqzRdUFqgL/jUh4ybTbiKyoqqfZ6Xg5d+FiPMTV4qae3VGoGqQH3BJ4s9iYyP4cOIhAAmTwbWrZMrBJcuBZ55Rs4+ExlFWRmQmFg5aGRkVH2+mxvQvLm8tWwJtGkj6zVcXe/YY0MIgfT8dIOAkXBdzmRcVF+s9nnWFtb6Dp+6FuKtvFqhlVcruNu7c38SIjPA8GFEc+cCK1bIr197TbZO56Zw1GDUanmZpGLQOHlSFiDdysJCzl5ERQHR0UCLFnJmIzxcpuXbrEIp1ZTiwo0LBgFD93V2UXa1z3O2cdYHjDDXMES6R6K1V2tEeUTB3tqexZ5EZozhw0jeegt4/XX59YwZcgbEwcGoQ6LGSgi5q2vFmYxjx4CkpKrPt7eXISMqSoaM1q1lQy9v79teOsktztVfHqk4m3Eu6xxKtaVVPkcBBZo5N5P7lLiGINw1HFEeUWjj3QYBzgGwsbLhklWiRoh/q43g44+B6dPl1+PHy11qnZyMOyZqJIqKgNOnKwcNtbrq8318yi+btGghW5S3bi3/QFbTLTQjP0PfPjw+I17+NzMel3MuVzssOys7fV+MMNcwWY/h0QItPVvCw8GDxZ5ETQzDRwP7+mvgxRfl1yNGyNkPZ2fjjonMVEaGYcA4ehSIjwc0VSwrtbKSl0giI2XIaNFCNu4KCpLVzbd88GuFFpeyL1YKGPEZ8bheeL3aIXnYeyDUNVR/qSTUNRStvVojwj0CTkonLlklIgAMHw3qxx/lfi1CAEOHAu+8w86lVAMaDXDuXOWgcfVq1eerVHImIzJSFoHqLpu4uVVqHFOqKcX564n6gKGb0UjITLjtjqvNnJrpZzDCXMMQ4RYhL5WoAmBvbc9ZDCK6LYaPBrJtm9yRVqORm8OtXCk/C4gM5OXJ3hkVg8bx47KnRlUCA2Vthm61SYcO8r6Dg8Flk4LSAiRknqw0i3E26yzKtGVVvrSVhRWCVcEIdZOzGGGuYWju0RytvVvDy8ELtlZcD05Ed4fhowHs3Qs88ohcRNCrF7B6NeDpaexRkVEJIWcubq3NOHtWPnYrW1sgIkIGi5Yt5S0mBvD1lUWgN2casgqzEJ9xAvHn4g1mM263dNXeyt6gFkPfStyzJVztXLlslYjqHMNHPTt+XM505OfLS+zr1gF+fsYeFTWo0lK5U2DFoHH0qNxgrSqenvKSSfPmQKtW5b0zVCrAygpCCFzJvSLDxbGfDGYzbtfl08XWBeGu4bIe42bIaOPdhvUYRNTgGD7q0ZkzwAMPANnZ8rL7+vWyvo8asexsw4Bx7Bhw6hRQUlL5XEtLICREBo3oaDmb0amTPGZnhzJokXQjSQaL+NUGl0tutyGar6OvvuAzzK28HiPYJZj1GERkEhg+6klKCtCvH5CeLj9bvvxSzppTIyGE7JNxa9C4WM3lDUfH8tkMXe+Mjh0BT08UWWiRmJl4M1h8i/j4eH0r8RJNFaEFgKXCEgGqAH0tRpibrMdo490GPo4+rMcgIpPG8FEPrl2TwePSJTnT8cUX8rOGzFRhoZy9qBg0jh8HcnKqPt/Pz7AItH17oGVLqK00iL9x5uZlkr8Qv+tTxGfEIyk7CVpRxR4pAGwsbQyWrlbcr8TNzo31GERklhg+6tiNG0D//rJu0NcX+M9/5Ew6Z7rNRHp65dqMxMSqe2dYW8veGbpOoC1aQMTE4JqbEvG5SYi/uU/J6XMbEb8nHql5qdV+W2cbZ33ICHcNR7jbza3dPZqzHoOIGh2GjzqUlwc8+KD8R7G7O7BmDdCjB4OHSdJoZEK8NWikpVV9vqtrecvxli2hbd0KyeHuiC9NRfyNszdnM7Yj/vv42+5X4uXgVakeo613W4S4hMBB6cB6DCJqEhg+6khRkVxOu3ev7Fj6+eey2JSfJSYgN1f2zqhYm3HihLycciuFQvbOuNlyvKRFFM5GeSDesQjxuRcQfz0RpzPWI/FQIor2Vb21u4XCAv5O/gh1K5/JaO7ZHG2928LPyY/1GETU5DF81IHSUuDJJ4E//pB7cn3yCfDQQ7fd6JPqgxDA5cuVe2ecO1f1+ba2+iLQ3BZhSAhVId7PGvFlN2czMrfgfNoH0KRWcckFcmt3/a6rbmEIdwtHtEc02ni3gbu9O5SWyiqfR0TU1DF83COtFhg5UrZOt7EBPvwQePzxavfkorpSUiL3Mbk1aGRlVX2+lxcQFYWMFkGIj3CVIcMuX85mZO7BpZz/Ay5A3m7hYO1gcKkkzC0MrTxboaVXS6hsVKzHICKqJYaPeyAEEBcHfPWV3Lfr3XeBf/+bwaPOZWVV3tfk9Gk55XQrKyuIkGBcah2E+EhXnPa1RrxLKeJLriI+6zSuF/4J5AE4U/mp7nbu+s3Qwt1k0Wcb7zaIcItgPQYRUR1i+LhLQgAzZgCrVskygSVLgOeflyGE7pJWK3tnVKzNOHpUrlmuQqnKEefbBSG+uTvim9kg3qUM8RbXkaA+j/zSm5dabty8VdDMqZl+JiPcLRyR7pFo690WgS6BsLG0YcggIqpn/Ki8S4sXA2+9Jb9+7TVg4kS58pJqoKREhoxz5+SKkzNn5BKhY8fkkqFbFFgDiS28EB/tifgAO5x2LUO8dTbOFV5BqfaUPEkLoMIVFysLKwSqAvWXSsJdwxHtKesxvBy8WI9BRGREDB93YcUK4NVX5dcvvyxnQJT8LDNUVgYkJ8twcevt4sUq+2Zk2QHxoVaIb+GF+EB7xHsIxCtzcLEsEwLpAG7uW1Jy8wbAzsrOoAlXuFs4Wnm1QiuvVnCxdWE9BhGRCWL4qKV16+QsBwCMHQvMmycXTTRJGo3sI19VwEhKkgGkgjIL4IoTkOwPJHtZIznQGcleNkhy0SLBJg/XkAegDMDVCk+S/3GxdTEIGLp6jEj3SDgqHWGh4NIiIiJzwfBRC999Bzz3nPz66aeBN9+US2sbNa1WLl+tKmBcuGCwYZpGAVxxBpJdgOSWQLK7JZL97JHkYYlkxzJcssqHBrrt4ksBVN7V1cfRxyBkRLlH6TdFs7WyZT0GEVEjwPBRQ7/+CgwbJj+LH3tMXnpxcjL2qOqIEMDVq5XDxZkzMmAUyWZaGgVw1elmuHABkrsCyW4KJHnbINkFuGRbjDKFqPDCGgCGu69aWVjBz8kP/k7+aObcDM2cmqGZczN9yPB18mU9BhFRI8fwcQdCAN9+K3t5lJYCAwbIJmIqlbFHVktCyB3vqgoY588DBQXQKIDUiuHCA0gOB5JdFUhyt0SKowZlFuLWFwZQ3unTysIKvo6+8Hf2RzOnZvK/zs0Q4hKCcLdwhLiGwMHagbUYRERNGMPHbezaBUyfDuzbJ+/36CHbpnt4GHdc1RICyMysNmBo83KR6lghXLgAycFAcjsgyRVIUQGlVWYCAV3xhS5c6GYtdOEiWBWMCPcIBKuC4WjjCCsL/tEiIqKq8ROiCqdOATNnAps3y/t2dnLmY9YsuVu60WVlVRkwtOfOIk2bYxguXIDk1kByD+CiCii5w/9xKwsr+Dj6wN/JH37O5ZdHglRBiHCLQIhrCJxsnBguiIjorvETpILLl2XPjrVrZW2HpSUwdCgwdSrQqVMD79VSUiI3PztzRh8uxNkzSLtyBskKdeWAEQtcHAQU3+H/qKXCUoaLm5dFmjnL2YtAVSDC3cIR6hoKJ6UTrC3ZtISIiOoHwwcAtRpYuhRYvlxfW4l+/YCXXpL/bZCupbm5wJ49wJ9/Iu/vHdh/ZT8OeJQgyfVmuHAFLj4AFN0hE1goLODj4KO/HKK7NGIQLmycWNRJRERG06TDR3Ex8NFHwOuvl+9H1qGDrPN45JF67t+Rng78+Sfw11+4tP937M47hb+bCewOAI7dB2iqmWWxgAI+9l5opgpAM2d/fcgIcA7QhwuVrYrhgoiITFaTDB9aLfD117JLaXKyPBYaCkybJjeGq/MltELIJat//YWyP3fi+Mlt2C1SsDsQ+DsAuNS/8lP8lO5o26wjQt3D0czZHwGqm+HCRYYLGyubOh4kERFRw2hy4WPrVtkO/cgRed/LC5gwARgzRn5dJzQa4ORJ4M8/kb17G/Ze2IndDjfwdwCwzx/IDzA83RIKRDuFoX1gZ3Tw74RuAd3Q0rMlHJQOdTQgIiIi09FkwkdCgmyLvnWrvO/oCIweDYwfL2c97klREXDwIMSuXbhw4FfsTjuAvz0KsTsAOBUFiOaGpzvDFm09W6NDcFe0b9YRPQN7wt/Zn0WeRETUJDSZ8FFWBmzbJneeHTYMmDIFaNsWuKtu3Wo18PffKP5zO44c+xW7c09ht58GfwcA19pVPj3I0h3tAzqhfWAXdA7ois7NOsPV1pWtwomIqElqMuGjVStg1Sqgc2f5tWVNG2wKITdJO3gQmX/9hr/PbMNubTL+9gcONAOKOxuebi0s0No+GG1DuqJjYFd0D+yBKI8o2Fnb1fnPREREZI6aTPgA5GWW2yopAeLjoT18CMkn/sSJi/txIvc8TqqKccQHOOMBINbwKW7CDh1co9EuvAfaB3RGz+Ce8HH0YRMuIiKiajTdT8icHOD4caQf2oUTiX/iZPoJnNCk4oSHFqe8gHwVgDaVnxYBd7T3bY92ET3RNfg+dPDtAGcbZ15CISIiqqF6Cx8ffvgh3nrrLaSlpaFt27ZYsWIFOnfufOcn1pOii+dx/Iu3cSLlIE7knMVJGzVOeAHpjgC8b94qUGotEKFwR4RLKMKDOyKyWRt0C7oP4W7hXOZKRER0D+olfGzYsAFTp07FqlWr0KVLFyxfvhwDBgxAYmIivOpsPWvtHE3ei1jNKqAZ5O0mhQCCyhwRZdMM4b4tEB7eGa0DYtDGpy1c7Vx5+YSIiKiOKYQQt+6Rfs+6dOmCTp064YMPPgAAaLVaBAQEYMKECZg5c+Ztn5uTkwOVSgW1Wg1nZ+c6G1Ne/g1ELPVDBNwR6RiMsNCOiI6IRfugrvBx8uVsBhER0T2ozed3nf+zvqSkBIcOHcKsWbP0xywsLNCvXz/s2bOnrr9djTk6uCJ1YaHRvj8RERFJdR4+MjMzodFo4O1tWETh7e2NhISESucXFxejuLhYfz8nJ6euh0REREQmpCE3ia/SkiVLoFKp9LeAgIA7P4mIiIjMVp2HDw8PD1haWuLatWsGx69duwYfH59K58+aNQtqtVp/u3TpUl0PiYiIiExInYcPpVKJjh07Ytu2bfpjWq0W27ZtQ2xsbKXzbWxs4OzsbHAjIiKixqte1pFOnToVI0aMQExMDDp37ozly5cjPz8fzz77bH18OyIiIjIj9RI+nnzySWRkZGDu3LlIS0tDu3bt8Msvv1QqQiUiIqKmp176fNyL+urzQURERPWnNp/fRl/tQkRERE0LwwcRERE1KIYPIiIialAMH0RERNSgGD6IiIioQTF8EBERUYNi+CAiIqIGVS9Nxu6Fru0Id7clIiIyH7rP7Zq0DzO58JGbmwsA3N2WiIjIDOXm5kKlUt32HJPrcKrVanH16lU4OTlBoVBUejwnJwcBAQG4dOkSO6BWg+/RnfE9ujO+R7fH9+fO+B7dWWN6j4QQyM3NhZ+fHywsbl/VYXIzHxYWFvD397/jedwB9874Ht0Z36M743t0e3x/7ozv0Z01lvfoTjMeOiw4JSIiogbF8EFEREQNyuzCh42NDV577TXY2NgYeygmi+/RnfE9ujO+R7fH9+fO+B7dWVN9j0yu4JSIiIgaN7Ob+SAiIiLzxvBBREREDYrhg4iIiBqU2YWPDz/8EMHBwbC1tUWXLl2wf/9+Yw/JaHbt2oXBgwfDz88PCoUC33//vcHjQgjMnTsXvr6+sLOzQ79+/XD27FnjDNYIlixZgk6dOsHJyQleXl549NFHkZiYaHBOUVER4uLi4O7uDkdHRwwdOhTXrl0z0ogb3sqVK9GmTRt9j4HY2Fhs2bJF/3hTf39utXTpUigUCkyePFl/rKm/R/PmzYNCoTC4NW/eXP94U39/dK5cuYJ///vfcHd3h52dHVq3bo2DBw/qH29qv6/NKnxs2LABU6dOxWuvvYbDhw+jbdu2GDBgANLT0409NKPIz89H27Zt8eGHH1b5+Jtvvon3338fq1atwr59++Dg4IABAwagqKiogUdqHDt37kRcXBz27t2LrVu3orS0FP3790d+fr7+nClTpuCnn37Ct99+i507d+Lq1asYMmSIEUfdsPz9/bF06VIcOnQIBw8exP33349HHnkEp06dAsD3p6IDBw7g448/Rps2bQyO8z0CWrZsidTUVP3tr7/+0j/G9we4ceMGunfvDmtra2zZsgWnT5/GsmXL4Orqqj+nyf2+Fmakc+fOIi4uTn9fo9EIPz8/sWTJEiOOyjQAEJs2bdLf12q1wsfHR7z11lv6Y9nZ2cLGxkZ8/fXXRhih8aWnpwsAYufOnUII+X5YW1uLb7/9Vn9OfHy8ACD27NljrGEanaurq1i9ejXfnwpyc3NFRESE2Lp1q+jVq5eYNGmSEIJ/hoQQ4rXXXhNt27at8jG+P9KMGTPEfffdV+3jTfH3tdnMfJSUlODQoUPo16+f/piFhQX69euHPXv2GHFkpikpKQlpaWkG75dKpUKXLl2a7PulVqsBAG5ubgCAQ4cOobS01OA9at68OQIDA5vke6TRaPDNN98gPz8fsbGxfH8qiIuLw0MPPWTwXgD8M6Rz9uxZ+Pn5ITQ0FMOHD0dKSgoAvj86P/74I2JiYvDEE0/Ay8sL7du3x6effqp/vCn+vjab8JGZmQmNRgNvb2+D497e3khLSzPSqEyX7j3h+yVptVpMnjwZ3bt3R6tWrQDI90ipVMLFxcXg3Kb2Hp04cQKOjo6wsbHB2LFjsWnTJrRo0YLvz03ffPMNDh8+jCVLllR6jO8R0KVLF6xduxa//PILVq5ciaSkJPTo0QO5ubl8f266cOECVq5ciYiICPz666948cUXMXHiRKxbtw5A0/x9bXIbyxHVh7i4OJw8edLgWjRJUVFROHr0KNRqNf773/9ixIgR2Llzp7GHZRIuXbqESZMmYevWrbC1tTX2cEzSoEGD9F+3adMGXbp0QVBQEP7v//4PdnZ2RhyZ6dBqtYiJicHixYsBAO3bt8fJkyexatUqjBgxwsijMw6zmfnw8PCApaVlpSrpa9euwcfHx0ijMl2694TvFzB+/Hhs3rwZ27dvN9gx2cfHByUlJcjOzjY4v6m9R0qlEuHh4ejYsSOWLFmCtm3b4r333uP7A3nZID09HR06dICVlRWsrKywc+dOvP/++7CysoK3t3eTf49u5eLigsjISJw7d45/hm7y9fVFixYtDI5FR0frL081xd/XZhM+lEolOnbsiG3btumPabVabNu2DbGxsUYcmWkKCQmBj4+PwfuVk5ODffv2NZn3SwiB8ePHY9OmTfjjjz8QEhJi8HjHjh1hbW1t8B4lJiYiJSWlybxHVdFqtSguLub7A6Bv3744ceIEjh49qr/FxMRg+PDh+q+b+nt0q7y8PJw/fx6+vr78M3RT9+7dKy3zP3PmDIKCggA00d/Xxq54rY1vvvlG2NjYiLVr14rTp0+LMWPGCBcXF5GWlmbsoRlFbm6uOHLkiDhy5IgAIN555x1x5MgRcfHiRSGEEEuXLhUuLi7ihx9+EMePHxePPPKICAkJEYWFhUYeecN48cUXhUqlEjt27BCpqan6W0FBgf6csWPHisDAQPHHH3+IgwcPitjYWBEbG2vEUTesmTNnip07d4qkpCRx/PhxMXPmTKFQKMRvv/0mhOD7U5WKq12E4Hv00ksviR07doikpCSxe/du0a9fP+Hh4SHS09OFEHx/hBBi//79wsrKSixatEicPXtWrF+/Xtjb24svv/xSf05T+31tVuFDCCFWrFghAgMDhVKpFJ07dxZ79+419pCMZvv27QJApduIESOEEHL51pw5c4S3t7ewsbERffv2FYmJicYddAOq6r0BINasWaM/p7CwUIwbN064uroKe3t78dhjj4nU1FTjDbqBPffccyIoKEgolUrh6ekp+vbtqw8eQvD9qcqt4aOpv0dPPvmk8PX1FUqlUjRr1kw8+eST4ty5c/rHm/r7o/PTTz+JVq1aCRsbG9G8eXPxySefGDze1H5fc1dbIiIialBmU/NBREREjQPDBxERETUohg8iIiJqUAwfRERE1KAYPoiIiKhBMXwQERFRg2L4ICIiogbF8EFEREQNiuGDqAno3bs3Jk+e3KDfMzk5GQqFAkePHq3z196xYwcUCkWlDcuIyDwwfBDRHZnah323bt2QmpoKlUpl7KEQ0V2wMvYAiIhqS6lUNtqtxomaAs58EDURZWVlGD9+PFQqFTw8PDBnzhzotnb64osvEBMTAycnJ/j4+OCpp55Ceno6AHn5pE+fPgAAV1dXKBQKjBw5EgCg1Wrx5ptvIjw8HDY2NggMDMSiRYsMvu+FCxfQp08f2Nvbo23bttizZ0+Nxnvx4kUMHjwYrq6ucHBwQMuWLfG///0PQOWZmN69e0OhUFS6JScnAwCys7Px/PPPw9PTE87Ozrj//vtx7Nixe3k7iegeMHwQNRHr1q2DlZUV9u/fj/feew/vvPMOVq9eDQAoLS3FwoULcezYMXz//fdITk7WB4yAgAB89913AIDExESkpqbivffeAwDMmjULS5cuxZw5c3D69Gl89dVX8Pb2Nvi+s2fPxrRp03D06FFERkZi2LBhKCsru+N44+LiUFxcjF27duHEiRN444034OjoWOW5GzduRGpqqv42ZMgQREVF6cfyxBNPID09HVu2bMGhQ4fQoUMH9O3bF1lZWXf1XhLRPTLyrrpE1AB69eoloqOjhVar1R+bMWOGiI6OrvL8AwcOCAAiNzdXCCHE9u3bBQBx48YN/Tk5OTnCxsZGfPrpp1W+RlJSkgAgVq9erT926tQpAUDEx8ffccytW7cW8+bNq/Kxqsaj88477wgXFxf9duR//vmncHZ2FkVFRQbnhYWFiY8//viO4yCiuseZD6ImomvXrlAoFPr7sbGxOHv2LDQaDQ4dOoTBgwcjMDAQTk5O6NWrFwAgJSWl2teLj49HcXEx+vbte9vv26ZNG/3Xvr6+AKC/pHM7EydOxOuvv47u3bvjtddew/Hjx+/4nC1btmDmzJnYsGEDIiMjAQDHjh1DXl4e3N3d4ejoqL8lJSXh/Pnzd3xNIqp7DB9ETVxRUREGDBgAZ2dnrF+/HgcOHMCmTZsAACUlJdU+z87Orkavb21trf9aF360Wu0dn/f888/jwoULePrpp3HixAnExMRgxYoV1Z5/+vRp/Otf/8LSpUvRv39//fG8vDz4+vri6NGjBrfExES8/PLLNfoZiKhuMXwQNRH79u0zuL93715EREQgISEB169fx9KlS9GjRw80b9680syEUqkEAGg0Gv2xiIgI2NnZYdu2bfU25oCAAIwdOxYbN27ESy+9hE8//bTK8zIzMzF48GAMHToUU6ZMMXisQ4cOSEtLg5WVFcLDww1uHh4e9TZ2IqoewwdRE5GSkoKpU6ciMTERX3/9NVasWIFJkyYhMDAQSqUSK1aswIULF/Djjz9i4cKFBs8NCgqCQqHA5s2bkZGRgby8PNja2mLGjBmYPn06/vOf/+D8+fPYu3cvPvvsszoZ7+TJk/Hrr78iKSkJhw8fxvbt2xEdHV3luUOHDoW9vT3mzZuHtLQ0/U2j0aBfv36IjY3Fo48+it9++w3Jycn4+++/MXv2bBw8eLBOxkpEtcM+H0RNxDPPPIPCwkJ07twZlpaWmDRpEsaMGQOFQoG1a9filVdewfvvv48OHTrg7bffxj/+8Q/9c5s1a4b58+dj5syZePbZZ/HMM89g7dq1mDNnDqysrDB37lxcvXoVvr6+GDt2bJ2MV6PRIC4uDpcvX4azszMGDhyId999t8pzd+3aBUCGpIqSkpIQHByM//3vf5g9ezaeffZZZGRkwMfHBz179qy0MoeIGoZCiJsL/YmIiIgaAC+7EBERUYNi+CAioxg0aJDB0teKt8WLFxt7eERUj3jZhYiM4sqVKygsLKzyMTc3N7i5uTXwiIiooTB8EBERUYPiZRciIiJqUAwfRERE1KAYPoiIiKhBMXwQERFRg2L4ICIiogbF8EFEREQNiuGDiIiIGhTDBxERETWo/wc7Ai9NcA+OQQAAAABJRU5ErkJggg==", - "text/plain": [ - "
" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "decoding benchmark:\n", - " batch_size flash-attn-v2 bitdecoding-kc-4 bitdecoding-kc-2\n", - "0 2.0 0.665600 0.245760 0.224256\n", - "1 4.0 1.295360 0.465920 0.427008\n", - "2 8.0 3.757568 1.016832 0.958464\n", - "3 16.0 6.727680 1.687552 1.541120\n", - "4 32.0 13.357056 3.412992 3.031040\n", - "5 64.0 26.876928 7.927808 5.938176\n" - ] - } - ], - "source": [ - "@triton.testing.perf_report(\n", - " triton.testing.Benchmark(\n", - " x_names=[\"batch_size\"],\n", - " x_vals=[2**i for i in range(1, 7, 1)],\n", - " line_arg='provider', # Argument name whose value corresponds to a different line in the plot.\n", - " line_vals=['flash-attn-v2', 'bitdecoding-kc-4', 'bitdecoding-kc-2'], # Possible values for `line_arg`.\n", - " line_names=['flash-attn-v2', 'bitdecoding-kc-4', 'bitdecoding-kc-2'], # Label name for the lines.\n", - " styles=[('blue', '-'), ('red', '-'), ('green', '-')], # Line color and style.\n", - " plot_name=\"decoding benchmark\",\n", - " args={}, # Values for function arguments not in `x_names` and `y_name`.\n", - " )\n", - ")\n", - "def benchmark(batch_size, provider):\n", - " torch.random.manual_seed(0)\n", - " device = \"cuda\"\n", - " dtype = torch.float16\n", - "\n", - " seq_len = 32768\n", - " nheads = 32\n", - " nheads_k = 32\n", - " d = 128\n", - " sm_scale = 1.0 / math.sqrt(d)\n", - "\n", - " # Quantization parameters\n", - " quant_mode = \"k-channel\"\n", - " group_size = 128\n", - "\n", - " q = torch.randn(batch_size, 1, nheads, d, device=device, dtype=dtype)\n", - " k_cache = torch.randn(batch_size, seq_len, nheads_k, d, device=device, dtype=dtype)\n", - " v_cache = torch.randn(batch_size, seq_len, nheads_k, d, device=device, dtype=dtype)\n", - "\n", - " quantiles = [0.5, 0.2, 0.8]\n", - "\n", - " if provider == 'flash-attn-v2':\n", - " ms, min_ms, max_ms = triton.testing.do_bench(lambda: flash_attn_with_kvcache(q, k_cache, v_cache), quantiles=quantiles)\n", - " elif provider == 'bitdecoding-kc-4':\n", - " num_bits = 4\n", - " pack_nums = 16 / num_bits\n", - " k_pack = torch.zeros((batch_size, int(seq_len // pack_nums), nheads_k, d), dtype=torch.uint16, device=device)\n", - " k_params = torch.zeros((batch_size, int(seq_len // group_size), nheads_k, d), dtype=torch.float32, device=device)\n", - " v_pack = torch.zeros((batch_size, seq_len, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device)\n", - " v_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seq_len), dtype=torch.float32, device=device)\n", - " ms, min_ms, max_ms = triton.testing.do_bench(lambda: fwd_kvcache_int(q,\n", - " k_pack, k_params, \n", - " v_pack, v_params,\n", - " None, # opt_block_table\n", - " sm_scale,\n", - " quant_mode, \n", - " group_size,\n", - " num_bits), quantiles=quantiles)\n", - " elif provider == 'bitdecoding-kc-2':\n", - " num_bits = 2\n", - " pack_nums = 16 / num_bits\n", - " k_pack = torch.zeros((batch_size, int(seq_len // pack_nums), nheads_k, d), dtype=torch.uint16, device=device)\n", - " k_params = torch.zeros((batch_size, int(seq_len // group_size), nheads_k, d), dtype=torch.float32, device=device)\n", - " v_pack = torch.zeros((batch_size, seq_len, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device)\n", - " v_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seq_len), dtype=torch.float32, device=device)\n", - " ms, min_ms, max_ms = triton.testing.do_bench(lambda: fwd_kvcache_int(q,\n", - " k_pack, k_params, \n", - " v_pack, v_params,\n", - " None, # opt_block_table\n", - " sm_scale,\n", - " quant_mode, \n", - " group_size,\n", - " num_bits), quantiles=quantiles)\n", - " \n", - " return ms, min_ms, max_ms\n", - "\n", - "benchmark.run(show_plots=True, print_data=True)" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "bitdecode", - "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.10.16" - } - }, - "nbformat": 4, - "nbformat_minor": 2 -} diff --git a/bit_decode/__init__.py b/bit_decode/__init__.py index bd3d949..d68dd31 100644 --- a/bit_decode/__init__.py +++ b/bit_decode/__init__.py @@ -4,3 +4,5 @@ kvcache_pack_int, fwd_kvcache_int ) + +from bit_decode.models.cache_utils import Cache, DynamicCache, StaticCache \ No newline at end of file diff --git a/bit_decode/bit_decode_interface.py b/bit_decode/bit_decode_interface.py index 39f4fc3..b23bc6d 100644 --- a/bit_decode/bit_decode_interface.py +++ b/bit_decode/bit_decode_interface.py @@ -5,6 +5,8 @@ import torch import torch.nn as nn +# isort: off +# We need to import the CUDA kernels after importing torch import bit_decode_cuda as bit_decode_cuda def kvcache_pack_int(k_cache: torch.Tensor, k_pack: torch.Tensor, k_params: torch.Tensor, @@ -12,7 +14,7 @@ def kvcache_pack_int(k_cache: torch.Tensor, k_pack: torch.Tensor, k_params: torc opt_block_table: Optional[torch.Tensor] = None, cu_seqlens_k: torch.Tensor = None, seqlen_k: int = 0, - quant_mode: str = "k-channel", + quant_mode: str = "k-tensor", group_size: int = 128, num_bits: int = 4): @@ -22,58 +24,53 @@ def kvcache_pack_int(k_cache: torch.Tensor, k_pack: torch.Tensor, k_params: torc V_unpad = v_cache.reshape(batch_size * seqlen_k, nheads_k, d) if num_bits == 4: - bit_decode_cuda.kvcache_pack_i4(K_unpad, k_pack, k_params, - V_unpad, v_pack, v_params, - opt_block_table, - cu_seqlens_k, - seqlen_k, - quant_mode, - group_size - ) - else: - bit_decode_cuda.kvcache_pack_i2(K_unpad, k_pack, k_params, - V_unpad, v_pack, v_params, - opt_block_table, - cu_seqlens_k, - seqlen_k, - quant_mode, - group_size - ) + bit_decode_cuda.kvcache_pack_int4(K_unpad, k_pack, k_params, + V_unpad, v_pack, v_params, + opt_block_table, + cu_seqlens_k, + seqlen_k, + quant_mode, + group_size + ) + # else: + # bit_decode_cuda.kvcache_pack_int2(K_unpad, k_pack, k_params, + # V_unpad, v_pack, v_params, + # opt_block_table, + # cu_seqlens_k, + # seqlen_k, + # quant_mode, + # group_size + # ) def fwd_kvcache_int(q: torch.Tensor, k_pack: torch.Tensor, k_params: torch.Tensor, v_pack: torch.Tensor, v_params: torch.Tensor, + opt_k_new: Optional[torch.Tensor] = None, + opt_v_new: Optional[torch.Tensor] = None, + opt_seqlens_k: Optional[torch.Tensor] = None, + k_pack_new: torch.Tensor = None, k_params_new: torch.Tensor = None, + v_pack_new: torch.Tensor = None, v_params_new: torch.Tensor = None, opt_block_table: Optional[torch.Tensor] = None, softmax_scale: float = 1.0, - quant_mode: str = "k-channel", + quant_mode: str = "k-tensor", group_size: int = 128, + residual_block_size: int = 128, + new_lens: int = 0, num_bits: int = 4): if num_bits == 4: - out_bit = bit_decode_cuda.fwd_kvcache_i4( - q, - k_pack, k_params, - v_pack, v_params, - opt_block_table, - softmax_scale, - quant_mode, - group_size, - False, # is_causal - -1, # window_size_left - -1, # window_size_right - 0.0, # softcap - True, # is_rotary_interleaved - 0 # num_splits - ) - else: - out_bit = bit_decode_cuda.fwd_kvcache_i2( + out_bit, k_pack_new, k_params_new, v_pack_new, v_params_new = bit_decode_cuda.fwd_kvcache_int4( q, k_pack, k_params, v_pack, v_params, + opt_k_new, opt_v_new, opt_seqlens_k, + k_pack_new, k_params_new, v_pack_new, v_params_new, opt_block_table, softmax_scale, quant_mode, group_size, + residual_block_size, + new_lens, False, # Added -1, # Added -1, # Added @@ -81,6 +78,26 @@ def fwd_kvcache_int(q: torch.Tensor, True, # Added 0 # Added ) + # else: + # out_bit, k_pack_new, k_params_new, v_pack_new, v_params_new = bit_decode_cuda.fwd_kvcache_int2( + # q, + # k_pack, k_params, + # v_pack, v_params, + # opt_k_new, opt_v_new, opt_seqlens_k, + # k_pack_new, k_params_new, v_pack_new, v_params_new, + # opt_block_table, + # softmax_scale, + # quant_mode, + # group_size, + # residual_block_size, + # new_lens, + # False, # Added + # -1, # Added + # -1, # Added + # 0.0, # Added + # True, # Added + # 0 # Added + # ) - return out_bit + return out_bit, k_pack_new, k_params_new, v_pack_new, v_params_new diff --git a/bit_decode/models/__init__.py b/bit_decode/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/bit_decode/models/cache_utils.py b/bit_decode/models/cache_utils.py new file mode 100644 index 0000000..6b54663 --- /dev/null +++ b/bit_decode/models/cache_utils.py @@ -0,0 +1,2645 @@ +import copy +import importlib.metadata +import json +import os +from dataclasses import dataclass +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union + +import torch +from packaging import version +from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_6 +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import ( + is_hqq_available, + is_optimum_quanto_available, + is_torchdynamo_compiling, + logging, + is_torch_greater_or_equal +) +from transformers.utils.deprecation import deprecate_kwarg + + +if is_hqq_available(): + from hqq.core.quantize import Quantizer as HQQQuantizer + +logger = logging.get_logger(__name__) + + +# Utility functions for static/sliding cache update logic +def _static_cache_update( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + cache_position: Optional[torch.LongTensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the static cache tensors in place. + + Args: + k_cache (`torch.Tensor`): The key cache tensor to update. + v_cache (`torch.Tensor`): The value cache tensor to update. + key_states (`torch.Tensor`): The new key states to add. + value_states (`torch.Tensor`): The new value states to add. + cache_position (`Optional[torch.LongTensor]`): The position indices where the new states should be inserted. + If None, the entire cache is overwritten (prefill). + + Returns: + Tuple[`torch.Tensor`, `torch.Tensor`]: The updated key and value cache tensors (modified in-place). + """ + if cache_position is None: + # Prefill phase where seq_len potentially equals max_cache_len. Directly copy. + k_cache.copy_(key_states) + v_cache.copy_(value_states) + else: + # Generation phase. Update specific positions. + # Use index_copy_ for in-place update (compile-friendly). + try: + k_cache.index_copy_(2, cache_position, key_states) + v_cache.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # Fallback for devices like MPS where index_copy_ might not be supported. + k_cache[:, :, cache_position] = key_states + v_cache[:, :, cache_position] = value_states + return k_cache, v_cache + + +def _sliding_cache_update( + k_cache: torch.Tensor, + v_cache: torch.Tensor, + key_states: torch.Tensor, + value_states: torch.Tensor, + cache_position: torch.LongTensor, + max_cache_len: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the sliding window cache tensors, returning the potentially modified tensors. + + Args: + k_cache (`torch.Tensor`): The key cache tensor to update. + v_cache (`torch.Tensor`): The value cache tensor to update. + key_states (`torch.Tensor`): The new key states to add. + value_states (`torch.Tensor`): The new value states to add. + cache_position (`torch.LongTensor`): The position indices where the new states should be inserted. + max_cache_len (`int`): The maximum length of the sliding window cache. + + Returns: + Tuple[`torch.Tensor`, `torch.Tensor`]: The key and value tensors representing the cache state after the update. + For prefill > window, these are the full input states. + Otherwise, they are the updated cache tensors. + """ + # Handle prefill phase when prompt length > sliding_window_size + if cache_position.shape[0] > max_cache_len: + new_k = key_states[:, :, -max_cache_len:, :] + new_v = value_states[:, :, -max_cache_len:, :] + k_cache.copy_(new_k) + v_cache.copy_(new_v) + return key_states, value_states + + # Sliding window logic for generation phase or prefill < window + slicing = torch.arange(max_cache_len, device=value_states.device) + current_seq_len = cache_position[-1] + 1 # Use last position to determine current length + to_shift = current_seq_len > max_cache_len + indices = (slicing + to_shift.sum()) % max_cache_len + + k_out_shifted = k_cache[:, :, indices] + v_out_shifted = v_cache[:, :, indices] + + # Clamp cache_position to determine the *target index* within the shifted cache view + update_position = cache_position.clamp(min=0, max=max_cache_len - 1) + + try: + k_out_updated = k_out_shifted.index_copy(2, update_position, key_states) + v_out_updated = v_out_shifted.index_copy(2, update_position, value_states) + except NotImplementedError: + # Fallback for MPS: clone and modify the clone + k_out_updated = k_out_shifted.clone() + v_out_updated = v_out_shifted.clone() + k_out_updated[:, :, update_position] = key_states + v_out_updated[:, :, update_position] = value_states + + k_cache.copy_(k_out_updated) + v_cache.copy_(v_out_updated) + return k_out_updated, v_out_updated + + +class Cache: + """ + Base, abstract class for all caches. The actual data structure is specific to each subclass. + """ + + is_compileable = False + + def __init__(self): + super().__init__() + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. These are specific to each subclass and allow new types of + cache to be created. + + Return: + A tuple containing the updated key and value states. + """ + raise NotImplementedError("Make sure to implement `update` in a subclass.") + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # TODO: deprecate this function in favor of `cache_position` + raise NotImplementedError("Make sure to implement `get_seq_length` in a subclass.") + + def get_max_cache_shape(self) -> Optional[int]: + """Returns the maximum sequence length (i.e. max capacity) of the cache object""" + raise NotImplementedError("Make sure to implement `get_max_cache_shape` in a subclass.") + + def get_usable_length(self, new_seq_length: int, layer_idx: Optional[int] = 0) -> int: + """Given the sequence length of the new inputs, returns the usable length of the cache.""" + # Cache without size limit -> all cache is usable + # Cache with size limit -> if the length cache plus the length of the new inputs is larger the maximum cache + # length, we will need to evict part of the cache (and thus not all cache is usable) + max_length = self.get_max_cache_shape() + previous_seq_length = self.get_seq_length(layer_idx) + if max_length is not None and previous_seq_length + new_seq_length > max_length: + return max_length - new_seq_length + return previous_seq_length + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + for layer_idx in range(len(self.key_cache)): + if self.key_cache[layer_idx].numel(): + device = self.key_cache[layer_idx].device + self.key_cache[layer_idx] = self.key_cache[layer_idx].index_select(0, beam_idx.to(device)) + if self.value_cache[layer_idx].numel(): + device = self.value_cache[layer_idx].device + self.value_cache[layer_idx] = self.value_cache[layer_idx].index_select(0, beam_idx.to(device)) + + @property + def seen_tokens(self): + logger.warning_once( + "The `seen_tokens` attribute is deprecated and will be removed in v4.41. Use the `cache_position` " + "model input instead." + ) + if hasattr(self, "_seen_tokens"): + return self._seen_tokens + else: + return None + + +@dataclass +class CacheConfig: + """ + Base class for cache configs + """ + + cache_implementation: None + + @classmethod + def from_dict(cls, config_dict, **kwargs): + """ + Constructs a CacheConfig instance from a dictionary of parameters. + Args: + config_dict (Dict[str, Any]): Dictionary containing configuration parameters. + **kwargs: Additional keyword arguments to override dictionary values. + + Returns: + CacheConfig: Instance of CacheConfig constructed from the dictionary. + """ + config = cls(**config_dict) + to_remove = [] + for key, value in kwargs.items(): + if hasattr(config, key): + setattr(config, key, value) + to_remove.append(key) + for key in to_remove: + kwargs.pop(key, None) + return config + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_json_file + def to_json_file(self, json_file_path: Union[str, os.PathLike]): + """ + Save this instance to a JSON file. + + Args: + json_file_path (`str` or `os.PathLike`): + Path to the JSON file in which this configuration instance's parameters will be saved. + use_diff (`bool`, *optional*, defaults to `True`): + If set to `True`, only the difference between the config instance and the default + `QuantizationConfig()` is serialized to JSON file. + """ + with open(json_file_path, "w", encoding="utf-8") as writer: + config_dict = self.to_dict() + json_string = json.dumps(config_dict, indent=2, sort_keys=True) + "\n" + + writer.write(json_string) + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.to_dict + def to_dict(self) -> Dict[str, Any]: + """ + Serializes this instance to a Python dictionary. Returns: + `Dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + return copy.deepcopy(self.__dict__) + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__iter__ + def __iter__(self): + """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" + for attr, value in copy.deepcopy(self.__dict__).items(): + yield attr, value + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.__repr__ + def __repr__(self): + return f"{self.__class__.__name__} {self.to_json_string()}" + + def to_json_string(self): + """ + Serializes this instance to a JSON formatted string. + Returns: + str: JSON formatted string representing the configuration instance. + """ + return json.dumps(self.__dict__, indent=2) + "\n" + + # Copied from transformers.utils.quantization_config.QuantizationConfigMixin.update + def update(self, **kwargs): + """ + Updates attributes of this class instance with attributes from `kwargs` if they match existing attributes, + returning all the unused kwargs. + + Args: + kwargs (`Dict[str, Any]`): + Dictionary of attributes to tentatively update this class. + + Returns: + `Dict[str, Any]`: Dictionary containing all the key-value pairs that were not used to update the instance. + """ + to_remove = [] + for key, value in kwargs.items(): + if hasattr(self, key): + setattr(self, key, value) + to_remove.append(key) + + # Remove all the attributes that were updated, without modifying the input dict + unused_kwargs = {key: value for key, value in kwargs.items() if key not in to_remove} + return unused_kwargs + + +@dataclass +class QuantizedCacheConfig(CacheConfig): + """ + Configuration class for quantized cache settings. + + Attributes: + backend (`str`, *optional*, defaults to `"quanto"`): + Backend to use when performing quantization, Can be one of [`quanto`, `HQQ`] + nbits (`Optional[int]`, *optional*, defaults to 4): + Number of bits, can be 2 or 4 for the `quanto` backend and one of [1, 2, 3, 4, 8] for the `HQQ` backend. Defaults to 2. + axis_key (`int`, *optional*, defaults to 0): + Axis over which to perform grouping for the key tensors. Can be [0, -1] for `quanto` backend and [0, 1] for `HQQ` backend. + axis_value (`int`, *optional*, defaults to 0): + Axis over which to perform grouping for the value tensors. Can be [0, -1] for `quanto` backend and [0, 1] for `HQQ` backend. + q_group_size (`Optional[int]`, *optional*, defaults to 64): + Size of the quantization group, should be a divisor of the model's hidden dimension. + Defaults to 64. + residual_length (`Optional[int]`, *optional*, defaults to 128): + Length of the residual cache which will always be stored in original precision. + Defaults to 128. + compute_dtype (`torch.dtype`, *optional*, defaults to `torch.float16`): + The default dtype used for computations in the model. Keys and Values will be cast to this dtype after dequantization. + device (`str`, *optional*, defaults to `"cpu"`): + Device on which to perform computations, should be same as the model's device. + """ + + def __init__( + self, + backend: str = "quanto", + nbits: Optional[int] = 4, + axis_key: Optional[int] = 0, + axis_value: Optional[int] = 0, + q_group_size: Optional[int] = 64, + residual_length: Optional[int] = 128, + compute_dtype: Optional[torch.dtype] = torch.float16, + device: Optional[str] = "cpu", + ): + self.backend = backend + self.nbits = nbits + self.axis_key = axis_key + self.axis_value = axis_value + self.q_group_size = q_group_size + self.residual_length = residual_length + self.compute_dtype = compute_dtype + self.device = device + + def validate(self): + """Validates if the arguments passed are correct""" + + incorrect_arg_msg = ( + "Some of the keys in `cache_config` are defined incorrectly. `{key}` should be {correct_value}` " + "but found {found_value}" + ) + # Check that the values are reasonable in general (nbits, axis) + # Later in QuantizedCache init we check if they are supported for that particular backend + if self.nbits not in [1, 2, 3, 4, 8]: + raise ValueError( + incorrect_arg_msg.format( + key="nbits", + correct_value="2 or 4 or 8", + found_value=self.nbits, + ), + ) + if self.q_group_size <= 0: + raise ValueError( + incorrect_arg_msg.format( + key="q_group_size", + correct_value="a positive integer", + found_value=self.q_group_size, + ), + ) + if self.residual_length < 0: + raise ValueError( + incorrect_arg_msg.format( + key="residual_length", + correct_value="a positive integer", + found_value=self.residual_length, + ), + ) + + if self.axis_key not in [0, 1, -1]: + raise ValueError( + incorrect_arg_msg.format( + key="axis_key", + correct_value="`1` or `0`, `-1`", + found_value=self.axis_key, + ), + ) + + if self.axis_value not in [0, 1, -1]: + raise ValueError( + incorrect_arg_msg.format( + key="axis_value", + correct_value="`1` or `0` or `-1`", + found_value=self.axis_value, + ), + ) + + +@dataclass +class StaticCacheConfig(CacheConfig): + """ + Configuration class for static cache settings. + """ + + cache_implementation = "static" + + def __init__(self, batch_size: int, max_cache_len: int, device="cpu"): + self.batch_size = batch_size + self.max_cache_len = max_cache_len + self.device = device + + def validate(self): + """Validates if the arguments passed are correct""" + + incorrect_arg_msg = ( + "Some of the keys in `cache_config` are defined incorrectly. `{key}` should be {correct_value}` " + "but found {found_value}" + ) + + if self.batch_size <= 0: + raise ValueError( + incorrect_arg_msg.format( + key="batch_size", + correct_value="> 0", + found_value=self.batch_size, + ), + ) + + if self.max_cache_len <= 0: + raise ValueError( + incorrect_arg_msg.format( + key="max_cache_len", + correct_value="> 0", + found_value=self.max_cache_len, + ), + ) + + +class DynamicCache(Cache): + """ + A cache that grows dynamically as more tokens are generated. This is the default for generative models. + + It stores the Key and Value states as a list of tensors, one for each layer. The expected shape for each tensor is + `[batch_size, num_heads, seq_len, head_dim]`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, DynamicCache + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + + >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> past_key_values = DynamicCache() + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + DynamicCache() + ``` + """ + + def __init__(self, _distributed_cache_data: Optional[Iterable] = None) -> None: + super().__init__() + self._seen_tokens = 0 # Used in `generate` to keep tally of how many tokens the cache has seen + self.key_cache: List[torch.Tensor] = [] + self.key_cache_pack: List[torch.Tensor] = [] + self.key_cache_params: List[torch.Tensor] = [] + + self.value_cache: List[torch.Tensor] = [] + self.value_cache_pack: List[torch.Tensor] = [] + self.value_cache_params: List[torch.Tensor] = [] + + # `_distributed_cache_data` was originally added for compatibility with `torch.distributed` (DDP). See #36121 + # and #36373 for more information. In a nutshell, it is `map(gather_map, zip(*caches))`, i.e. each item in the + # iterable contains the key and value states for a layer gathered across replicas by torch.distributed + # (shape=[global batch size, num_heads, seq_len, head_dim]). + # WARNING: `_distributed_cache_data` must be the first argument in `__init__`, otherwise we'll break + # compatibility. The name of the argument doesn't matter. + if _distributed_cache_data is not None: + for key_states, value_states in _distributed_cache_data: + self.key_cache.append(key_states) + self.value_cache.append(value_states) + + def __getitem__(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Support for backwards-compatible `past_key_value` indexing, e.g. `past_key_value[0][0].shape[2]` to get the + sequence length. + """ + if layer_idx < len(self): + return (self.key_cache[layer_idx], self.value_cache[layer_idx]) + else: + raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") + + def __iter__(self): + """ + Support for backwards-compatible `past_key_value` iteration, e.g. `for x in past_key_value:` to iterate over + keys and values + """ + for layer_idx in range(len(self)): + yield (self.key_cache[layer_idx], self.value_cache[layer_idx]) + + def __len__(self): + """ + Support for backwards-compatible `past_key_value` length, e.g. `len(past_key_value)`. This value corresponds + to the number of layers in the model. + """ + return len(self.key_cache) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`. + + Return: + A tuple containing the updated key and value states. + """ + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + + # Update the cache + if key_states is not None: + if len(self.key_cache) <= layer_idx: + # There may be skipped layers, fill them with empty lists + for _ in range(len(self.key_cache), layer_idx): + self.key_cache.append(torch.tensor([])) + self.value_cache.append(torch.tensor([])) + self.key_cache.append(key_states) + self.value_cache.append(value_states) + elif ( + not self.key_cache[layer_idx].numel() # prefers not t.numel() to len(t) == 0 to export the model + ): # fills previously skipped layers; checking for tensor causes errors + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def update_residual( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`. + + Return: + A tuple containing the updated key and value states. + """ + # Update the cache + if key_states is not None: + + # Update the number of seen tokens + # if layer_idx == 0: + # self._seen_tokens += value_states.shape[-3] + + if len(self.key_cache) <= layer_idx: + # There may be skipped layers, fill them with empty lists + for _ in range(len(self.key_cache), layer_idx): + self.key_cache.append([]) + self.value_cache.append([]) + self.key_cache.append(key_states) + self.value_cache.append(value_states) + elif ( + len(self.key_cache[layer_idx]) == 0 + ): # fills previously skipped layers; checking for tensor causes errors + self.key_cache[layer_idx] = key_states + self.value_cache[layer_idx] = value_states + else: + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-3) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-3) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def update_pack( + self, + key_pack: torch.Tensor, + key_params: torch.Tensor, + value_pack: torch.Tensor, + value_params: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. No additional arguments are used in `DynamicCache`. + + Return: + A tuple containing the updated key and value states. + """ + # Update the cache + if key_pack is not None: + # Update the number of seen tokens + # if layer_idx == 0: + # self._seen_tokens += value_pack.shape[-3] + + if len(self.key_cache_pack) <= layer_idx: + # There may be skipped layers, fill them with empty lists + for _ in range(len(self.key_cache_pack), layer_idx): + self.key_cache_pack.append([]) + self.value_cache_pack.append([]) + self.key_cache_params.append([]) + self.value_cache_params.append([]) + self.key_cache_pack.append(key_pack) + self.value_cache_pack.append(value_pack) + self.key_cache_params.append(key_params) + self.value_cache_params.append(value_params) + elif ( + len(self.key_cache_pack[layer_idx]) == 0 + ): # fills previously skipped layers; checking for tensor causes errors + self.key_cache_pack[layer_idx] = key_pack + self.value_cache_pack[layer_idx] = value_pack + self.key_cache_params[layer_idx] = key_params + self.value_cache_params[layer_idx] = value_params + else: + self.key_cache_pack[layer_idx] = torch.cat([self.key_cache_pack[layer_idx], key_pack], dim=-3).contiguous() + self.value_cache_pack[layer_idx] = torch.cat([self.value_cache_pack[layer_idx], value_pack], dim=-3).contiguous() + self.key_cache_params[layer_idx] = torch.cat([self.key_cache_params[layer_idx], key_params], dim=-3).contiguous() + self.value_cache_params[layer_idx] = torch.cat([self.value_cache_params[layer_idx], value_params], dim=-1).contiguous() + + return self.key_cache_pack[layer_idx], self.key_cache_params[layer_idx], self.value_cache_pack[layer_idx], self.value_cache_params[layer_idx] + + def clear_residual(self, layer_idx: int): + self.key_cache[layer_idx] = [] + self.value_cache[layer_idx] = [] + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # TODO: deprecate this function in favor of `cache_position` + is_empty_layer = ( + len(self.key_cache) == 0 # no cache in any layer + or len(self.key_cache) <= layer_idx # skipped `layer_idx` and hasn't run a layer with cache after it + or not self.key_cache[layer_idx].numel() # the layer has no cache + ) + layer_seq_length = self.key_cache[layer_idx].shape[-2] if not is_empty_layer else 0 + return layer_seq_length + + def get_max_cache_shape(self) -> Optional[int]: + """Returns the maximum sequence length of the cache object. DynamicCache does not have a maximum length.""" + return None + + def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]: + """Converts the `DynamicCache` instance into the its equivalent in the legacy cache format. Used for + backward compatibility.""" + legacy_cache = () + for layer_idx in range(len(self)): + legacy_cache += ((self.key_cache[layer_idx], self.value_cache[layer_idx]),) + return legacy_cache + + @classmethod + def from_legacy_cache( + cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor, torch.FloatTensor]]] = None + ) -> "DynamicCache": + """Converts a cache in the legacy cache format into an equivalent `DynamicCache`. Used for + backward compatibility.""" + cache = cls() + if past_key_values is not None: + for layer_idx in range(len(past_key_values)): + key_states, value_states = past_key_values[layer_idx] + cache.update(key_states, value_states, layer_idx) + return cache + + def crop(self, max_length: int): + """Crop the past key values up to a new `max_length` in terms of tokens. `max_length` can also be + negative to remove `max_length` tokens. This is used in assisted decoding and contrastive search.""" + # In case it is negative + if max_length < 0: + max_length = self.get_seq_length() - abs(max_length) + + if self.get_seq_length() <= max_length: + return + + self._seen_tokens = max_length + for idx in range(len(self.key_cache)): + if self.key_cache[idx].numel(): + self.key_cache[idx] = self.key_cache[idx][..., :max_length, :] + self.value_cache[idx] = self.value_cache[idx][..., :max_length, :] + + def batch_split(self, full_batch_size: int, split_size: int) -> List["DynamicCache"]: + """Split the current instance into a list of `DynamicCache` by the batch size. This will be used by + `_split_model_inputs()` in `generation.utils`""" + out = [] + for i in range(0, full_batch_size, split_size): + current_split = DynamicCache() + current_split._seen_tokens = self._seen_tokens + current_split.key_cache = [tensor[i : i + split_size] for tensor in self.key_cache] + current_split.value_cache = [tensor[i : i + split_size] for tensor in self.value_cache] + out.append(current_split) + return out + + @classmethod + def from_batch_splits(cls, splits: List["DynamicCache"]) -> "DynamicCache": + """This is the opposite of the above `batch_split()` method. This will be used by `stack_model_outputs` in + `generation.utils`""" + cache = cls() + for idx in range(len(splits[0])): + key_cache = [current.key_cache[idx] for current in splits if current.key_cache[idx].numel()] + value_cache = [current.value_cache[idx] for current in splits if current.value_cache[idx].numel()] + if key_cache != []: + layer_keys = torch.cat(key_cache, dim=0) + layer_values = torch.cat(value_cache, dim=0) + cache.update(layer_keys, layer_values, idx) + return cache + + def batch_repeat_interleave(self, repeats: int): + """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search.""" + for layer_idx in range(len(self)): + self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(repeats, dim=0) + self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(repeats, dim=0) + + def batch_select_indices(self, indices: torch.Tensor): + """Only keep the `indices` in the batch dimension of the cache. Used in contrastive search.""" + for layer_idx in range(len(self)): + self.key_cache[layer_idx] = self.key_cache[layer_idx][indices, ...] + self.value_cache[layer_idx] = self.value_cache[layer_idx][indices, ...] + + +# Utilities for `DynamicCache` <> torch.export support +def _flatten_dynamic_cache( + dynamic_cache: DynamicCache, +): + """Flattens DynamicCache into flat list of tensors for `torch.export.export` to consume""" + if not isinstance(dynamic_cache, DynamicCache): + raise RuntimeError("This pytree flattening function should only be applied to DynamicCache") + + if not is_torch_greater_or_equal_than_2_6: + logger.warning_once( + "DynamicCache + torch.export is tested on torch 2.6.0+ and may not work on earlier versions." + ) + + # NOTE it seems _seen_tokens is deprecated, so probably doesn't need tracking + dictionary = { + "key_cache": getattr(dynamic_cache, "key_cache"), + "value_cache": getattr(dynamic_cache, "value_cache"), + } + return torch.utils._pytree._dict_flatten(dictionary) + + +def _flatten_with_keys_dynamic_cache(dynamic_cache: DynamicCache): + dictionary = { + "key_cache": getattr(dynamic_cache, "key_cache"), + "value_cache": getattr(dynamic_cache, "value_cache"), + } + return torch.utils._pytree._dict_flatten_with_keys(dictionary) + + +def _unflatten_dynamic_cache( + values, + context: torch.utils._pytree.Context, +): + dictionary = torch.utils._pytree._dict_unflatten(values, context) + cache = DynamicCache() + for k, v in dictionary.items(): + setattr(cache, k, v) + return cache + + +def _flatten_dynamic_cache_for_fx(cache, spec): + dictionary = { + "key_cache": getattr(cache, "key_cache"), + "value_cache": getattr(cache, "value_cache"), + } + return torch.utils._pytree.tree_flatten(dictionary)[0] + + +if is_torch_greater_or_equal("2.3"): + torch.utils._pytree.register_pytree_node( + DynamicCache, + _flatten_dynamic_cache, + _unflatten_dynamic_cache, + serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}", + flatten_with_keys_fn=_flatten_with_keys_dynamic_cache, + ) + # TODO (tmanlaibaatar) This won't be needed in torch 2.7. + torch.fx._pytree.register_pytree_flatten_spec(DynamicCache, _flatten_dynamic_cache_for_fx) + + +class OffloadedCache(DynamicCache): + """ + A drop-in replacement for DynamicCache that conserves accelerator(GPU, XPU) memory at the expense of more CPU memory. + Useful for generating from models with very long context. + + In addition to the default accelerator stream, where all forward() computations happen, + this class uses another stream, the prefetch stream, which it creates itself. + Since scheduling of operations on separate streams happens independently, this class uses + the prefetch stream to asynchronously prefetch the KV cache of layer k+1 when layer k is executing. + The movement of the layer k-1 cache to the CPU is handled by the default stream as a simple way to + ensure the eviction is scheduled after all computations on that cache are finished. + """ + + def __init__(self) -> None: + if not ( + torch.cuda.is_available() + or (is_torch_greater_or_equal("2.7", accept_dev=True) and torch.xpu.is_available()) + ): + raise RuntimeError( + "OffloadedCache can only be used with a GPU" + + (" or XPU" if is_torch_greater_or_equal("2.7", accept_dev=True) else "") + ) + + super().__init__() + self.original_device = [] + self.prefetch_stream = None + self.prefetch_stream = ( + torch.Stream() if is_torch_greater_or_equal("2.7", accept_dev=True) else torch.cuda.Stream() + ) + self.beam_idx = None # used to delay beam search operations + + def prefetch_layer(self, layer_idx: int): + "Starts prefetching the next layer cache" + if layer_idx < len(self): + with ( + self.prefetch_stream + if is_torch_greater_or_equal("2.7", accept_dev=True) + else torch.cuda.stream(self.prefetch_stream) + ): + # Prefetch next layer tensors to GPU + device = self.original_device[layer_idx] + self.key_cache[layer_idx] = self.key_cache[layer_idx].to(device, non_blocking=True) + self.value_cache[layer_idx] = self.value_cache[layer_idx].to(device, non_blocking=True) + + def evict_previous_layer(self, layer_idx: int): + "Moves the previous layer cache to the CPU" + if len(self) > 2: + # We do it on the default stream so it occurs after all earlier computations on these tensors are done + prev_layer_idx = (layer_idx - 1) % len(self) + self.key_cache[prev_layer_idx] = self.key_cache[prev_layer_idx].to("cpu", non_blocking=True) + self.value_cache[prev_layer_idx] = self.value_cache[prev_layer_idx].to("cpu", non_blocking=True) + + def __getitem__(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + "Gets the cache for this layer to the device. Prefetches the next and evicts the previous layer." + if layer_idx < len(self): + # Evict the previous layer if necessary + if is_torch_greater_or_equal("2.7", accept_dev=True): + torch.accelerator.current_stream().synchronize() + else: + torch.cuda.current_stream().synchronize() + self.evict_previous_layer(layer_idx) + # Load current layer cache to its original device if not already there + original_device = self.original_device[layer_idx] + self.prefetch_stream.synchronize() + key_tensor = self.key_cache[layer_idx] + value_tensor = self.value_cache[layer_idx] + # Now deal with beam search ops which were delayed + if self.beam_idx is not None: + self.beam_idx = self.beam_idx.to(original_device) + key_tensor = key_tensor.index_select(0, self.beam_idx) + value_tensor = value_tensor.index_select(0, self.beam_idx) + # Prefetch the next layer + self.prefetch_layer((layer_idx + 1) % len(self)) + return (key_tensor, value_tensor) + else: + raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Saves the beam indices and reorders the cache when the tensor is back to its device.""" + # We delay this operation until the tensors are back to their original + # device because performing torch.index_select on the CPU is very slow + del self.beam_idx + self.beam_idx = beam_idx.clone() + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. No additional arguments are used in `OffloadedCache`. + Return: + A tuple containing the updated key and value states. + """ + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + + # Update the cache + if len(self.key_cache) < layer_idx: + raise ValueError("OffloadedCache does not support model usage where layers are skipped. Use DynamicCache.") + elif len(self.key_cache) == layer_idx: + self.key_cache.append(key_states) + self.value_cache.append(value_states) + self.original_device.append(key_states.device) + self.evict_previous_layer(layer_idx) + else: + key_tensor, value_tensor = self[layer_idx] + self.key_cache[layer_idx] = torch.cat([key_tensor, key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([value_tensor, value_states], dim=-2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + # According to https://docs.python.org/3/library/exceptions.html#NotImplementedError + # if a method is not supposed to be supported in a subclass we should set it to None + from_legacy_cache = None + + to_legacy_cache = None + + +class QuantizedCache(DynamicCache): + """ + A quantizer cache similar to what is described in the [KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache paper](https://arxiv.org/abs/2402.02750). + It allows the model to generate longer sequence length without allocating too much memory for Key and Value cache by applying quantization. + + The cache has two types of storage, one for original precision and one for the quantized cache. A `residual length` is set as a maximum capacity for the + original precision cache. When the length goes beyond maximum capacity, the original precision cache is discarded and moved into the quantized cache. The + quantization is done per-channel with a set `q_group_size` for both Keys and Values, in contrast to what was described in the paper. + + It stores Keys and Values a list of quantized tensors (tuples in case we need to store metadata), one for each layer. Additionally, it stores the Key and + Value in original precision states as a list of tensors, one for each layer. The size of each tensor + is `[batch_size, num_heads, seq_len - residual_length, head_dim]` + """ + + def __init__(self, cache_config: QuantizedCacheConfig) -> None: + super().__init__() + self._quantized_key_cache: List[torch.Tensor] = [] + self._quantized_value_cache: List[torch.Tensor] = [] + + self.nbits = cache_config.nbits + self.residual_length = cache_config.residual_length + self.q_group_size = cache_config.q_group_size + self.axis_key = cache_config.axis_key + self.axis_value = cache_config.axis_value + self.compute_dtype = cache_config.compute_dtype + self.device = cache_config.device + + super().__init__() + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + + if len(self.key_cache) < layer_idx: + raise ValueError("QuantizedCache does not support model usage where layers are skipped. Use DynamicCache.") + elif len(self.key_cache) == layer_idx: + self._quantized_key_cache.append(self._quantize(key_states.contiguous(), axis=self.axis_key)) + self._quantized_value_cache.append(self._quantize(value_states.contiguous(), axis=self.axis_value)) + self.key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device)) + self.value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device)) + keys_to_return, values_to_return = key_states, value_states + else: + dequant_key = self._dequantize(self._quantized_key_cache[layer_idx]) + dequant_value = self._dequantize(self._quantized_value_cache[layer_idx]) + keys_to_return = [dequant_key, self.key_cache[layer_idx], key_states] + values_to_return = [dequant_value, self.value_cache[layer_idx], value_states] + + keys_to_return = torch.cat(keys_to_return, dim=-2) + values_to_return = torch.cat(values_to_return, dim=-2) + if ( + self.key_cache[layer_idx].dim() == 4 + and self.key_cache[layer_idx].shape[-2] + 1 >= self.residual_length + ): + self._quantized_key_cache[layer_idx] = self._quantize(keys_to_return.contiguous(), axis=self.axis_key) + self._quantized_value_cache[layer_idx] = self._quantize( + values_to_return.contiguous(), axis=self.axis_value + ) + self.key_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device) + self.value_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device) + else: + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) + + return keys_to_return, values_to_return + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + if len(self.key_cache) <= layer_idx: + return 0 + # since we cannot get the seq_length of each layer directly and rely on `_seen_tokens` which is + # updated every "layer_idx" == 0, this is a hack to get the actual seq_length for the given layer_idx + # this part of code otherwise fails when used to verify attn_weight shape in some models + return self._seen_tokens if layer_idx == 0 else self._seen_tokens - 1 + + def _quantize(self, tensor, axis): + """Quantizes a key/value using a defined quantization method.""" + raise NotImplementedError("Make sure to implement `_quantize` in a subclass.") + + def _dequantize(self, q_tensor): + """Dequantizes back the tensor that was quantized by `self._quantize()`""" + raise NotImplementedError("Make sure to implement `_dequantize` in a subclass.") + + +class QuantoQuantizedCache(QuantizedCache): + """ + Quantized Cache class that uses `quanto` as a backend to perform quantization. Current implementation supports `int2` and `int4` dtypes only. + + Parameters: + cache_config (`QuantizedCacheConfig`): + A configuration containing all the arguments to be used by the quantizer, including axis, qtype and group size. + + Example: + + ```python + >>> # Run pip install quanto first if you don't have it yet + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, QuantoQuantizedCache, QuantizedCacheConfig + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + + >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> cache_config = QuantizedCacheConfig(nbits=4) + >>> past_key_values = QuantoQuantizedCache(cache_config=cache_config) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + QuantoQuantizedCache() + ``` + """ + + def __init__(self, cache_config: CacheConfig) -> None: + super().__init__(cache_config) + + if is_optimum_quanto_available(): + optimum_quanto_version = version.parse(importlib.metadata.version("optimum-quanto")) + if optimum_quanto_version <= version.parse("0.2.5"): + raise ImportError( + f"You need optimum-quanto package version to be greater or equal than 0.2.5 to use `QuantoQuantizedCache`. Detected version {optimum_quanto_version}." + ) + from optimum.quanto import MaxOptimizer, qint2, qint4 + + if self.nbits not in [2, 4]: + raise ValueError(f"`nbits` for `quanto` backend has to be one of [`2`, `4`] but got {self.nbits}") + + if self.axis_key not in [0, -1]: + raise ValueError(f"`axis_key` for `quanto` backend has to be one of [`0`, `-1`] but got {self.axis_key}") + + if self.axis_value not in [0, -1]: + raise ValueError( + f"`axis_value` for `quanto` backend has to be one of [`0`, `-1`] but got {self.axis_value}" + ) + + self.qtype = qint4 if self.nbits == 4 else qint2 + self.optimizer = MaxOptimizer() # hardcode as it's the only one for per-channel quantization + + def _quantize(self, tensor, axis): + # We have two different API since in optimum-quanto, we don't use AffineQuantizer anymore + if is_optimum_quanto_available(): + from optimum.quanto import quantize_weight + + scale, zeropoint = self.optimizer(tensor, self.qtype, axis, self.q_group_size) + qtensor = quantize_weight(tensor, self.qtype, axis, scale, zeropoint, self.q_group_size) + return qtensor + + def _dequantize(self, qtensor): + return qtensor.dequantize() + + +class HQQQuantizedCache(QuantizedCache): + """ + Quantized Cache class that uses `HQQ` as a backend to perform quantization. Current implementation supports `int2`, `int4`, `int8` dtypes. + + Parameters: + cache_config (`QuantizedCacheConfig`): + A configuration containing all the arguments to be used by the quantizer, including axis, qtype and group size. + + Example: + + ```python + >>> # Run pip install hqq first if you don't have it yet + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, HQQQuantizedCache, QuantizedCacheConfig + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + + >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> cache_config = QuantizedCacheConfig(nbits=4, axis_key=1, axis_value=1) + >>> past_key_values = HQQQuantizedCache(cache_config=cache_config) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + HQQQuantizedCache() + ``` + """ + + def __init__(self, cache_config: CacheConfig) -> None: + super().__init__(cache_config) + if self.nbits not in [1, 2, 3, 4, 8]: + raise ValueError( + f"`nbits` for `HQQ` backend has to be one of [`1`, `2`, `3`, `4`, `8`] but got {self.nbits}" + ) + + if self.axis_key not in [0, 1]: + raise ValueError(f"`axis_key` for `HQQ` backend has to be one of [`0`, `1`] but got {self.axis_key}") + + if self.axis_value not in [0, 1]: + raise ValueError(f"`axis_value` for `HQQ` backend has to be one of [`0`, `1`] but got {self.axis_value}") + + self.quantizer = HQQQuantizer + + def _quantize(self, tensor, axis): + qtensor, meta = self.quantizer.quantize( + tensor, + axis=axis, + device=self.device, + compute_dtype=self.compute_dtype, + nbits=self.nbits, + group_size=self.q_group_size, + ) + meta["compute_dtype"] = self.compute_dtype + self.quantizer.cuda(qtensor, meta=meta, device=self.device) # Move to device and cast to dtype + meta["scale"] = meta["scale"].to(qtensor.device) + meta["zero"] = meta["zero"].to(qtensor.device) + return qtensor, meta + + def _dequantize(self, qtensor): + quant_tensor, meta = qtensor + tensor = self.quantizer.dequantize(quant_tensor, meta) + return tensor + + +class SinkCache(Cache): + """ + Deprecated. + + A cache that as described in the [Attention Sinks paper](https://arxiv.org/abs/2309.17453). It allows the model to + generate beyond the length of its context window, without losing fluency in the conversation. As it discards past + tokens, the model will lose the ability to generate tokens that depend on the context that was discarded. + + It stores the Key and Value states as a list of tensors, one for each layer. The expected shape for each tensor is + `[batch_size, num_heads, seq_len, head_dim]`. + + Parameters: + window_length (`int`): + The length of the context window. + num_sink_tokens (`int`): + The number of sink tokens. See the original paper for more information. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, SinkCache + + >>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct") + + >>> inputs = tokenizer(text="My name is Qwen2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> past_key_values = SinkCache(window_length=256, num_sink_tokens=4) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + SinkCache() + ``` + """ + + is_sliding = True + + def __init__(self, window_length: int, num_sink_tokens: int) -> None: + super().__init__() + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + self.window_length = window_length + self.num_sink_tokens = num_sink_tokens + self.cos_sin_rerotation_cache = {} + self._cos_cache = None + self._sin_cache = None + self._seen_tokens = 0 # Used in `generate` to keep tally of how many tokens the cache has seen + + warnings.warn( + "`SinkCache` is deprecated and will be removed in v4.53.0. You can achieve similar functionality by " + "using a model with a sliding window attention mechanism, or by expanding RoPE and optionally using an " + "offloaded cache implementation.", + FutureWarning, + ) + + @staticmethod + def _rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _apply_key_rotary_pos_emb( + self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> torch.Tensor: + rotated_key_states = (key_states * cos) + (self._rotate_half(key_states) * sin) + return rotated_key_states + + def _get_rerotation_cos_sin( + self, key_states: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + if key_states.shape[-2] not in self.cos_sin_rerotation_cache: + # Upcast to float32 temporarily for better accuracy + cos = cos.to(torch.float32) + sin = sin.to(torch.float32) + + # Compute the cos and sin required for back- and forward-rotating to one position earlier in the sequence + original_cos = cos[self.num_sink_tokens + key_states.shape[-2] :] + shifted_cos = cos[self.num_sink_tokens : -key_states.shape[-2]] + original_sin = sin[self.num_sink_tokens + key_states.shape[-2] :] + shifted_sin = sin[self.num_sink_tokens : -key_states.shape[-2]] + rerotation_cos = original_cos * shifted_cos + original_sin * shifted_sin + rerotation_sin = -original_sin * shifted_cos + original_cos * shifted_sin + + self.cos_sin_rerotation_cache[key_states.shape[-2]] = ( + rerotation_cos.to(key_states.dtype).unsqueeze(0), + rerotation_sin.to(key_states.dtype).unsqueeze(0), + ) + return self.cos_sin_rerotation_cache[key_states.shape[-2]] + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # TODO: deprecate this function in favor of `cache_position` + # Workaround to make 'key_states.shape[-2] + past_key_value.get_seq_length(self.layer_idx)' <= window_length + if len(self.key_cache) <= layer_idx: + return 0 + return self.key_cache[layer_idx].shape[-2] + + def get_max_cache_shape(self) -> Optional[int]: + """Returns the maximum sequence length of the cache object, in case of SinkCache it is the window length.""" + return self.window_length + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. The following arguments can be used in `SinkCache`: `sin`, + `cos` and `partial_rotation_size`. These arguments are used with models using RoPE, to recompute the + rotation as the tokens are shifted. + + Return: + A tuple containing the updated key and value states. + """ + # Optional kwargs for `SinkCache` -- needed on models using RoPE. `partial_rotation_size` is used on models + # with partially rotated position embeddings, like Phi or Persimmon. + if cache_kwargs is None: + cache_kwargs = {} + sin = cache_kwargs.get("sin") + cos = cache_kwargs.get("cos") + partial_rotation_size = cache_kwargs.get("partial_rotation_size") + using_rope = cos is not None and sin is not None + + # Update the number of seen tokens + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + + # Update the sin/cos cache, which holds sin/cos values for all possible positions + if using_rope and layer_idx == 0: + # BC: some models still pass `sin`/`cos` with 2 dims. In those models, they are the full sin/cos. Remove + # after all RoPE models have a llama-like cache utilization. + if cos.dim() == 2: + self._cos_cache = cos + self._sin_cache = sin + else: + if self._cos_cache is None: + self._cos_cache = cos[0, ...] + self._sin_cache = sin[0, ...] + elif self._cos_cache.shape[0] < self.window_length: + self._cos_cache = torch.cat([self._cos_cache, cos[0, ...]], dim=0) + self._sin_cache = torch.cat([self._sin_cache, sin[0, ...]], dim=0) + + # [bsz, num_heads, seq_len, head_dim] + if len(self.key_cache) <= layer_idx: + # Empty cache + self.key_cache.append(key_states) + self.value_cache.append(value_states) + + elif key_states.shape[-2] + self.get_seq_length(layer_idx) < self.window_length: + # Growing cache + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) + + else: + # Shifting cache + keys_to_keep = self.key_cache[layer_idx][ + :, :, -self.window_length + self.num_sink_tokens + key_states.shape[-2] : + ] + + # On RoPE models, we need to recompute the Key rotation as the tokens are shifted + if using_rope: + rerotation_cos, rerotation_sin = self._get_rerotation_cos_sin( + key_states, self._cos_cache[: self.window_length], self._sin_cache[: self.window_length] + ) + if partial_rotation_size is not None: + keys_to_keep, keys_pass = ( + keys_to_keep[..., :partial_rotation_size], + keys_to_keep[..., partial_rotation_size:], + ) + keys_to_keep = self._apply_key_rotary_pos_emb(keys_to_keep, rerotation_cos, rerotation_sin) + if partial_rotation_size is not None: + keys_to_keep = torch.cat((keys_to_keep, keys_pass), dim=-1) + + # Concatenate sink tokens, shifted & rotated tokens (if needed), and new tokens + sink_keys = self.key_cache[layer_idx][:, :, : self.num_sink_tokens] + self.key_cache[layer_idx] = torch.cat([sink_keys, keys_to_keep, key_states], dim=-2) + + sink_values = self.value_cache[layer_idx][:, :, : self.num_sink_tokens] + values_to_keep = self.value_cache[layer_idx][ + :, :, -self.window_length + self.num_sink_tokens + value_states.shape[-2] : + ] + self.value_cache[layer_idx] = torch.cat([sink_values, values_to_keep, value_states], dim=-2) + + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + +class StaticCache(Cache): + """ + Static Cache class to be used with `torch.compile(model)` and `torch.export()`. + + Parameters: + config (`PretrainedConfig`): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a + smaller batch size is used. If you are manually setting the batch size, make sure to take into account the + number of beams if you are running beam search + max_cache_len (`int`, *optional*): + The maximum sequence length with which the model will be used. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): + The default `dtype` to use when initializing the layer. + layer_device_map (`Optional[Dict[int, Union[str, torch.device, int]]]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache + + >>> model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf") + + >>> inputs = tokenizer(text="My name is Llama", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = StaticCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + StaticCache() + ``` + """ + + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.float32, + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + super().__init__() + self.max_batch_size = max_batch_size + self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len + + # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads + self.head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads + + self._dtype = dtype + self.num_key_value_heads = ( + config.num_attention_heads + if getattr(config, "num_key_value_heads", None) is None + else config.num_key_value_heads + ) + + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + # Note: There will be significant perf decrease if switching to use 5D tensors instead. + cache_shape = (self.max_batch_size, self.num_key_value_heads, self.max_cache_len, self.head_dim) + device = torch.device(device) if device is not None else None + for idx in range(config.num_hidden_layers): + if layer_device_map is not None: + layer_device = layer_device_map[idx] + else: + layer_device = device + new_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device) + new_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device) + # Note: `mark_static_address` is used to tag the cache as a fixed data pointer, + # preventing compiled graph breaks when updating the cache. + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + self.key_cache.append(new_layer_key_cache) + self.value_cache.append(new_layer_value_cache) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + It is VERY important to index using a tensor, otherwise you introduce a copy to the device. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, `optional`): + Additional arguments for the cache subclass. The `StaticCache` needs the `cache_position` input + to know how where to write in the cache. + + Return: + A tuple containing the updated key and value states. + """ + if cache_kwargs is None: + cache_kwargs = {} + + key_states = key_states.to(self.key_cache[layer_idx].dtype) + value_states = value_states.to(self.value_cache[layer_idx].dtype) + return _static_cache_update( + self.key_cache[layer_idx], + self.value_cache[layer_idx], + key_states, + value_states, + cache_kwargs.get("cache_position"), + ) + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states that were seen by the model.""" + # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's + # limit the check to the first batch member and head dimension. + # TODO: deprecate this function in favor of `cache_position` + return (self.key_cache[layer_idx][0, 0].any(dim=-1)).sum() + + def get_max_cache_shape(self) -> Optional[int]: + return self.max_cache_len + + def reset(self): + """Resets the cache values while preserving the objects""" + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + +class SlidingWindowCache(StaticCache): + """ + Sliding Window Cache class to be used with `torch.compile` for models like Mistral that support sliding window attention. + Every time when we try to update the cache, we compute the `indices` based on `cache_position >= self.config.sliding_window - 1`, + if true(which means the cache can not hold all the old key value states and new states together because of the sliding window constraint), + we need to do a cycle shift based on `indices` to replace the oldest states by the new key value states passed in. + + The `to_shift` is only true once we are above sliding_window. Thus with `sliding_window==64`: + + indices = (slicing + to_shift[-1].sum()-1) % self.config.sliding_window + tensor([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, + 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, + 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, + 55, 56, 57, 58, 59, 60, 61, 62, 63, 0]) + + We overwrite the cache using these, then we always write at cache_position (clamped to `sliding_window`) + + Parameters: + config (`PretrainedConfig`): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a + smaller batch size is used. + max_cache_len (`int`, *optional*): + The maximum sequence length with which the model will be used. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): + The default `dtype` to use when initializing the layer. + layer_device_map (`Optional[Dict[int, Union[str, torch.device, int]]]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, SlidingWindowCache + + >>> model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3") + + >>> inputs = tokenizer(text="My name is Mistral", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = SlidingWindowCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + SlidingWindowCache() + ``` + """ + + is_sliding = True + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.float32, + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + if not hasattr(config, "sliding_window") or config.sliding_window is None: + raise ValueError( + "Setting `cache_implementation` to 'sliding_window' requires the model config supporting " + "sliding window attention, please check if there is a `sliding_window` field in the model " + "config and it's not set to None." + ) + max_cache_len = min(config.sliding_window, max_cache_len) + super().__init__( + config=config, + max_batch_size=max_batch_size, + max_cache_len=max_cache_len, + device=device, + dtype=dtype, + layer_device_map=layer_device_map, + ) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if cache_kwargs is None: + cache_kwargs = {} + cache_position = cache_kwargs.get("cache_position") + + if cache_position is None: + raise ValueError("`cache_position` must be provided for SlidingWindowCache.") + + key_states = key_states.to(self.key_cache[layer_idx].dtype) + value_states = value_states.to(self.value_cache[layer_idx].dtype) + + return _sliding_cache_update( + self.key_cache[layer_idx], + self.value_cache[layer_idx], + key_states, + value_states, + cache_position, + self.max_cache_len, + ) + + def get_max_cache_shape(self) -> Optional[int]: + return self.max_cache_len + + def reset(self): + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + +class EncoderDecoderCache(Cache): + """ + Base, abstract class for all encoder-decoder caches. Can be used to hold combinations of self-attention and + cross-attention caches. + + Example: + + ```python + >>> from transformers import AutoProcessor, AutoModelForCausalLM, DynamicCache, EncoderDecoderCache + + >>> model = AutoModelForCausalLM.from_pretrained("openai/whisper-small") + >>> processor = AutoProcessor.from_pretrained("openai/whisper-small") + + >>> inputs = processor(audio=YOUR-AUDIO, return_tensors="pt") + + >>> # Prepare cache classes for encoder and decoder and pass it to model's forward + >>> self_attention_cache = DynamicCache() + >>> cross_attention_cache = DynamicCache() + >>> past_key_values = EncoderDecoderCache(self_attention_cache, cross_attention_cache) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + EncoderDecoderCache() + ``` + + """ + + def __init__(self, self_attention_cache: Cache, cross_attention_cache: Cache): + super().__init__() + self.self_attention_cache = self_attention_cache + self.cross_attention_cache = cross_attention_cache + self.is_compileable = getattr(self.self_attention_cache, "is_compileable", False) + + self.is_updated = {} + for layer_idx in range(len(cross_attention_cache.key_cache)): + self.is_updated[layer_idx] = bool(cross_attention_cache.get_seq_length(layer_idx) > 0) + + def __getitem__(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Support for backwards-compatible `past_key_value` indexing, e.g. `past_key_value[0][0].shape[2]` to get the + sequence length. + """ + if layer_idx < len(self): + return ( + self.self_attention_cache.key_cache[layer_idx], + self.self_attention_cache.value_cache[layer_idx], + self.cross_attention_cache.key_cache[layer_idx], + self.cross_attention_cache.value_cache[layer_idx], + ) + else: + raise KeyError(f"Cache only has {len(self)} layers, attempted to access layer with index {layer_idx}") + + def __len__(self): + """ + Support for backwards-compatible `past_key_value` length, e.g. `len(past_key_value)`. This value corresponds + to the number of layers in the model. + """ + return len(self.self_attention_cache) + + def to_legacy_cache(self) -> Tuple[Tuple[torch.Tensor]]: + """Converts the `EncoderDecoderCache` instance into its equivalent in the legacy cache format.""" + legacy_cache = () + if len(self.cross_attention_cache) > 0: + for self_attn, cross_attn in zip( + self.self_attention_cache.to_legacy_cache(), self.cross_attention_cache.to_legacy_cache() + ): + legacy_cache += (self_attn + cross_attn,) + else: + legacy_cache = self.self_attention_cache.to_legacy_cache() + return legacy_cache + + @classmethod + def from_legacy_cache( + cls, past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None + ) -> "EncoderDecoderCache": + """Converts a cache in the legacy cache format into an equivalent `EncoderDecoderCache`.""" + cache = cls( + self_attention_cache=DynamicCache(), + cross_attention_cache=DynamicCache(), + ) + if past_key_values is not None: + for layer_idx in range(len(past_key_values)): + key_states, value_states = past_key_values[layer_idx][:2] + cache.self_attention_cache.update(key_states, value_states, layer_idx) + if len(past_key_values[layer_idx]) > 2: + key_states, value_states = past_key_values[layer_idx][2:] + cache.cross_attention_cache.update(key_states, value_states, layer_idx) + cache.is_updated[layer_idx] = True + return cache + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states. A layer index can be optionally passed.""" + # check if empty list because in case of static cache it will be a tensors and we can't check `if not torch.Tensor` + return self.self_attention_cache.get_seq_length(layer_idx) + + def reset(self): + if hasattr(self.self_attention_cache, "reset"): + self.self_attention_cache.reset() + if hasattr(self.cross_attention_cache, "reset"): + self.cross_attention_cache.reset() + elif not hasattr(self.self_attention_cache, "reset") and not hasattr(self.cross_attention_cache, "reset"): + raise ValueError( + "Neither self nor cross-attention cache have valid `.reset()` methods. `.reset()` should " + "only be called on compatible cache classes, such as `StaticCache` or `SlidingWindowCache`. " + f"Got {self.self_attention_cache.__str__()} for the self attention cache and " + f"{self.cross_attention_cache.__str__()} for the cross attention cache." + ) + for layer_idx in self.is_updated: + self.is_updated[layer_idx] = False + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorders the cache for beam search, given the selected beam indices.""" + self.self_attention_cache.reorder_cache(beam_idx) + self.cross_attention_cache.reorder_cache(beam_idx) + + def check_dynamic_cache(self, method: str): + if not ( + isinstance(self.self_attention_cache, DynamicCache) + and isinstance(self.cross_attention_cache, DynamicCache) + ): + raise ValueError( + f"`{method}` is only defined for dynamic cache, got {self.self_attention_cache.__str__()} for the self " + f"attention cache and {self.cross_attention_cache.__str__()} for the cross attention cache." + ) + + # TODO(gante, sanchit-gandhi): move following functionality into `.generate` + def crop(self, maximum_length: int): + """Crop the past key values up to a new `maximum_length` in terms of tokens. `maximum_length` can also be + negative to remove `maximum_length` tokens. This is used in assisted decoding and contrastive search.""" + self.check_dynamic_cache(self.crop.__name__) + self.self_attention_cache.crop(maximum_length) + + def batch_split(self, full_batch_size: int, split_size: int) -> "List[EncoderDecoderCache]": + """Split the current instance into a list of `DynamicCache` by the batch size. This will be used by + `_split_model_inputs()` in `generation.utils`""" + self.check_dynamic_cache(self.batch_split.__name__) + self_attention_cache = self.self_attention_cache.batch_split(full_batch_size, split_size) + cross_attention_cache = self.cross_attention_cache.batch_split(full_batch_size, split_size) + + out = [] + for self_attn, cross_attn in zip(self_attention_cache, cross_attention_cache): + out.append(EncoderDecoderCache(self_attn, cross_attn)) + return out + + @classmethod + def from_batch_splits(cls, splits: List["EncoderDecoderCache"]) -> "EncoderDecoderCache": + """This is the opposite of the above `batch_split()` method. This will be used by `stack_model_outputs` in + `generation.utils`""" + self_attention_cache = DynamicCache() + cross_attention_cache = DynamicCache() + for idx in range(len(splits[0])): + layer_keys = torch.cat([current.self_attention_cache.key_cache[idx] for current in splits], dim=0) + layer_values = torch.cat([current.self_attention_cache.value_cache[idx] for current in splits], dim=0) + self_attention_cache.update(layer_keys, layer_values, idx) + + layer_keys = torch.cat([current.cross_attention_cache.key_cache[idx] for current in splits], dim=0) + layer_values = torch.cat([current.cross_attention_cache.value_cache[idx] for current in splits], dim=0) + cross_attention_cache.update(layer_keys, layer_values, idx) + return cls(self_attention_cache, cross_attention_cache) + + def batch_repeat_interleave(self, repeats: int): + """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search.""" + self.check_dynamic_cache(self.batch_repeat_interleave.__name__) + self.self_attention_cache.batch_repeat_interleave(repeats) + self.cross_attention_cache.batch_repeat_interleave(repeats) + + def batch_select_indices(self, indices: torch.Tensor): + """Only keep the `indices` in the batch dimension of the cache. Used in contrastive search.""" + self.check_dynamic_cache(self.batch_select_indices.__name__) + self.self_attention_cache.batch_select_indices(indices) + self.cross_attention_cache.batch_select_indices(indices) + + +class HybridCache(Cache): + """ + Hybrid Cache class to be used with `torch.compile` for models that alternate between a local sliding window + attention and global attention in every other layer (originally implemented for Gemma2). + Under the hood, Hybrid Cache leverages ["SlidingWindowCache"] for sliding window attention and ["StaticCache"] + for global attention.For more information, see the documentation of each subcomponent cache class. + + Parameters: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a + smaller batch size is used. + max_cache_len (`int`, *optional*): + The maximum sequence length with which the model will be used. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (torch.dtype, *optional*, defaults to `torch.float32`): + The default `dtype` to use when initializing the layer. + layer_device_map (`Optional[Dict[int, Union[str, torch.device, int]]]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, HybridCache + + >>> model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b") + + >>> inputs = tokenizer(text="My name is Gemma", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = HybridCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + HybridCache() + ``` + """ + + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.float32, + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + super().__init__() + if not hasattr(config, "sliding_window") or config.sliding_window is None: + raise ValueError( + "Setting `cache_implementation` to 'hybrid' requires the model config supporting " + "sliding window attention, please check if there is a `sliding_window` field in the model " + "config and it's not set to None." + ) + self.max_cache_len = max_cache_len if max_cache_len is not None else config.max_position_embeddings + # Sliding layers can't be larger than the overall max cache len + self.sliding_window_len = min(config.sliding_window, self.max_cache_len) + self.max_batch_size = max_batch_size + # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads + self.head_dim = ( + config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads + ) + + self._dtype = dtype + self.num_key_value_heads = ( + config.num_attention_heads + if getattr(config, "num_key_value_heads", None) is None + else config.num_key_value_heads + ) + + layer_switch = config.sliding_window_pattern if hasattr(config, "sliding_window_pattern") else 2 # 2 is for BC + self.is_sliding_list = [bool((i + 1) % layer_switch) for i in range(config.num_hidden_layers)] + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + global_cache_shape = (self.max_batch_size, self.num_key_value_heads, self.max_cache_len, self.head_dim) + sliding_cache_shape = (self.max_batch_size, self.num_key_value_heads, self.sliding_window_len, self.head_dim) + device = torch.device(device) if device is not None else None + for i in range(config.num_hidden_layers): + if layer_device_map is not None: + layer_device = layer_device_map[i] + else: + layer_device = device + # Note: `mark_static_address` is used to tag the cache as an fixed data pointer, preventing cuda graph + # breaks when updating the cache. + cache_shape = sliding_cache_shape if self.is_sliding_list[i] else global_cache_shape + new_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device) + new_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=layer_device) + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + self.key_cache.append(new_layer_key_cache) + self.value_cache.append(new_layer_value_cache) + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if cache_kwargs is None: + cache_kwargs = {} + cache_position = cache_kwargs.get("cache_position") + if cache_position is None: + raise ValueError("`cache_position` must be provided for HybridCache.") + + is_sliding_layer = self.is_sliding_list[layer_idx] + + # These two `if` blocks are only reached in multigpu and if `layer_device_map` is not passed. They are used + # when the cache is initialized in the forward pass (e.g. Gemma2) + if self.key_cache[layer_idx].device != key_states.device: + self.key_cache[layer_idx] = self.key_cache[layer_idx].to(key_states.device) + if self.value_cache[layer_idx].device != value_states.device: + self.value_cache[layer_idx] = self.value_cache[layer_idx].to(value_states.device) + + k_cache = self.key_cache[layer_idx] + v_cache = self.value_cache[layer_idx] + key_states = key_states.to(k_cache.dtype) + value_states = value_states.to(v_cache.dtype) + + if is_sliding_layer: + return _sliding_cache_update( + k_cache, + v_cache, + key_states, + value_states, + cache_position, + k_cache.shape[2], # Use actual cache dim as max cache len + ) + else: + return _static_cache_update(k_cache, v_cache, key_states, value_states, cache_position) + + def get_max_cache_shape(self) -> Optional[int]: + return self.max_cache_len + + def get_seq_length(self, layer_idx: Optional[int] = 0): + # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's + # limit the check to the first batch member and head dimension. + # TODO: deprecate this function in favor of `cache_position` + if layer_idx != 0: + raise ValueError( + "`get_seq_length` on `HybridCache` may get inconsistent results depending on the layer index. " + "Using the `layer_idx` argument is not supported." + ) + return (self.key_cache[layer_idx][0, 0].any(dim=-1)).sum() + + def reset(self): + """Resets the cache values while preserving the objects""" + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + +class HybridChunkedCache(Cache): + """ + Hybrid Cache class to be used with `torch.compile` for models that alternate between a local sliding window + attention and global attention in every other layer, with support for chunked attention (originally implemented + for Llama4). + Under the hood, Hybrid Cache leverages ["SlidingWindowCache"] for sliding window attention and ["StaticCache"] + for global attention. For more information, see the documentation of each subcomponent cache class. + + Parameters: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a + smaller batch size is used. + max_cache_len (`int`, *optional*): + The maximum sequence length with which the model will be used. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (torch.dtype, *optional*, defaults to `torch.bfloat16`): + The default `dtype` to use when initializing the layer. + layer_device_map (`Optional[Dict[int, Union[str, torch.device, int]]]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, HybridCache + + >>> model = AutoModelForCausalLM.from_pretrained("google/gemma-2-2b") + >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b") + + >>> inputs = tokenizer(text="My name is Gemma", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = HybridCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values # access cache filled with key/values from generation + HybridCache() + ``` + """ + + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.bfloat16, + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + super().__init__() + if not hasattr(config, "sliding_window") or config.sliding_window is None: + self.sliding_window = getattr(config.get_text_config(), "attention_chunk_size", 8192) + else: + self.sliding_window = config.sliding_window + self.max_cache_len = max_cache_len + self._sliding_window_max_len = min(self.sliding_window, max_cache_len) + self.max_batch_size = max_batch_size + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self._dtype = dtype + + if hasattr(config.get_text_config(), "no_rope_layers"): + self.is_sliding = config.no_rope_layers + else: + layer_switch = getattr(config, "sliding_window_pattern", 2) + self.is_sliding = [bool((i + 1) % layer_switch) for i in range(config.num_hidden_layers)] + + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + self.cumulative_length = [0 for _ in range(config.num_hidden_layers)] + + def initialise_cache_layer(self, layer_idx, key_states): + if len(self.key_cache) > layer_idx: + return + + num_key_value_heads = key_states.shape[1] + device = key_states.device + global_cache_shape = (self.max_batch_size, num_key_value_heads, self.max_cache_len, self.head_dim) + sliding_cache_shape = (self.max_batch_size, num_key_value_heads, self._sliding_window_max_len, self.head_dim) + # Note: `mark_static_address` is used to tag the cache as an fixed data pointer, preventing cuda graph + # breaks when updating the cache. + cache_shape = sliding_cache_shape if self.is_sliding[layer_idx] else global_cache_shape + new_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=device) + new_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=device) + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + self.key_cache.append(new_layer_key_cache) + self.value_cache.append(new_layer_value_cache) + + def _sliding_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len): + cumulative_length = self.cumulative_length[layer_idx] + # Update it now that we saved the value above + self.cumulative_length[layer_idx] += key_states.shape[-2] + is_full = cumulative_length >= max_cache_len + if is_full: + full_key_states = torch.cat((k_out[:, :, 1:, :], key_states), dim=-2) + full_value_states = torch.cat((v_out[:, :, 1:, :], value_states), dim=-2) + # Fast decoding path -> here as the effective size is still sliding window, it is extremely important + # to return `self.key_cache[layer_idx]` and `self.value_cache[layer_idx]`, as they have the fixed address + # in memory (the values are the same as the full states, but not the address!!) + if key_states.shape[-2] == 1: + self.key_cache[layer_idx].copy_(full_key_states) + self.value_cache[layer_idx].copy_(full_value_states) + return self.key_cache[layer_idx], self.value_cache[layer_idx] + elif not is_full and cumulative_length + key_states.shape[2] > max_cache_len: + # Fast prefill path, no need to cat() in this case (which creates a copy even if cating from 0 dim) + if cumulative_length == 0: + full_key_states = key_states + full_value_states = value_states + else: + full_key_states = torch.cat((k_out[:, :, :cumulative_length, :], key_states), dim=-2) + full_value_states = torch.cat((v_out[:, :, :cumulative_length, :], value_states), dim=-2) + else: + self.key_cache[layer_idx].index_copy_(2, cache_position, key_states) + self.value_cache[layer_idx].index_copy_(2, cache_position, value_states) + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + self.key_cache[layer_idx].copy_(full_key_states[:, :, -max_cache_len:, :]) + self.value_cache[layer_idx].copy_(full_value_states[:, :, -max_cache_len:, :]) + # we should return the whole states instead of k_out, v_out to take the whole prompt + # into consideration when building kv cache instead of just throwing away tokens outside of the window + return full_key_states, full_value_states + + def _static_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len): + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + + self.key_cache[layer_idx] = k_out + self.value_cache[layer_idx] = v_out + return k_out, v_out + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if cache_kwargs is None: + cache_kwargs = {} + cache_position = cache_kwargs.get("cache_position") + self.initialise_cache_layer(layer_idx, key_states) + + k_out = self.key_cache[layer_idx] + v_out = self.value_cache[layer_idx] + key_states = key_states.to(k_out.dtype) + value_states = value_states.to(v_out.dtype) + + if self.is_sliding[layer_idx]: + update_fn = self._sliding_update + else: + update_fn = self._static_update + + return update_fn( + cache_position, + layer_idx, + key_states, + value_states, + k_out, + v_out, + k_out.shape[2], + ) + + def get_max_cache_shape(self) -> Optional[int]: + return self.max_cache_len + + def get_seq_length(self, layer_idx: Optional[int] = 0): + # Occupied cache == any slot in the 3rd dim (sequence length) holds a non-zero value. To save on compute, let's + # limit the check to the first batch member and head dimension. + # TODO: deprecate this function in favor of `cache_position` + if layer_idx != 0: + raise ValueError( + "`get_seq_length` on `HybridCache` may get inconsistent results depending on the layer index. " + "Using the `layer_idx` argument is not supported." + ) + if len(self.key_cache) == 0: + return 0 + return (self.key_cache[layer_idx][0, 0].any(dim=-1)).sum() + + def reset(self): + """Resets the cache values while preserving the objects""" + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + self.cumulative_length = [0 for _ in range(len(self.cumulative_length))] + + +class OffloadedHybridCache(HybridChunkedCache): + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int] = None, + device: Union[torch.device, str, None] = None, + dtype: torch.dtype = torch.bfloat16, + offload_device: Union[str, torch.device] = torch.device("cpu"), + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ): + super().__init__(config, max_batch_size, max_cache_len, device, dtype, layer_device_map) + + # TODO (joao): to enable this cache on multiple devicesuse the pattern from `OffloadedCache`, which keeps + # track of the original device of each layer + unique_devices = set(layer_device_map.values()) if layer_device_map else set() + if len(unique_devices) > 1: + raise ValueError(f"OffloadedHybridCache does not support multiple devices. Got devices: {unique_devices}") + + self.offload_device = torch.device(offload_device) + # Create new CUDA stream for parallel prefetching. + self._prefetch_stream = torch.cuda.Stream() if torch._C._get_accelerator().type == "cuda" else None + # Those will be dynamically created as the other layers (for TP) + self.device_key_cache = None + self.device_value_cache = None + # This gives the index of which on-device full layer to use (we need 2 to avoid race conditions when prefetching) + self.active_device_layer = 0 + + def initialise_cache_layer(self, layer_idx, key_states): + """Overridden to use the correct device if offloaded layer (and pin memory).""" + if len(self.key_cache) > layer_idx: + return + + num_key_value_heads = key_states.shape[1] + device = key_states.device if self.is_sliding[layer_idx] else self.offload_device + pin_memory = not self.is_sliding[layer_idx] + global_cache_shape = (self.max_batch_size, num_key_value_heads, self.max_cache_len, self.head_dim) + sliding_cache_shape = (self.max_batch_size, num_key_value_heads, self._sliding_window_max_len, self.head_dim) + # Note: `mark_static_address` is used to tag the cache as an fixed data pointer, preventing cuda graph + # breaks when updating the cache. + cache_shape = sliding_cache_shape if self.is_sliding[layer_idx] else global_cache_shape + new_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=device, pin_memory=pin_memory) + new_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=device, pin_memory=pin_memory) + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + self.key_cache.append(new_layer_key_cache) + self.value_cache.append(new_layer_value_cache) + + # Make sure to initialize the on-device layer if it does not already exist + if self.device_key_cache is None and not self.is_sliding[layer_idx]: + self.device_key_cache = [] + self.device_value_cache = [] + # We need 2 layers to avoid race conditions when prefetching the next one + for _ in range(2): + device_layer_key_cache = torch.zeros(cache_shape, dtype=self._dtype, device=key_states.device) + device_layer_value_cache = torch.zeros(cache_shape, dtype=self._dtype, device=key_states.device) + torch._dynamo.mark_static_address(new_layer_key_cache) + torch._dynamo.mark_static_address(new_layer_value_cache) + self.device_key_cache.append(device_layer_key_cache) + self.device_value_cache.append(device_layer_value_cache) + + def _static_update(self, cache_position, layer_idx, key_states, value_states, k_out, v_out, max_cache_len): + # Wait for prefetch stream if needed + if self._prefetch_stream is not None: + torch.cuda.default_stream(key_states.device).wait_stream(self._prefetch_stream) + + # Get correct on-device layer + k_out = self.device_key_cache[self.active_device_layer] + v_out = self.device_value_cache[self.active_device_layer] + + # Let's prefetch the next layer as soon as possible + self._prefetch_next_layer(layer_idx) + + # Copy to on-device layer + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + + # Copy to offloaded device + self.key_cache[layer_idx][:, :, cache_position] = key_states.to(self.offload_device) + self.value_cache[layer_idx][:, :, cache_position] = value_states.to(self.offload_device) + + return k_out, v_out + + def _prefetch_next_layer(self, layer_idx: int) -> None: + """Based on current layer_idx, prefetch next full layer to the device.""" + + # Switch the active layer + self.active_device_layer = 0 if self.active_device_layer == 1 else 1 + + # Find the next non-sliding layer + try: + next_layer = layer_idx + 1 + self.is_sliding[layer_idx + 1 :].index(False) + # In this case, we are at the last layer, and we go back to prefect the first one + except ValueError: + next_layer = self.is_sliding.index(False) + + # Alternate between two on-device caches. + if self._prefetch_stream is not None: + with torch.cuda.stream(self._prefetch_stream): + self._prefetch_layer_in_context(next_layer) + else: + self._prefetch_layer_in_context(next_layer) + + def _prefetch_layer_in_context(self, layer_idx: int) -> None: + """Performs the actual copy of the layer to device cache.""" + if len(self.key_cache) >= layer_idx: + self.device_key_cache[self.active_device_layer].copy_(self.key_cache[layer_idx], non_blocking=True) + self.device_value_cache[self.active_device_layer].copy_(self.value_cache[layer_idx], non_blocking=True) + # The layer was not yet initialized + else: + self.device_key_cache[self.active_device_layer].fill_(0.0) + self.device_value_cache[self.active_device_layer].fill_(0.0) + + +class MambaCache: + """ + Cache for mamba model which does not have attention mechanism and key value states. + + Arguments: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. Note that a new instance must be instantiated if a smaller batch size is used. + dtype (`torch.dtype`, *optional*, defaults to `torch.float16`): + The default `dtype` to use when initializing the layer. + device (`torch.device` or `str`, *optional*): + The device on which the cache should be initialized. Should be the same as the layer. + + Example: + + ```python + >>> from transformers import AutoTokenizer, MambaForCausalLM, MambaCache + + >>> model = MambaForCausalLM.from_pretrained("state-spaces/mamba-130m-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("state-spaces/mamba-130m-hf") + + >>> inputs = tokenizer(text="My name is Mamba", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> past_key_values = MambaCache(config=model.config, max_batch_size=1, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> outputs.past_key_values + MambaCache() + ``` + """ + + is_compileable = True + + # TODO (joao): add layer_device_map arg and update code in `generate` accordingly + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + dtype: torch.dtype = torch.float16, + device: Union[torch.device, str, None] = None, + ): + self.max_batch_size = max_batch_size + self._dtype = dtype + self.intermediate_size = config.intermediate_size + self.ssm_state_size = config.state_size + self.conv_kernel_size = config.conv_kernel + + self.conv_states: List[torch.Tensor] = [] + self.ssm_states: List[torch.Tensor] = [] + device = torch.device(device) if device is not None else None + for _ in range(config.num_hidden_layers): + conv_state: torch.Tensor = torch.zeros( + self.max_batch_size, + self.intermediate_size, + self.conv_kernel_size, + device=device, + dtype=self._dtype, + ) + ssm_state: torch.Tensor = torch.zeros( + self.max_batch_size, + self.intermediate_size, + self.ssm_state_size, + device=device, + dtype=self._dtype, + ) + + torch._dynamo.mark_static_address(conv_state) + torch._dynamo.mark_static_address(ssm_state) + self.conv_states.append(conv_state) + self.ssm_states.append(ssm_state) + + def update_conv_state( + self, layer_idx: int, new_conv_state: torch.Tensor, cache_position: torch.LongTensor + ) -> torch.Tensor: + # This `if` blocks is only reached in multigpu and if `layer_device_map` is not passed. It is used + # when the cache is initialized in the forward pass (e.g. Mamba) + if self.conv_states[layer_idx].device != new_conv_state.device: + self.conv_states[layer_idx] = self.conv_states[layer_idx].to(new_conv_state.device) + + conv_state = self.conv_states[layer_idx] + cache_position = cache_position.clamp(0, self.conv_kernel_size - 1) + + conv_state = conv_state.roll(shifts=-1, dims=-1) + conv_state[:, :, cache_position] = new_conv_state.to(device=conv_state.device, dtype=conv_state.dtype) + self.conv_states[layer_idx].zero_() + self.conv_states[layer_idx] += conv_state + return self.conv_states[layer_idx] + + def update_ssm_state(self, layer_idx: int, new_ssm_state: torch.Tensor): + self.ssm_states[layer_idx] = new_ssm_state.to(self.ssm_states[layer_idx].device) + return self.ssm_states[layer_idx] + + def reset(self): + for layer_idx in range(len(self.conv_states)): + # In-place ops prevent breaking the static address + self.conv_states[layer_idx].zero_() + self.ssm_states[layer_idx].zero_() + + +class OffloadedStaticCache(StaticCache): + """ + Static cache class to be used with `torch.compile(model)` that offloads to the CPU or + another device. + + Args: + config (`PretrainedConfig): + The configuration file defining the shape-related attributes required to initialize + the static cache. + max_batch_size (`int`): + The maximum batch size with which the model will be used. + max_cache_len (`int`): + The maximum sequence length with which the model will be used. + device (`Union[str, torch.device]`): + The device on which the cache should be initialized. If you're using more than 1 computation device, you + should pass the `layer_device_map` argument instead. + dtype (`torch.dtype`, *optional*): + The default `dtype` to use when initializing the cache. + offload_device (`Union[str, torch.device]`, *optional*, defaults to `cpu`): + The device to offload to. Defaults to CPU. + layer_device_map (`Dict[int, Union[str, torch.device, int]]`, *optional*): + Mapping between the layers and its device. This is required when you are manually initializing the cache + and the model is split between different gpus. You can know which layers mapped to which device by + checking the associated device_map: `model.hf_device_map`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, AutoModelForCausalLM, OffloadedStaticCache + + >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") + >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") + + >>> inputs = tokenizer(text="My name is GPT2", return_tensors="pt") + + >>> # Prepare a cache class and pass it to model's forward + >>> # Leave empty space for 10 new tokens, which can be used when calling forward iteratively 10 times to generate + >>> max_generated_length = inputs.input_ids.shape[1] + 10 + >>> past_key_values = OffloadedStaticCache(config=model.config, max_batch_size=1, max_cache_len=max_generated_length, device=model.device, dtype=model.dtype) + >>> outputs = model(**inputs, past_key_values=past_key_values, use_cache=True) + >>> past_kv_length = outputs.past_key_values # access cache filled with key/values from generation + ``` + """ + + is_compileable = True + + def __init__( + self, + config: PretrainedConfig, + max_batch_size: int, + max_cache_len: Optional[int], + device: Union[str, torch.device], + dtype: Optional[torch.dtype] = None, + offload_device: Union[str, torch.device] = torch.device("cpu"), + layer_device_map: Optional[Dict[int, Union[str, torch.device, int]]] = None, + ) -> None: + super(Cache, self).__init__() + + # TODO (joao): to enable this cache on multiple devicesuse the pattern from `OffloadedCache`, which keeps + # track of the original device of each layer + unique_devices = set(layer_device_map.values()) if layer_device_map else set() + if len(unique_devices) > 1: + raise ValueError(f"OffloadedStaticCache does not support multiple devices. Got devices: {unique_devices}") + + self.max_batch_size = max_batch_size + self.max_cache_len = config.max_position_embeddings if max_cache_len is None else max_cache_len + self.device = torch.device(device) if layer_device_map is None else torch.device(layer_device_map[0]) + self.offload_device = torch.device(offload_device) + self._dtype = dtype if dtype is not None else torch.float32 + + # Some model define a custom `head_dim` != config.hidden_size // config.num_attention_heads + head_dim = config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads + + num_key_value_heads = ( + config.num_attention_heads + if getattr(config, "num_key_value_heads", None) is None + else config.num_key_value_heads + ) + + cache_shape = (max_batch_size, num_key_value_heads, self.max_cache_len, head_dim) + + # Create offloaded CPU tensors. + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + + for i in range(config.num_hidden_layers): + # First layer is always on-device. + device = self.device if i == 0 else self.offload_device + + key_cache, value_cache = self._create_key_value_cache_tensors(cache_shape, device) + + self.key_cache.append(key_cache) + self.value_cache.append(value_cache) + + # Create device tensors. + self._device_key_cache: List[torch.Tensor] = [] + self._device_value_cache: List[torch.Tensor] = [] + + for i in range(2): + key_cache, value_cache = self._create_key_value_cache_tensors(cache_shape, self.device) + + self._device_key_cache.append(key_cache) + self._device_value_cache.append(value_cache) + + # For backwards compatibility. + # TODO(gante): Remove this. + self._seen_tokens = 0 + + # Create new CUDA stream for parallel prefetching. + self._prefetch_stream = torch.cuda.Stream() if self.device.type == "cuda" else None + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + cache_kwargs: Optional[Dict[str, Any]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. + It is VERY important to index using a tensor, otherwise you introduce a copy to the device. + + Parameters: + key_states (`torch.Tensor`): + The new key states to cache. + value_states (`torch.Tensor`): + The new value states to cache. + layer_idx (`int`): + The index of the layer to cache the states for. + cache_kwargs (`Dict[str, Any]`, *optional*): + Additional arguments for the cache subclass. The `OffloadedStaticCache` needs the + `cache_position` input to know how where to write in the cache. + + Return: + A tuple containing the updated key and value states. + """ + + key_states = key_states.to(self.key_cache[layer_idx].dtype) + value_states = value_states.to(self.value_cache[layer_idx].dtype) + + if layer_idx == 0: + # Update seen tokens. + # TODO(gante): Remove this. + self._seen_tokens += key_states.shape[-2] + + # Always there. + k_out = self.key_cache[0] + v_out = self.value_cache[0] + else: + # Wait for prefetch stream. + if self._prefetch_stream is not None: + torch.cuda.default_stream(self.device).wait_stream(self._prefetch_stream) + + k_out = self._device_key_cache[layer_idx & 1] + v_out = self._device_value_cache[layer_idx & 1] + + self._prefetch_layer(layer_idx + 1) + + cache_position = cache_kwargs.get("cache_position") if cache_kwargs is not None else None + if cache_position is None: + k_out.copy_(key_states) + v_out.copy_(value_states) + + # Copy the values to the offloaded device as well. + if layer_idx == 0: + self.key_cache[layer_idx].copy_(key_states.to(self.offload_device)) + self.value_cache[layer_idx].copy_(value_states.to(self.offload_device)) + else: + # Note: here we use `tensor.index_copy_(dim, index, tensor)` that is equivalent to + # `tensor[:, :, index] = tensor`, but the first one is compile-friendly and it does + # explicitly an in-place operation, that avoids copies and uses less memory. + try: + k_out.index_copy_(2, cache_position, key_states) + v_out.index_copy_(2, cache_position, value_states) + except NotImplementedError: + # The operator 'aten::index_copy.out' is not currently implemented for the MPS + # device. + k_out[:, :, cache_position] = key_states + v_out[:, :, cache_position] = value_states + + # Copy the values to the offloaded device as well. + if layer_idx != 0: + cache_position = cache_position.to(self.offload_device) + key_states = key_states.to(self.offload_device) + value_states = value_states.to(self.offload_device) + + try: + self.key_cache[layer_idx].index_copy_(2, cache_position, key_states) + self.value_cache[layer_idx].index_copy_(2, cache_position, value_states) + except NotImplementedError: + # The operator 'aten::index_copy.out' is not currently implemented for the MPS + # device. + self.key_cache[layer_idx][:, :, cache_position] = key_states + self.value_cache[layer_idx][:, :, cache_position] = value_states + + return k_out, v_out + + def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: + """Returns the sequence length of the cached states that were seen by the model.""" + + # TODO(gante): Remove this. + return self._seen_tokens + + def get_max_cache_shape(self) -> Optional[int]: + """Returns the maximum sequence length of the cached states.""" + + return self.max_cache_len + + def reset(self) -> None: + """Resets the cache values while preserving the objects.""" + + # For backwards compatibility. + # TODO(gante): Remove this. + self._seen_tokens = 0 + + # Zero out cache. + for layer_idx in range(len(self.key_cache)): + # In-place ops prevent breaking the static address. + self.key_cache[layer_idx].zero_() + self.value_cache[layer_idx].zero_() + + @property + def seen_tokens(self) -> int: + # For backwards compatibility. + # TODO(gante): Remove this. + return self._seen_tokens + + def _create_key_value_cache_tensors( + self, shape: Tuple[int, ...], device: torch.device + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Creates K/V cache tensors on a device. Pins memory for CPU tensors. Marks them as static + addresses for non-CPU tensors. + + Args: + shape (`Tuple[int, ...]`): Shape. + device (`torch.device`): Device. + + Returns: + Key and value cache tensors as a tuple. + """ + + is_cpu_device = device == torch.device("cpu") + + key_cache = torch.zeros(shape, dtype=self._dtype, device=device, pin_memory=is_cpu_device) + value_cache = torch.zeros(shape, dtype=self._dtype, device=device, pin_memory=is_cpu_device) + + # Note: `mark_static_address` is used to tag the cache as a fixed data pointer, + # preventing compiled graph breaks when updating the cache. + torch._dynamo.mark_static_address(key_cache) + torch._dynamo.mark_static_address(value_cache) + + return key_cache, value_cache + + def _prefetch_layer(self, layer_idx: int) -> None: + """Prefetch a layer to the device. Needs to be called in order of layer indices.""" + + # Don't fetch layers that do not exist. + if layer_idx >= len(self.key_cache): + return + + # Alternate between two on-device caches. + if self._prefetch_stream is not None: + with torch.cuda.stream(self._prefetch_stream): + self._prefetch_layer_in_context(layer_idx) + else: + self._prefetch_layer_in_context(layer_idx) + + def _prefetch_layer_in_context(self, layer_idx: int) -> None: + """Performs the actual copy of the layer to device cache.""" + + self._device_key_cache[layer_idx & 1].copy_(self.key_cache[layer_idx], non_blocking=True) + self._device_value_cache[layer_idx & 1].copy_(self.value_cache[layer_idx], non_blocking=True) diff --git a/csrc/bit_decode/CMakeLists.txt b/csrc/bit_decode/CMakeLists.txt index ae2e5d1..dea0cbe 100644 --- a/csrc/bit_decode/CMakeLists.txt +++ b/csrc/bit_decode/CMakeLists.txt @@ -6,7 +6,8 @@ set(CMAKE_CXX_STANDARD 17) set(CMAKE_CUDA_STANDARD 17) set(CMAKE_CUDA_ARCHITECTURES 80) -set(INCLUDE_DIR ${PROJECT_SOURCE_DIR}/../../libs/cutlass/include) +# set(INCLUDE_DIR ${PROJECT_SOURCE_DIR}/../../libs/cutlass/include) +set(INCLUDE_DIR /home/ddy/Projects/BitDecoding/libs/cutlass) # Enable ccache if available find_program(CCACHE_PROGRAM ccache) @@ -18,42 +19,68 @@ endif() find_package(Torch REQUIRED) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}") -message(STATUS "Compile testing packdecode kernel.") -add_executable(test_single_packdecode - ${PROJECT_SOURCE_DIR}/src/test_single_packdecode.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_hdim128_fp16_sm80.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_2bit.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_2bit.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu -) -target_link_libraries(test_single_packdecode "${TORCH_LIBRARIES}") -target_include_directories(test_single_packdecode PRIVATE ${INCLUDE_DIR}) -target_compile_options(test_single_packdecode PRIVATE $<$:-maxrregcount=255 -gencode arch=compute_80,code=sm_80 -w>) +# message(STATUS "Compile testing packdecode kernel.") +# add_executable(test_single_packdecode +# ${PROJECT_SOURCE_DIR}/src/test_single_packdecode.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_hdim128_fp16_sm80.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_2bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_2bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu +# ) +# target_link_libraries(test_single_packdecode "${TORCH_LIBRARIES}") +# target_include_directories(test_single_packdecode PRIVATE ${INCLUDE_DIR}) +# target_compile_options(test_single_packdecode PRIVATE $<$:-maxrregcount=255 -gencode arch=compute_80,code=sm_80 -w>) -message(STATUS "Compile testing packdecode kernel.") -add_executable(test_batch_packdecode - ${PROJECT_SOURCE_DIR}/src/test_batch_packdecode.cu +message(STATUS "Compile testing residual kernel.") +add_executable(test_single_residual + ${PROJECT_SOURCE_DIR}/src/test_single_residual.cu ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_hdim128_fp16_sm80.cu ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_2bit.cu ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_2bit.cu ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu ) -target_link_libraries(test_batch_packdecode "${TORCH_LIBRARIES}") -target_include_directories(test_batch_packdecode PRIVATE ${INCLUDE_DIR}) -target_compile_options(test_batch_packdecode PRIVATE $<$:-maxrregcount=255 -gencode arch=compute_80,code=sm_80 -w>) +target_link_libraries(test_single_residual "${TORCH_LIBRARIES}") +target_include_directories(test_single_residual PRIVATE ${INCLUDE_DIR}) +target_compile_options(test_single_residual PRIVATE $<$:-maxrregcount=255 -gencode arch=compute_80,code=sm_80 -w>) -message(STATUS "Compile benchmarking kernel.") -add_executable(bench_single_packdecode - ${PROJECT_SOURCE_DIR}/src/bench_single_packdecode.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_hdim128_fp16_sm80.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_2bit.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_2bit.cu - ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu -) -target_link_libraries(bench_single_packdecode "${TORCH_LIBRARIES}") -target_include_directories(bench_single_packdecode PRIVATE ${INCLUDE_DIR}) -target_compile_options(bench_single_packdecode PRIVATE $<$:-maxrregcount=255 -gencode arch=compute_80,code=sm_80 -w>) +# message(STATUS "Compile testing packdecode kernel.") +# add_executable(test_batch_packdecode +# ${PROJECT_SOURCE_DIR}/src/test_batch_packdecode.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_hdim128_fp16_sm80.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_2bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_2bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu +# ) +# target_link_libraries(test_batch_packdecode "${TORCH_LIBRARIES}") +# target_include_directories(test_batch_packdecode PRIVATE ${INCLUDE_DIR}) +# target_compile_options(test_batch_packdecode PRIVATE $<$:-maxrregcount=255 -gencode arch=compute_80,code=sm_80 -w>) + +# message(STATUS "Compile benchmarking kernel.") +# add_executable(bench_single_packdecode +# ${PROJECT_SOURCE_DIR}/src/bench_single_packdecode.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_hdim128_fp16_sm80.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_2bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_2bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu +# ) +# target_link_libraries(bench_single_packdecode "${TORCH_LIBRARIES}") +# target_include_directories(bench_single_packdecode PRIVATE ${INCLUDE_DIR}) +# target_compile_options(bench_single_packdecode PRIVATE $<$:-maxrregcount=255 -gencode arch=compute_80,code=sm_80 -w>) + +# message(STATUS "Compile benchmarking kernel.") +# add_executable(bench_single_residual +# ${PROJECT_SOURCE_DIR}/src/bench_single_residual.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_hdim128_fp16_sm80.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_2bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_2bit.cu +# ${PROJECT_SOURCE_DIR}/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu +# ) +# target_link_libraries(bench_single_residual "${TORCH_LIBRARIES}") +# target_include_directories(bench_single_residual PRIVATE ${INCLUDE_DIR}) +# target_compile_options(bench_single_residual PRIVATE $<$:-maxrregcount=255 -gencode arch=compute_80,code=sm_80 -w>) diff --git a/csrc/bit_decode/decode_api.cpp b/csrc/bit_decode/decode_api.cpp index 82bd49f..0a9c388 100644 --- a/csrc/bit_decode/decode_api.cpp +++ b/csrc/bit_decode/decode_api.cpp @@ -8,6 +8,7 @@ #include #include #include + #include #include @@ -20,35 +21,35 @@ #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") void set_params_fprop(Flash_fwd_params ¶ms, - // sizes - const size_t b, - const size_t seqlen_q, - const size_t seqlen_k, - const size_t seqlen_q_rounded, - const size_t seqlen_k_rounded, - const size_t h, - const size_t h_k, - const size_t d, - const size_t d_rounded, - // device pointers - const at::Tensor q, - const at::Tensor k, const at::Tensor k_pack, const at::Tensor k_params, - const at::Tensor v, const at::Tensor v_pack, const at::Tensor v_params, - at::Tensor out, - void *cu_seqlens_q_d, - void *cu_seqlens_k_d, - void *seqused_k, - void *p_d, - void *softmax_lse_d, - float p_dropout, - float softmax_scale, - int window_size_left, - int window_size_right, - const float softcap, - const std::string quant_mode, - const int group_size, - bool seqlenq_ngroups_swapped=false, - const bool unpadded_lse=false) { + // sizes + const size_t b, + const size_t seqlen_q, + const size_t seqlen_k, + const size_t seqlen_q_rounded, + const size_t seqlen_k_rounded, + const size_t h, + const size_t h_k, + const size_t d, + const size_t d_rounded, + // device pointers + const at::Tensor q, + const at::Tensor k, const at::Tensor k_pack, const at::Tensor k_params, + const at::Tensor v, const at::Tensor v_pack, const at::Tensor v_params, + at::Tensor out, + void *cu_seqlens_q_d, + void *cu_seqlens_k_d, + void *seqused_k, + void *p_d, + void *softmax_lse_d, + float p_dropout, + float softmax_scale, + int window_size_left, + int window_size_right, + const float softcap, + const std::string quant_mode, + const int group_size, + bool seqlenq_ngroups_swapped=false, + const bool unpadded_lse=false) { // Reset the parameters params = {}; @@ -89,19 +90,19 @@ void set_params_fprop(Flash_fwd_params ¶ms, params.o_head_stride = out.stride(-2); if (cu_seqlens_q_d == nullptr) { - params.q_batch_stride = q.stride(0); - // params.k_batch_stride = k.stride(0); - params.K_pack_batch_stride = k_pack.stride(0); - params.k_params_batch_stride = k_params.stride(0); - // params.v_batch_stride = v.stride(0); - params.v_pack_batch_stride = v_pack.stride(0); - params.v_params_batch_stride = v_params.stride(0); - params.o_batch_stride = out.stride(0); - - if (seqlenq_ngroups_swapped) { - params.q_batch_stride *= seqlen_q; - params.o_batch_stride *= seqlen_q; - } + params.q_batch_stride = q.stride(0); + // params.k_batch_stride = k.stride(0); + params.K_pack_batch_stride = k_pack.stride(0); + params.k_params_batch_stride = k_params.stride(0); + // params.v_batch_stride = v.stride(0); + params.v_pack_batch_stride = v_pack.stride(0); + params.v_params_batch_stride = v_params.stride(0); + params.o_batch_stride = out.stride(0); + + if (seqlenq_ngroups_swapped) { + params.q_batch_stride *= seqlen_q; + params.o_batch_stride *= seqlen_q; + } } params.cu_seqlens_q = static_cast(cu_seqlens_q_d); @@ -131,14 +132,14 @@ void set_params_fprop(Flash_fwd_params ¶ms, TORCH_CHECK(softcap <= 0.0, "This flash attention build does not support softcap."); #endif if (softcap > 0.0) { - params.softcap = softmax_scale / softcap; - params.scale_softmax = softcap; - params.scale_softmax_log2 = softcap * M_LOG2E; + params.softcap = softmax_scale / softcap; + params.scale_softmax = softcap; + params.scale_softmax_log2 = softcap * M_LOG2E; } else{ - // Remove potential NaN - params.softcap = 0.0; - params.scale_softmax = softmax_scale; - params.scale_softmax_log2 = softmax_scale * M_LOG2E; + // Remove potential NaN + params.softcap = 0.0; + params.scale_softmax = softmax_scale; + params.scale_softmax_log2 = softmax_scale * M_LOG2E; } // Set this to probability of keeping an element to simplify things. @@ -181,6 +182,17 @@ void set_params_fprop(Flash_fwd_params ¶ms, template void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream, bool force_split_kernel=false) { + // FP16_SWITCH(!params.is_bf16, [&] { + // HEADDIM_SWITCH(params.d, [&] { + // BOOL_SWITCH(params.is_causal, Is_causal, [&] { + // if (params.num_splits <= 1 && !force_split_kernel) { // If we don't set it num_splits == 0 + // run_mha_fwd_(params, stream); + // } else { + // run_mha_fwd_splitkv_dispatch(params, stream); + // } + // }); + // }); + // }); if (params.num_splits <= 1 && !force_split_kernel) { // If we don't set it num_splits == 0 run_mha_fwd_(params, stream); } else { @@ -190,7 +202,7 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream, bool force_split } else if (params.group_size == 64) { // run_mha_fwd_splitkv_dispatch(params, stream); } else if (params.group_size == 32) { - // run_mha_fwd_splitkv_dispatch(params, stream); + run_mha_fwd_splitkv_dispatch(params, stream); } } else { if (params.group_size == 128) { @@ -207,20 +219,20 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream, bool force_split template void run_kvcache_qpack(Flash_fwd_params ¶ms, cudaStream_t stream) { if (params.quant_mode == "k-channel") { - if (params.group_size == 128) { - run_kvcache_qpack_(params, stream); + if (params.group_size == 32) { + run_kvcache_qpack_(params, stream); } else if (params.group_size == 64) { // run_kvcache_qpack_(params, stream); - } else if (params.group_size == 32) { - // run_kvcache_qpack_(params, stream); + } else if (params.group_size == 128) { + run_kvcache_qpack_(params, stream); } } else { - if (params.group_size == 128) { - // run_kvcache_qpack_(params, stream); + if (params.group_size == 32) { + // run_kvcache_qpack_(params, stream); } else if (params.group_size == 64) { // run_kvcache_qpack_(params, stream); - } else if (params.group_size == 32) { - // run_kvcache_qpack_(params, stream); + } else if (params.group_size == 128) { + // run_kvcache_qpack_(params, stream); } } } @@ -274,8 +286,8 @@ void set_params_splitkv(Flash_fwd_params ¶ms, const int batch_size, // This needs to match with run_mha_fwd_splitkv_dispatch // TODO - const int block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); - // const int block_n = 256; + // const int block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); + const int block_n = 256; const int num_n_blocks = (max_seqlen_k + block_n - 1) / block_n; // Technically kBlockM = 64 only for the splitKV kernels, not the standard kernel. // In any case we don't expect seqlen_q to be larger than 64 for inference. @@ -288,6 +300,8 @@ void set_params_splitkv(Flash_fwd_params ¶ms, const int batch_size, // printf("num_splits = %d\n", params.num_splits); // params.num_splits= 1; } + params.num_splits += 1; // We need to add 1 for residual kernel. + // printf("num_splits = %d\n", params.num_splits); if (params.num_splits > 1) { at::Tensor softmax_lse_accum = torch::empty({params.num_splits, batch_size, num_heads, max_seqlen_q}, opts.dtype(at::kFloat)); at::Tensor out_accum = torch::empty({params.num_splits, batch_size, num_heads, max_seqlen_q, head_size_rounded}, opts.dtype(at::kFloat)); @@ -299,16 +313,25 @@ void set_params_splitkv(Flash_fwd_params ¶ms, const int batch_size, } template -at::Tensor +std::tuple mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size const at::Tensor &k_pack, // batch_size_c x seqlen_k / 4 x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. const at::Tensor &k_params, // batch_size_c x num_groups x num_heads_k x head_size const at::Tensor &v_pack, // batch_size_c x seqlen_k / 4 x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. const at::Tensor &v_params, // batch_size_c x num_groups x num_heads_k x head_size + c10::optional &k_, // batch_size x seqlen_knew x num_heads_k x head_size + c10::optional &v_, // batch_size x seqlen_knew x num_heads_k x head_size + c10::optional &seqlens_k_, // batch_size + at::Tensor &k_pack_new, + at::Tensor &k_params_new, + at::Tensor &v_pack_new, + at::Tensor &v_params_new, c10::optional &block_table_, // batch_size x max_num_blocks_per_seq const float softmax_scale=1.0, const std::string quant_mode="k-tensor", const int group_size=128, + const int residual_block_size=128, + const int new_lens=0, bool is_causal=false, int window_size_left=-1, int window_size_right=-1, @@ -323,7 +346,7 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x bool is_sm90 = dprops->major == 9 && dprops->minor == 0; TORCH_CHECK(is_sm90 || is_sm8x, "FlashAttention only supports Ampere GPUs or newer."); - CHECK_DEVICE(q); CHECK_DEVICE(k_pack); CHECK_DEVICE(v_pack); + CHECK_DEVICE(q); // CHECK_DEVICE(kcache); CHECK_DEVICE(vcache); at::Tensor block_table; const bool paged_KV = block_table_.has_value(); @@ -414,7 +437,69 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x group_size ); - params.is_seqlens_k_cumulative = true; + at::Tensor k, v; + if (k_.has_value()) { + TORCH_CHECK(v_.has_value(), "If key is supplied, value must also be passed in"); + TORCH_CHECK(seqlens_k_.has_value(), "If key is supplied, seqlens_k must also be passed in"); + + k = k_.value(); + v = v_.value(); + int seqlen_knew = k.size(1); + auto seqlens_k = seqlens_k_.value(); + + TORCH_CHECK(k.dtype() == q_dtype, "Key must have the same dtype as query"); + TORCH_CHECK(v.dtype() == q_dtype, "Value must have the same dtype as query"); + TORCH_CHECK(k.stride(-1) == 1, "Key tensor must have contiguous last dimension"); + TORCH_CHECK(v.stride(-1) == 1, "Value tensor must have contiguous last dimension"); + CHECK_SHAPE(k, batch_size, seqlen_knew, num_heads_k, head_size_og); + CHECK_SHAPE(v, batch_size, seqlen_knew, num_heads_k, head_size_og); + CHECK_DEVICE(k); CHECK_DEVICE(v); + TORCH_CHECK(seqlens_k.dtype() == torch::kInt32, "seqlens_k must have dtype int32"); + CHECK_DEVICE(seqlens_k); + CHECK_CONTIGUOUS(seqlens_k); + CHECK_SHAPE(seqlens_k, batch_size); + + params.new_lens = new_lens; + + params.seqlen_knew = seqlen_knew; + params.knew_ptr = k.data_ptr(); + params.vnew_ptr = v.data_ptr(); + params.knew_batch_stride = k.stride(0); + params.vnew_batch_stride = v.stride(0); + params.knew_row_stride = k.stride(-3); + params.vnew_row_stride = v.stride(-3); + params.knew_head_stride = k.stride(-2); + params.vnew_head_stride = v.stride(-2); + params.cu_seqlens_k = static_cast(seqlens_k.data_ptr()); + + const int pack_nums = 16 / num_bits; + + params.k_pack_new_ptr = k_pack_new.data_ptr(); + params.k_params_new_ptr = k_params_new.data_ptr(); + params.v_pack_new_ptr = v_pack_new.data_ptr(); + params.v_params_new_ptr = v_params_new.data_ptr(); + + params.k_pack_new_batch_stride = k_pack_new.stride(0); + params.k_params_new_batch_stride = k_params_new.stride(0); + params.v_pack_new_batch_stride = v_pack_new.stride(0); + params.v_params_new_batch_stride = v_params_new.stride(0); + + params.k_pack_new_row_stride = k_pack_new.stride(-3); + params.k_params_new_row_stride = k_params_new.stride(-3); + params.v_pack_new_row_stride = v_pack_new.stride(-3); + params.v_params_new_row_stride = v_params_new.stride(-1); + + params.k_pack_new_head_stride = k_pack_new.stride(-2); + params.k_params_new_head_stride = k_params_new.stride(-2); + params.v_pack_new_head_stride = v_pack_new.stride(-2); + params.v_params_new_head_stride = v_params_new.stride(-2); + + params.k_params_new_dim_stride = k_params_new.stride(-1); + params.v_params_new_dim_stride = v_params_new.stride(-3); + + } + + params.is_seqlens_k_cumulative = !(seqlens_k_.has_value()); params.rotary_dim = 0; if (paged_KV) { @@ -437,7 +522,7 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x out = out.transpose(1, 2).reshape({batch_size, 1, num_heads_k * seqlen_q, head_size_og}); } - return out; + return std::make_tuple(out, k_pack_new, k_params_new, v_pack_new, v_params_new); } @@ -457,7 +542,8 @@ void set_params_fprop_qpack(Flash_fwd_params ¶ms, const at::Tensor v, at::Tensor v_pack, at::Tensor v_params, void *cu_seqlens_k_d, const std::string quant_mode, - const int group_size + const int group_size, + bool page_kv ) { // Reset the parameters @@ -490,12 +576,12 @@ void set_params_fprop_qpack(Flash_fwd_params ¶ms, params.v_pack_head_stride = v_pack.stride(-2); params.v_params_head_stride = v_params.stride(-2); - // params.k_batch_stride = k.stride(0); - params.k_batch_stride = seqlen_k * k.size(-2) * k.size(-1); + if (page_kv) params.k_batch_stride = k.stride(0); + else params.k_batch_stride = seqlen_k * k.size(-2) * k.size(-1); params.K_pack_batch_stride = k_pack.stride(0); params.k_params_batch_stride = k_params.stride(0); - // params.v_batch_stride = v.stride(0); - params.v_batch_stride = seqlen_k * v.size(-2) * v.size(-1); + if (page_kv) params.v_batch_stride = v.stride(0); + else params.v_batch_stride = seqlen_k * v.size(-2) * v.size(-1); params.v_pack_batch_stride = v_pack.stride(0); params.v_params_batch_stride = v_params.stride(0); @@ -513,20 +599,15 @@ void set_params_fprop_qpack(Flash_fwd_params ¶ms, params.group_size = group_size; } - template -void kvcache_qpack(const at::Tensor &k, - at::Tensor &k_pack, - at::Tensor &k_params, - const at::Tensor &v, - at::Tensor &v_pack, - at::Tensor &v_params, - c10::optional &block_table_, - const at::Tensor &cu_seqlens_k, - const int max_seqlen_k, - const std::string quant_mode, - const int group_size - ) { +void kvcache_qpack(const at::Tensor &k, at::Tensor &k_pack, at::Tensor &k_params, + const at::Tensor &v, at::Tensor &v_pack, at::Tensor &v_params, + c10::optional &block_table_, + const at::Tensor &cu_seqlens_k, + const int max_seqlen_k, + const std::string quant_mode, + const int group_size + ) { auto k_dtype = k.dtype(); TORCH_CHECK(k_dtype == torch::kFloat16 || k_dtype == torch::kBFloat16, @@ -584,7 +665,8 @@ void kvcache_qpack(const at::Tensor &k, v, v_pack, v_params, /*cu_seqlens_k_d=*/nullptr, quant_mode, - group_size + group_size, + paged_KV ); if (paged_KV) { @@ -602,10 +684,11 @@ void kvcache_qpack(const at::Tensor &k, return; } + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "BitDecoding"; - m.def("kvcache_pack_i2", &kvcache_qpack<2>, "Forward pass, kvcache quantization and packing (2-bit)"); - m.def("kvcache_pack_i4", &kvcache_qpack<4>, "Forward pass, kvcache quantization and packing (4-bit)"); - m.def("fwd_kvcache_i2", &mha_fwd_kvcache<2>, "Forward pass, with 2-bit KV-cache"); - m.def("fwd_kvcache_i4", &mha_fwd_kvcache<4>, "Forward pass, with 4-bit KV-cache"); + // m.def("kvcache_pack_i2", &kvcache_qpack<2>, "Forward pass, kvcache quantization and packing (2-bit)"); + m.def("kvcache_pack_int4", &kvcache_qpack<4>, "Forward pass, kvcache quantization and packing (4-bit)"); + // m.def("fwd_kvcache_i2", &mha_fwd_kvcache<2>, "Forward pass, with 2-bit KV-cache"); + m.def("fwd_kvcache_int4", &mha_fwd_kvcache<4>, "Forward pass, with 4-bit KV-cache"); } diff --git a/csrc/bit_decode/src/bench_single_residual.cu b/csrc/bit_decode/src/bench_single_residual.cu new file mode 100644 index 0000000..80abaeb --- /dev/null +++ b/csrc/bit_decode/src/bench_single_residual.cu @@ -0,0 +1,167 @@ +#include "flash_api.h" +#include +#include + + +template +double TestDecodingKernelPerformance(int seqlen_kv, const std::string& quant_mode, const int group_size, const int repeat) { + const int bs = 1; + const int seqlen_q = 1; + const int pack_nums = 16 / num_bits; + const int residual_block_size = num_bits == 4 ? 128 : 256; + int residual_len = seqlen_kv % residual_block_size == 0 ? residual_block_size : seqlen_kv % residual_block_size; + seqlen_kv = seqlen_kv - residual_len; + + bool residual = residual_len > 0 ? true : false; + + torch::Tensor Q_host = torch::rand({bs, seqlen_q, num_heads, head_dim}, torch::dtype(torch::kHalf)); + torch::Tensor K_host = torch::ones({bs, seqlen_kv, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + torch::Tensor V_host = torch::ones({bs, seqlen_kv, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + + torch::Tensor Q_device = Q_host.to(torch::kCUDA); + torch::Tensor K_device = K_host.to(torch::kCUDA); + torch::Tensor V_device = V_host.to(torch::kCUDA); + + at::Tensor k_pack, k_params, v_pack, v_params, k_pack_new, k_params_new, v_pack_new, v_params_new; + if (quant_mode == "k-channel") { + k_pack = torch::empty({bs, seqlen_kv / pack_nums, num_heads_kv, head_dim}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + k_params = torch::empty({bs, seqlen_kv / group_size, num_heads_kv, head_dim}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + k_pack_new = torch::empty({bs, residual_block_size / pack_nums, num_heads_kv, k_pack.size(-1)}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + k_params_new = torch::empty({bs, residual_block_size / group_size, num_heads_kv, k_params.size(-1)}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + } else { + k_pack = torch::empty({bs, seqlen_kv, num_heads_kv, head_dim / pack_nums}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + k_params = torch::empty({bs, head_dim / group_size, num_heads_kv, seqlen_kv}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + k_pack_new = torch::empty({bs, residual_block_size, num_heads_kv, k_pack.size(-1)}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + k_params_new = torch::empty({bs, k_params.size(1), num_heads_kv, residual_block_size}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + } + + v_pack = torch::empty({bs, seqlen_kv, num_heads_kv, head_dim / pack_nums}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + v_params = torch::empty({bs, head_dim / group_size, num_heads_kv, seqlen_kv}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + v_pack_new = torch::empty({bs, residual_block_size, num_heads_kv, v_pack.size(-1)}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + v_params_new = torch::empty({bs, v_params.size(1), num_heads_kv, residual_block_size}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + + // Convert K, V to unpadded format + torch::Tensor K_unpad = K_device.reshape({bs * seqlen_kv, num_heads_kv, head_dim}); + torch::Tensor V_unpad = V_device.reshape({bs * seqlen_kv, num_heads_kv, head_dim}); + + auto cu_seqlens_k = torch::arange(0, (bs + 1) * seqlen_kv, seqlen_kv, torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA)); + std::optional opt_block_table = std::nullopt; + + kvcache_qpack( + K_unpad, k_pack, k_params, + V_unpad, v_pack, v_params, + opt_block_table, + cu_seqlens_k, + seqlen_kv, + quant_mode, + group_size + ); + + at::Tensor K_residual_host, V_residual_host, K_new_host, V_new_host, K_new_device, V_new_device, seqlens_k; + int new_lens = 0; + + if (residual) { + seqlens_k = torch::full({bs}, seqlen_kv, torch::dtype(torch::kInt32).device(torch::kCUDA)); + + K_residual_host = torch::zeros({bs, residual_block_size, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + V_residual_host = torch::zeros({bs, residual_block_size, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + + K_new_host = torch::randn({bs, residual_len, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + V_new_host = torch::randn({bs, residual_len, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + + new_lens = residual_len; + + // Copy data from K_new_host to K_residual_host + K_residual_host.slice(1, 0, residual_len).copy_(K_new_host); + V_residual_host.slice(1, 0, residual_len).copy_(V_new_host); + + K_new_device = K_residual_host.to(torch::kCUDA); + V_new_device = V_residual_host.to(torch::kCUDA); + } + + const float sm_scale = 1 / std::sqrt(float(head_dim)); + std::optional opt_K_new_device = residual ? std::make_optional(K_new_device) : std::nullopt; + std::optional opt_V_new_device = residual ? std::make_optional(V_new_device) : std::nullopt; + std::optional opt_seqlens_k = std::make_optional(seqlens_k); + + // Warm up + for (int i = 0; i < 10; ++i) + mha_fwd_kvcache(Q_device, + k_pack, k_params, + v_pack, v_params, + opt_K_new_device, opt_V_new_device, opt_seqlens_k, + k_pack_new, k_params_new, v_pack_new, v_params_new, + opt_block_table, + sm_scale, + quant_mode, + group_size, + residual_block_size, + new_lens); + + // Benchmark + cudaEvent_t start, end; + cudaEventCreate(&start); + cudaEventCreate(&end); + cudaEventRecord(start); + for (int i = 0; i < repeat; i++) { + mha_fwd_kvcache(Q_device, + k_pack, k_params, + v_pack, v_params, + opt_K_new_device, opt_V_new_device, opt_seqlens_k, + k_pack_new, k_params_new, v_pack_new, v_params_new, + opt_block_table, + sm_scale, + quant_mode, + group_size, + residual_block_size, + new_lens); + } + cudaEventRecord(end); + cudaEventSynchronize(end); + + float msec, sec; + cudaEventElapsedTime(&msec, start, end); + msec = msec / repeat; + + return msec; +} + +int main() { + const int num_heads = 32; + const int num_heads_kv = 32; + const int head_dim = 128; + + const std::string quant_mode = "k-channel"; + const int num_bits = 4; + const int group_size = 128; + + const int test_num = 10; + int len_list[test_num]; + len_list[0] = 1024; + for (int i = 1; i < test_num; i++) { + len_list[i] = len_list[i - 1] * 2; + } + + const int outer_repeat = 3, inner_repeat = 3; + printf("\n######## Benchmark single decode ########\n"); + for (int j = 0; j < test_num; j++) { + + int seqlen_kv = len_list[j] + 1; + double max_msec = 0.0; + double min_msec = DBL_MAX; + double total_msec = 0.0; + + for (int k = 0; k < outer_repeat; k++) { + double this_sec = TestDecodingKernelPerformance(seqlen_kv, quant_mode, group_size, inner_repeat); + max_msec = max(max_msec, this_sec); + min_msec = min(min_msec, this_sec); + total_msec += this_sec; + } + + double avg_msec = total_msec / outer_repeat; + printf("seqlen_kv num_heads head_dim = %6d %6d %6d, ", seqlen_kv, num_heads, head_dim); + printf("Time = %12.8lf %12.8lf %12.8lf ms, \n", min_msec, avg_msec, max_msec); + } + + return 0; +} \ No newline at end of file diff --git a/csrc/bit_decode/src/flash_api.h b/csrc/bit_decode/src/flash_api.h index b12d600..8a0f308 100644 --- a/csrc/bit_decode/src/flash_api.h +++ b/csrc/bit_decode/src/flash_api.h @@ -19,35 +19,35 @@ #define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous") void set_params_fprop(Flash_fwd_params ¶ms, - // sizes - const size_t b, - const size_t seqlen_q, - const size_t seqlen_k, - const size_t seqlen_q_rounded, - const size_t seqlen_k_rounded, - const size_t h, - const size_t h_k, - const size_t d, - const size_t d_rounded, - // device pointers - const at::Tensor q, - const at::Tensor k, const at::Tensor k_pack, const at::Tensor k_params, - const at::Tensor v, const at::Tensor v_pack, const at::Tensor v_params, - at::Tensor out, - void *cu_seqlens_q_d, - void *cu_seqlens_k_d, - void *seqused_k, - void *p_d, - void *softmax_lse_d, - float p_dropout, - float softmax_scale, - int window_size_left, - int window_size_right, - const float softcap, - const std::string quant_mode, - const int group_size, - bool seqlenq_ngroups_swapped=false, - const bool unpadded_lse=false) { + // sizes + const size_t b, + const size_t seqlen_q, + const size_t seqlen_k, + const size_t seqlen_q_rounded, + const size_t seqlen_k_rounded, + const size_t h, + const size_t h_k, + const size_t d, + const size_t d_rounded, + // device pointers + const at::Tensor q, + const at::Tensor k, const at::Tensor k_pack, const at::Tensor k_params, + const at::Tensor v, const at::Tensor v_pack, const at::Tensor v_params, + at::Tensor out, + void *cu_seqlens_q_d, + void *cu_seqlens_k_d, + void *seqused_k, + void *p_d, + void *softmax_lse_d, + float p_dropout, + float softmax_scale, + int window_size_left, + int window_size_right, + const float softcap, + const std::string quant_mode, + const int group_size, + bool seqlenq_ngroups_swapped=false, + const bool unpadded_lse=false) { // Reset the parameters params = {}; @@ -180,6 +180,17 @@ void set_params_fprop(Flash_fwd_params ¶ms, template void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream, bool force_split_kernel=false) { + // FP16_SWITCH(!params.is_bf16, [&] { + // HEADDIM_SWITCH(params.d, [&] { + // BOOL_SWITCH(params.is_causal, Is_causal, [&] { + // if (params.num_splits <= 1 && !force_split_kernel) { // If we don't set it num_splits == 0 + // run_mha_fwd_(params, stream); + // } else { + // run_mha_fwd_splitkv_dispatch(params, stream); + // } + // }); + // }); + // }); if (params.num_splits <= 1 && !force_split_kernel) { // If we don't set it num_splits == 0 run_mha_fwd_(params, stream); } else { @@ -189,7 +200,7 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream, bool force_split } else if (params.group_size == 64) { // run_mha_fwd_splitkv_dispatch(params, stream); } else if (params.group_size == 32) { - // run_mha_fwd_splitkv_dispatch(params, stream); + run_mha_fwd_splitkv_dispatch(params, stream); } } else { if (params.group_size == 128) { @@ -206,20 +217,20 @@ void run_mha_fwd(Flash_fwd_params ¶ms, cudaStream_t stream, bool force_split template void run_kvcache_qpack(Flash_fwd_params ¶ms, cudaStream_t stream) { if (params.quant_mode == "k-channel") { - if (params.group_size == 128) { - run_kvcache_qpack_(params, stream); + if (params.group_size == 32) { + run_kvcache_qpack_(params, stream); } else if (params.group_size == 64) { // run_kvcache_qpack_(params, stream); - } else if (params.group_size == 32) { - // run_kvcache_qpack_(params, stream); + } else if (params.group_size == 128) { + run_kvcache_qpack_(params, stream); } } else { - if (params.group_size == 128) { - // run_kvcache_qpack_(params, stream); + if (params.group_size == 32) { + // run_kvcache_qpack_(params, stream); } else if (params.group_size == 64) { // run_kvcache_qpack_(params, stream); - } else if (params.group_size == 32) { - // run_kvcache_qpack_(params, stream); + } else if (params.group_size == 128) { + // run_kvcache_qpack_(params, stream); } } } @@ -273,8 +284,8 @@ void set_params_splitkv(Flash_fwd_params ¶ms, const int batch_size, // This needs to match with run_mha_fwd_splitkv_dispatch // TODO - const int block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); - // const int block_n = 256; + // const int block_n = head_size <= 64 ? 256 : (head_size <= 128 ? 128 : 64); + const int block_n = 256; const int num_n_blocks = (max_seqlen_k + block_n - 1) / block_n; // Technically kBlockM = 64 only for the splitKV kernels, not the standard kernel. // In any case we don't expect seqlen_q to be larger than 64 for inference. @@ -287,6 +298,8 @@ void set_params_splitkv(Flash_fwd_params ¶ms, const int batch_size, // printf("num_splits = %d\n", params.num_splits); // params.num_splits= 1; } + params.num_splits += 1; // We need to add 1 for residual kernel. + // printf("num_splits = %d\n", params.num_splits); if (params.num_splits > 1) { at::Tensor softmax_lse_accum = torch::empty({params.num_splits, batch_size, num_heads, max_seqlen_q}, opts.dtype(at::kFloat)); at::Tensor out_accum = torch::empty({params.num_splits, batch_size, num_heads, max_seqlen_q, head_size_rounded}, opts.dtype(at::kFloat)); @@ -298,16 +311,27 @@ void set_params_splitkv(Flash_fwd_params ¶ms, const int batch_size, } template -at::Tensor +std::tuple mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size + // const at::Tensor &kcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. const at::Tensor &k_pack, // batch_size_c x seqlen_k / 4 x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. const at::Tensor &k_params, // batch_size_c x num_groups x num_heads_k x head_size + // const at::Tensor &vcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. const at::Tensor &v_pack, // batch_size_c x seqlen_k / 4 x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table. const at::Tensor &v_params, // batch_size_c x num_groups x num_heads_k x head_size + c10::optional &k_, // batch_size x seqlen_knew x num_heads_k x head_size + c10::optional &v_, // batch_size x seqlen_knew x num_heads_k x head_size + c10::optional &seqlens_k_, // batch_size + at::Tensor &k_pack_new, + at::Tensor &k_params_new, + at::Tensor &v_pack_new, + at::Tensor &v_params_new, c10::optional &block_table_, // batch_size x max_num_blocks_per_seq const float softmax_scale=1.0, const std::string quant_mode="k-tensor", const int group_size=128, + const int residual_block_size=128, + const int new_lens=0, bool is_causal=false, int window_size_left=-1, int window_size_right=-1, @@ -322,7 +346,7 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x bool is_sm90 = dprops->major == 9 && dprops->minor == 0; TORCH_CHECK(is_sm90 || is_sm8x, "FlashAttention only supports Ampere GPUs or newer."); - CHECK_DEVICE(q); CHECK_DEVICE(k_pack); CHECK_DEVICE(v_pack); + CHECK_DEVICE(q); // CHECK_DEVICE(kcache); CHECK_DEVICE(vcache); at::Tensor block_table; const bool paged_KV = block_table_.has_value(); @@ -337,7 +361,7 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x const auto sizes = q.sizes(); const int batch_size = sizes[0]; - int seqlen_q = sizes[1]; + int seqlen_q = sizes[1]; int num_heads = sizes[2]; const int head_size_og = sizes[3]; // dim @@ -413,7 +437,69 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x group_size ); - params.is_seqlens_k_cumulative = true; + at::Tensor k, v; + if (k_.has_value()) { + TORCH_CHECK(v_.has_value(), "If key is supplied, value must also be passed in"); + TORCH_CHECK(seqlens_k_.has_value(), "If key is supplied, seqlens_k must also be passed in"); + + k = k_.value(); + v = v_.value(); + int seqlen_knew = k.size(1); + auto seqlens_k = seqlens_k_.value(); + + TORCH_CHECK(k.dtype() == q_dtype, "Key must have the same dtype as query"); + TORCH_CHECK(v.dtype() == q_dtype, "Value must have the same dtype as query"); + TORCH_CHECK(k.stride(-1) == 1, "Key tensor must have contiguous last dimension"); + TORCH_CHECK(v.stride(-1) == 1, "Value tensor must have contiguous last dimension"); + CHECK_SHAPE(k, batch_size, seqlen_knew, num_heads_k, head_size_og); + CHECK_SHAPE(v, batch_size, seqlen_knew, num_heads_k, head_size_og); + CHECK_DEVICE(k); CHECK_DEVICE(v); + TORCH_CHECK(seqlens_k.dtype() == torch::kInt32, "seqlens_k must have dtype int32"); + CHECK_DEVICE(seqlens_k); + CHECK_CONTIGUOUS(seqlens_k); + CHECK_SHAPE(seqlens_k, batch_size); + + params.new_lens = new_lens; + + params.seqlen_knew = seqlen_knew; + params.knew_ptr = k.data_ptr(); + params.vnew_ptr = v.data_ptr(); + params.knew_batch_stride = k.stride(0); + params.vnew_batch_stride = v.stride(0); + params.knew_row_stride = k.stride(-3); + params.vnew_row_stride = v.stride(-3); + params.knew_head_stride = k.stride(-2); + params.vnew_head_stride = v.stride(-2); + params.cu_seqlens_k = static_cast(seqlens_k.data_ptr()); + + const int pack_nums = 16 / num_bits; + + params.k_pack_new_ptr = k_pack_new.data_ptr(); + params.k_params_new_ptr = k_params_new.data_ptr(); + params.v_pack_new_ptr = v_pack_new.data_ptr(); + params.v_params_new_ptr = v_params_new.data_ptr(); + + params.k_pack_new_batch_stride = k_pack_new.stride(0); + params.k_params_new_batch_stride = k_params_new.stride(0); + params.v_pack_new_batch_stride = v_pack_new.stride(0); + params.v_params_new_batch_stride = v_params_new.stride(0); + + params.k_pack_new_row_stride = k_pack_new.stride(-3); + params.k_params_new_row_stride = k_params_new.stride(-3); + params.v_pack_new_row_stride = v_pack_new.stride(-3); + params.v_params_new_row_stride = v_params_new.stride(-1); + + params.k_pack_new_head_stride = k_pack_new.stride(-2); + params.k_params_new_head_stride = k_params_new.stride(-2); + params.v_pack_new_head_stride = v_pack_new.stride(-2); + params.v_params_new_head_stride = v_params_new.stride(-2); + + params.k_params_new_dim_stride = k_params_new.stride(-1); + params.v_params_new_dim_stride = v_params_new.stride(-3); + + } + + params.is_seqlens_k_cumulative = !(seqlens_k_.has_value()); params.rotary_dim = 0; if (paged_KV) { @@ -436,7 +522,7 @@ mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x out = out.transpose(1, 2).reshape({batch_size, 1, num_heads_k * seqlen_q, head_size_og}); } - return out; + return std::make_tuple(out, k_pack_new, k_params_new, v_pack_new, v_params_new); } @@ -513,20 +599,19 @@ void set_params_fprop_qpack(Flash_fwd_params ¶ms, params.group_size = group_size; } - template void kvcache_qpack(const at::Tensor &k, - at::Tensor &k_pack, - at::Tensor &k_params, - const at::Tensor &v, - at::Tensor &v_pack, - at::Tensor &v_params, - c10::optional &block_table_, - const at::Tensor &cu_seqlens_k, - const int max_seqlen_k, - const std::string quant_mode, - const int group_size - ) { + at::Tensor &k_pack, + at::Tensor &k_params, + const at::Tensor &v, + at::Tensor &v_pack, + at::Tensor &v_params, + c10::optional &block_table_, + const at::Tensor &cu_seqlens_k, + const int max_seqlen_k, + const std::string quant_mode, + const int group_size + ) { auto k_dtype = k.dtype(); TORCH_CHECK(k_dtype == torch::kFloat16 || k_dtype == torch::kBFloat16, diff --git a/csrc/bit_decode/src/flash_fwd_kernel.h b/csrc/bit_decode/src/flash_fwd_kernel.h index 1e0681a..226cfc6 100644 --- a/csrc/bit_decode/src/flash_fwd_kernel.h +++ b/csrc/bit_decode/src/flash_fwd_kernel.h @@ -63,6 +63,501 @@ inline __device__ void compute_attn_1rowblock(const Params ¶ms, const int bi //////////////////////////////////////////////////////////////////////////////////////////////////// +template +inline __device__ void compute_attn_1rowblock_residualkv(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int num_n_splits) { + + using Element = typename Kernel_traits::Element; + using ElementKVPack = typename Kernel_traits::ElementKVPack; + using ElementAccum = typename Kernel_traits::ElementAccum; + using index_t = typename Kernel_traits::index_t; + using SharedStorage = typename Kernel_traits::SharedStorage; + + // Shared memory. + extern __shared__ char smem_[]; + SharedStorage& shared_storage = *reinterpret_cast(smem_); + + // The thread index. + const int tidx = threadIdx.x; + + constexpr int kBlockM = Kernel_traits::kBlockM; + constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kBlockN_residual = Kernel_traits::kBlockN_residual; + constexpr int kBlockN_pack = Kernel_traits::kBlockN_pack; + constexpr int kBlockP = Kernel_traits::kBlockP; + constexpr int kBlockP_new_pack = Kernel_traits::kBlockP_new_pack; + constexpr int kBlockK_params = Kernel_traits::kBlockK_params; + constexpr int kBlockK_params_new = Kernel_traits::kBlockK_params_new; + constexpr int kHeadDim = Kernel_traits::kHeadDim; + constexpr int kHeadDim_pack = Kernel_traits::kHeadDim_pack; + constexpr int kHeadDim_k = Kernel_traits::kHeadDim_k; + constexpr int kHeadDim_k_params = Kernel_traits::kHeadDim_k_params; + constexpr int kHeadDim_v_params = Kernel_traits::kHeadDim_v_params; + constexpr int kNWarps = Kernel_traits::kNWarps; + constexpr int tile_paramsk_j = Kernel_traits::tile_paramsk_j; + constexpr int tile_paramsk_k = Kernel_traits::tile_paramsk_k; + constexpr int tile_paramsk_m = Kernel_traits::tile_paramsk_m; + constexpr int tile_paramsk_g = Kernel_traits::tile_paramsk_g; + constexpr int tile_paramsk_g_r = Kernel_traits::tile_paramsk_g_r; + constexpr int tile_paramsv_k = Kernel_traits::tile_paramsv_k; + constexpr int tile_paramsv_k_r = Kernel_traits::tile_paramsv_k_r; + constexpr int num_params = Kernel_traits::num_params; + constexpr int num_bits = Kernel_traits::num_bits; + constexpr int group_size = Kernel_traits::group_size; + constexpr int k_pack_div = Kernel_traits::k_pack_div; + constexpr int residual_block_size = Kernel_traits::residual_block_size; + + using GmemTiledCopyO = std::conditional_t< + !Split, + typename Kernel_traits::GmemTiledCopyO, + typename Kernel_traits::GmemTiledCopyOaccum + >; + using ElementO = std::conditional_t; + + const BlockInfo binfo(params, bidb); + if (m_block * kBlockM >= binfo.actual_seqlen_q) return; + + // Tensor, global memory + + // Q + Tensor mQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb)), + make_shape(binfo.actual_seqlen_q, params.h, params.d), + make_stride(params.q_row_stride, params.q_head_stride, _1{})); + Tensor gQ = local_tile(mQ(_, bidh, _), Shape, Int>{}, + make_coord(m_block, 0)); // (kBlockM, kHeadDim) + + // + // Tensor, Residual + // + + const int residual_len = params.new_lens; + const int n_blocks_residual = cute::ceil_div(residual_len, kBlockN_residual); + + const index_t row_offset_knew = bidb * params.knew_batch_stride + ((n_blocks_residual - 1) * kBlockN_residual) * params.knew_row_stride + + (bidh / params.h_h_k_ratio) * params.knew_head_stride; + const index_t row_offset_knew_pack = bidb * params.k_pack_new_batch_stride + 0 + + (bidh / params.h_h_k_ratio) * params.k_pack_new_head_stride; + const index_t row_offset_knew_params = bidb * params.k_params_new_batch_stride + 0 + + (bidh / params.h_h_k_ratio) * params.k_params_new_head_stride; + + const index_t row_offset_vnew = bidb * params.vnew_batch_stride + ((n_blocks_residual - 1) * kBlockN_residual) * params.vnew_row_stride + + (bidh / params.h_h_k_ratio) * params.vnew_head_stride; + const index_t row_offset_vnew_pack = bidb * params.v_pack_new_batch_stride + 0 + + (bidh / params.h_h_k_ratio) * params.v_pack_new_head_stride; + const index_t row_offset_vnew_params = bidb * params.v_params_new_batch_stride + 0 + + (bidh / params.h_h_k_ratio) * params.v_params_new_head_stride; + + Tensor gK_residual = make_tensor(make_gmem_ptr(reinterpret_cast(params.knew_ptr) + row_offset_knew), + Shape, Int>{}, + make_stride(params.knew_row_stride, _1{})); + Tensor gK_new_pack = make_tensor(make_gmem_ptr(reinterpret_cast(params.k_pack_new_ptr) + row_offset_knew_pack), + Shape, Int>{}, + make_stride(params.k_pack_new_row_stride, _1{})); + Tensor gK_new_params = make_tensor(make_gmem_ptr(reinterpret_cast<__half2*>(params.k_params_new_ptr) + row_offset_knew_params), + Shape, Int>{}, + make_stride(params.k_params_new_row_stride, params.k_params_new_dim_stride)); + + Tensor gV_residual = make_tensor(make_gmem_ptr(reinterpret_cast(params.vnew_ptr) + row_offset_vnew), + Shape, Int>{}, + make_stride(params.vnew_row_stride, _1{})); + Tensor gV_new_pack = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_pack_new_ptr) + row_offset_vnew_pack), + Shape, Int>{}, + make_stride(params.v_pack_new_row_stride, _1{})); + Tensor gV_new_params = make_tensor(make_gmem_ptr(reinterpret_cast<__half2*>(params.v_params_new_ptr) + row_offset_vnew_params), + Shape, Int>{}, + make_stride(params.v_params_new_row_stride, params.v_params_new_dim_stride)); + + #if DEBUG + if (threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { + printf("residual kernel \n"); + printf("params.new_lens = %d, binfo.actual_seqlen_k = %d, params.seqlen_k = %d, binfo.seqlen_k_cache = %d\n", params.new_lens, binfo.actual_seqlen_k, params.seqlen_k, binfo.seqlen_k_cache); + } + #endif + + // + // Tensor, shared memory + // + + // Q + Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_Q .data()), typename Kernel_traits::SmemLayoutQ{}); + + // K + Tensor sK_residual = make_tensor(make_smem_ptr(reinterpret_cast(shared_storage.smem_Kpack.data())), typename Kernel_traits::SmemLayoutKResidual{}); + Tensor sK_new_pack = make_tensor(make_smem_ptr(shared_storage.smem_Kpack .data()), typename Kernel_traits::SmemLayoutKNewPack{}); + Tensor sK_new_pack_transposed = make_tensor(sK_new_pack .data(), typename Kernel_traits::SmemLayoutKNewPacktransposed{}); + + // Acc + Tensor sAcc_residual = make_tensor(make_smem_ptr(shared_storage.smem_acc .data()), typename Kernel_traits::SmemLayoutAcc_residual{}); + Tensor sReduce_tmp = make_tensor(make_smem_ptr(shared_storage.smem_acc .data()), typename Kernel_traits::SmemLayoutReduce_tmp{}); + + // V + Tensor sV_residual = make_tensor(make_smem_ptr(reinterpret_cast(shared_storage.smem_Vpack.data())), typename Kernel_traits::SmemLayoutVResidual{}); + Tensor sVt_residual = make_tensor(sV_residual .data(), typename Kernel_traits::SmemLayoutVResidualtransposed{}); + Tensor sVtNoSwizzle_residual = make_tensor(sV_residual .data().get(), typename Kernel_traits::SmemLayoutVResidualtransposedNoSwizzle{}); + Tensor sV_new_pack = make_tensor(make_smem_ptr(shared_storage.smem_Vpack .data()), typename Kernel_traits::SmemLayoutVNewPack{}); + Tensor sVt_new_pack = make_tensor(sV_new_pack .data(), typename Kernel_traits::SmemLayoutVNewPacktransposed{}); + + // + // Copy, Global memory to shared memory + // + + typename Kernel_traits::GmemTiledCopyQKV gmem_tiled_copy_QKV; + typename Kernel_traits::GmemTileCopyKV_Residual gmem_tiled_copy_kv_residual; + typename Kernel_traits::GmemTileCopyK_Pack gmem_tiled_copy_k_pack; + typename Kernel_traits::GmemTileCopyV_Pack gmem_tiled_copy_v_pack; + typename Kernel_traits::GmemTileCopyKParams gmem_tiled_copy_k_params; + typename Kernel_traits::GmemTileCopyVParams gmem_tiled_copy_v_params; + + auto gmem_thr_copy_QKV = gmem_tiled_copy_QKV.get_thread_slice(tidx); + auto gmem_thr_copy_kv_residual = gmem_tiled_copy_kv_residual.get_thread_slice(tidx); + + Tensor tQgQ = gmem_thr_copy_QKV.partition_S(gQ); + Tensor tQsQ = gmem_thr_copy_QKV.partition_D(sQ); + + Tensor tKgK_residual = gmem_thr_copy_kv_residual.partition_S(gK_residual); + Tensor tKsK_residual = gmem_thr_copy_kv_residual.partition_D(sK_residual); + + Tensor tVgV_residual = gmem_thr_copy_kv_residual.partition_S(gV_residual); + Tensor tVsV_residual = gmem_thr_copy_kv_residual.partition_D(sV_residual); + + // + // Tensor, quantization params + // + + using TensorParamsKC = decltype(make_tensor(make_shape(Int<4 * num_params>{}, Int{}, Int{}))); + using TensorParamsVG = decltype(make_tensor(make_shape(Int{}, Int{}))); + using TensorParamsG = decltype(make_tensor(make_shape(Int{}))); + using TensorParamsKC_residual = decltype(make_tensor(make_shape(Int<4 * num_params>{}, Int{}))); + using TensorParamsKT_residual = decltype(make_tensor(make_shape(Int{}))); + using TensorParamsVT_residual = decltype(make_tensor(make_shape(Int{}, Int{}))); + + TensorParamsKC tScales_k_c, tZeros_k_c; + TensorParamsVG tScales_v_c, tZeros_v_c; + TensorParamsG tScales_k_g, tZeros_k_g; + TensorParamsKC_residual tScales_k_cr, tZeros_k_cr; + TensorParamsKT_residual tScales_k_tr, tZeros_k_tr; + TensorParamsVT_residual tScales_v_tr, tZeros_v_tr; + + auto tScales_k_h2_c = cute::recast<__half2>(tScales_k_c); + auto tZeros_k_h2_c = cute::recast<__half2>(tZeros_k_c); + auto tScales_k_h2_g = cute::recast<__half2>(tScales_k_g); + auto tZeros_k_h2_g = cute::recast<__half2>(tZeros_k_g); + auto tScales_k_h2_cr = cute::recast<__half2>(tScales_k_cr); + auto tZeros_k_h2_cr = cute::recast<__half2>(tZeros_k_cr); + auto tScales_k_h2_tr = cute::recast<__half2>(tScales_k_tr); + auto tZeros_k_h2_tr = cute::recast<__half2>(tZeros_k_tr); + + auto tScales_v_h2 = cute::recast<__half2>(tScales_v_c); + auto tZeros_v_h2 = cute::recast<__half2>(tZeros_v_c); + auto tScales_v_h2_tr = cute::recast<__half2>(tScales_v_tr); + auto tZeros_v_h2_tr = cute::recast<__half2>(tZeros_v_tr); + + // + // MMA Atom partitioning + // + + typename Kernel_traits::TiledMma tiled_mma; + typename Kernel_traits::TiledMmaKV_i4 tiled_mma_KV_i4; + typename Kernel_traits::TiledMma_residual tiled_mma_residual; + + auto thr_mma = tiled_mma.get_thread_slice(tidx); + auto thr_mma_KV_i4 = tiled_mma_KV_i4.get_thread_slice(tidx); + auto thr_mma_residual = tiled_mma_residual.get_thread_slice(tidx); + + Tensor tSrQ_residual = thr_mma_residual.partition_fragment_A(sQ); + + Tensor tSrK_residual = thr_mma_residual.partition_fragment_B(sK_residual); + // new pack + Tensor tSrK_new_pack_tmp = thr_mma_KV_i4.partition_fragment_B(sK_new_pack_transposed); + Tensor tSrK_new_pack = make_fragment_like(tSrK_new_pack_tmp); + + Tensor tOrVt_residual = thr_mma_residual.partition_fragment_B(sVtNoSwizzle_residual); + // new pack + Tensor tOrVt_new_pack_tmp = thr_mma_KV_i4.partition_fragment_B(sVt_new_pack); + Tensor tOrVt_new_pack = make_fragment_like(tOrVt_new_pack_tmp); + + Tensor acc_o = partition_fragment_C(tiled_mma, Shape, Int>{}); // MMA, MMA_M, MMA_K + + // + // Copy, Shared memory to register + // + + // Q + auto smem_tiled_copy_Q_residual = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtom{}, tiled_mma_residual); + auto smem_thr_copy_Q_residual = smem_tiled_copy_Q_residual.get_thread_slice(tidx); + Tensor tSsQ_residual = smem_thr_copy_Q_residual.partition_S(sQ); + + // K + auto smem_tiled_copy_K_residual = make_tiled_copy_B(typename Kernel_traits::S2RCopyAtomK{}, tiled_mma_residual); + auto smem_thr_copy_K_residual = smem_tiled_copy_K_residual.get_thread_slice(tidx); + + Tensor tSsK_residual = smem_thr_copy_K_residual.partition_S(sK_residual); + Tensor tSrK_residual_view = smem_thr_copy_K_residual.retile_D(tSrK_residual); + + // V + auto smem_tiled_copy_V_residual = make_tiled_copy_B(typename Kernel_traits::SmemCopyAtomTransposed_residual{}, tiled_mma_residual); + auto smem_thr_copy_V_residual = smem_tiled_copy_V_residual.get_thread_slice(tidx); + + Tensor tOsVt_residual = smem_thr_copy_V_residual.partition_S(sVt_residual); + Tensor tOrVt_residual_view = smem_thr_copy_V_residual.retile_D(tOrVt_residual); + + // + // Copy, Register to Shared memory + // + + auto smem_tiled_copy_kv_pack = make_tiled_copy_B(typename Kernel_traits::R2SCopyAtomPack{}, tiled_mma_KV_i4); + auto smem_thr_copy_kv_pack = smem_tiled_copy_kv_pack.get_thread_slice(tidx); + + Tensor tSrK_pack_r2s_view = smem_thr_copy_kv_pack.retile_S(tSrK_new_pack); + Tensor tSsK_pack_r2s = smem_thr_copy_kv_pack.partition_D(sK_new_pack); + + Tensor tOrVt_new_pack_r2s_view = smem_thr_copy_kv_pack.retile_S(tOrVt_new_pack); + Tensor tSsVt_new_pack_r2s = smem_thr_copy_kv_pack.partition_D(sVt_new_pack); + + // + // Copy, Shared memory to Global memory + // + + typename Kernel_traits::GmemTileCopyKV_NewPack gmem_tiled_copy_kv_newpack; + auto gmem_thr_copy_kv_newpack = gmem_tiled_copy_kv_newpack.get_thread_slice(tidx); + Tensor tKsK_new_pack_g2s = gmem_thr_copy_kv_newpack.partition_S(sK_new_pack); + Tensor tKgK_new_pack_g2s = gmem_thr_copy_kv_newpack.partition_D(gK_new_pack); + + Tensor tSsVt_new_pack_g2s = gmem_thr_copy_kv_newpack.partition_S(sV_new_pack); + Tensor tVgV_new_pack_g2s = gmem_thr_copy_kv_newpack.partition_D(gV_new_pack); + + // + // PREDICATES + // + + // Construct identity layout for sQ and sK + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + Tensor cKV_residual = make_identity_tensor(make_shape(size<0>(sK_residual), size<1>(sK_residual))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_QKV.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKVcKV_residual = gmem_thr_copy_QKV.partition_S(cKV_residual); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + // Allocate predicate tensors for k we assume is_even_k = true + Tensor tQpQ = make_tensor(make_shape(size<2>(tQsQ))); + Tensor tKVpKV_residual = make_tensor(make_shape(size<2>(tKsK_residual))); + + // + // Start of the main loop + // + + clear(acc_o); + + // Read Q from gmem to smem, optionally apply rotary embedding. + if (!Append_KV || params.rotary_dim == 0) { + // We don't need to clear the sQ smem tiles since we'll only write out the valid outputs + flash::copy(gmem_tiled_copy_QKV, tQgQ, tQsQ, tQcQ, tQpQ, + binfo.actual_seqlen_q - m_block * kBlockM); + } + cute::cp_async_fence(); + + flash::Softmax<2 * size<1>(acc_o)> softmax; + const float alibi_slope = !Has_alibi ? 0.0f : reinterpret_cast(params.alibi_slopes_ptr)[bidb * params.alibi_slopes_batch_stride + bidh] / params.scale_softmax; + flash::Mask mask_residual(params.new_lens, binfo.actual_seqlen_q, params.window_size_left, params.window_size_right, alibi_slope); + + // + // Residual loop + // + + int n_block_r = n_blocks_residual - 1; + + clear(tKsK_residual); + clear(tVsV_residual); + + // Advance gK_residual + flash::copy( + gmem_tiled_copy_kv_residual, tKgK_residual, tKsK_residual, tKVcKV_residual, tKVpKV_residual, params.new_lens - n_block_r * kBlockN_residual + ); + cute::cp_async_fence(); + + // Current Residual loops only one step + for (int residual_steps = 0; n_block_r >= 0; --n_block_r, ++residual_steps) { + Tensor acc_s = partition_fragment_C(tiled_mma_residual, Shape, Int>{}); // (MMA=4, MMA_M, MMA_N) + clear(acc_s); + flash::cp_async_wait<0>(); + __syncthreads(); + + // Advance gV_residual + if (residual_steps > 0) { + // TODO + } else { + // TODO: check clear_oob_mn + flash::copy( + gmem_tiled_copy_kv_residual, tVgV_residual, tVsV_residual, tKVcKV_residual, tKVpKV_residual, params.new_lens - n_block_r * kBlockN_residual + ); + } + cute::cp_async_fence(); + + // QK + flash::gemm_residual( + acc_s, tSrQ_residual, tSrK_residual, tSsQ_residual, tSsK_residual, tiled_mma_residual, smem_tiled_copy_Q_residual, smem_tiled_copy_K_residual, + smem_thr_copy_Q_residual, smem_thr_copy_K_residual + ); + + if (params.new_lens == residual_block_size) { + if (Kernel_traits::quant_mode == 1) { + quant::qpack_Kchannel_Vtensor(tSrK_residual, tSrK_new_pack, tScales_k_cr, tZeros_k_cr, sReduce_tmp, num_params); + quant::pack_Kchannel_store( + smem_tiled_copy_kv_pack, tSrK_pack_r2s_view, tSsK_pack_r2s, + gmem_tiled_copy_kv_newpack, tKsK_new_pack_g2s, tKgK_new_pack_g2s, + tScales_k_h2_cr, tZeros_k_h2_cr, gK_new_params, + num_params + ); + } else { + // Quant, Pack and Store K + quant::quant_Ktensor(tSrK_residual, tSrK_new_pack, tScales_k_tr, tZeros_k_tr, num_params); + quant::pack_Ktensor_store( + smem_tiled_copy_kv_pack, tSrK_pack_r2s_view, tSsK_pack_r2s, + gmem_tiled_copy_kv_newpack, tKsK_new_pack_g2s, tKgK_new_pack_g2s, + tScales_k_h2_tr, tZeros_k_h2_tr, gK_new_params, + num_params + ); + } + } + + // Mask + mask_residual.template apply_mask( + acc_s, 0, 0, 0 + ); + + flash::cp_async_wait<0>(); + __syncthreads(); + + // Advance gK + if (n_block_r > 0) { + // TODO + } + + residual_steps == 0 + ? softmax.template softmax_rescale_o(acc_s, acc_o, sReduce_tmp, params.scale_softmax_log2) + : softmax.template softmax_rescale_o(acc_s, acc_o, sReduce_tmp, params.scale_softmax_log2); + + // Convert acc_s from fp32 to fp16/bf16 + Tensor acc_s_fp16 = flash::convert_type(acc_s); + + auto r2s_tiled_copy_c = make_tiled_copy_C(typename Kernel_traits::R2SCopyAtomAcc{}, tiled_mma_residual); + auto r2s_thr_copy_c = r2s_tiled_copy_c.get_slice(tidx); + auto tCrAcc_r2s = r2s_thr_copy_c.retile_S(acc_s_fp16); + auto tCsAcc_r2s = r2s_thr_copy_c.partition_D(sAcc_residual); + + // copy acc_s from register to shared memory + // clear(tCsAcc_r2s); + cute::copy(r2s_tiled_copy_c, tCrAcc_r2s, tCsAcc_r2s); + __syncthreads(); + + Tensor tSrAcc = thr_mma_residual.partition_fragment_A(sAcc_residual); // (MMA,MMA_M,MMA_K) + auto smem_tiled_copy_Acc = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtom{}, tiled_mma_residual); + auto smem_thr_copy_Acc = smem_tiled_copy_Acc.get_thread_slice(tidx); + Tensor tSsAcc_view = smem_thr_copy_Acc.partition_S(sAcc_residual); + Tensor tSrAcc_view = smem_thr_copy_Acc.retile_D(tSrAcc); + + // PV + flash::gemm_residual( + acc_o, tSrAcc, tOrVt_residual, + tSsAcc_view, tOsVt_residual, + tiled_mma_residual, smem_tiled_copy_Acc, smem_tiled_copy_V_residual, + smem_thr_copy_Acc, smem_thr_copy_V_residual + ); + + // Quant, Pack and Store V + if (params.new_lens == residual_block_size) { + quant::qpack_Kchannel_Vtensor(tOrVt_residual, tOrVt_new_pack, tScales_v_tr, tZeros_v_tr, sReduce_tmp, num_params); + quant::pack_Vtensor_store( + smem_tiled_copy_kv_pack, tOrVt_new_pack_r2s_view, tSsVt_new_pack_r2s, + gmem_tiled_copy_kv_newpack, tSsVt_new_pack_g2s, tVgV_new_pack_g2s, + tScales_v_h2_tr, tZeros_v_h2_tr, gV_new_params, + num_params + ); + } + } + + // Epilogue + __syncthreads(); + Tensor lse = softmax.template normalize_softmax_lse(acc_o, sReduce_tmp, params.scale_softmax); + __syncthreads(); + // if (cute::thread0()) { print(lse); } + + Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) + // Partition sO to match the accumulator partitioning + using SmemTiledCopyO = std::conditional_t< + !Split, + typename Kernel_traits::SmemCopyAtomO, + typename Kernel_traits::SmemCopyAtomOaccum + >; + auto smem_tiled_copy_Oaccum = make_tiled_copy_C(SmemTiledCopyO{}, tiled_mma); + auto smem_thr_copy_Oaccum = smem_tiled_copy_Oaccum.get_thread_slice(tidx); + Tensor rO = flash::convert_type(acc_o); + Tensor taccOrOaccum = smem_thr_copy_Oaccum.retile_S(rO); // ((Atom,AtomNum), MMA_M, MMA_N) + Tensor taccOsOaccum = smem_thr_copy_Oaccum.partition_D(sOaccum); // ((Atom,AtomNum),PIPE_M,PIPE_N) + + // sOaccum is larger than sQ, so we need to syncthreads here + // TODO: allocate enough smem for sOaccum + if constexpr (Split) { __syncthreads(); } + + cute::copy(smem_tiled_copy_Oaccum, taccOrOaccum, taccOsOaccum); + + const index_t row_offset_o = binfo.q_offset(params.o_batch_stride, params.o_row_stride, bidb) + + m_block * kBlockM * params.o_row_stride + bidh * params.o_head_stride; + const index_t row_offset_oaccum = (((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q + + m_block * kBlockM) * params.d_rounded; + const index_t row_offset_lseaccum = (Split || !params.unpadded_lse ? + ((n_split_idx * params.b + bidb) * params.h + bidh) * params.seqlen_q : bidh * params.total_q + binfo.q_offset(params.seqlen_q, 1, bidb) + ) + m_block * kBlockM; + + Tensor gOaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.oaccum_ptr : params.o_ptr) + (Split ? row_offset_oaccum : row_offset_o)), + Shape, Int>{}, + make_stride(Split ? kHeadDim : params.o_row_stride, _1{})); + Tensor gLSEaccum = make_tensor(make_gmem_ptr(reinterpret_cast(Split ? params.softmax_lseaccum_ptr : params.softmax_lse_ptr) + row_offset_lseaccum), + Shape>{}, Stride<_1>{}); + // if (tidx == 0) { printf("row_offset_o = %d, bidh = %d, gOaccum = %p\n", row_offset_o, bidh, gOaccum.data()); } + + GmemTiledCopyO gmem_tiled_copy_Oaccum; + auto gmem_thr_copy_Oaccum = gmem_tiled_copy_Oaccum.get_thread_slice(tidx); + Tensor tOsOaccum = gmem_thr_copy_Oaccum.partition_S(sOaccum); // ((Atom,AtomNum),ATOM_M,ATOM_N) + Tensor tOgOaccum = gmem_thr_copy_Oaccum.partition_D(gOaccum); + + __syncthreads(); + + Tensor tOrOaccum = make_tensor(shape(tOgOaccum)); + cute::copy(gmem_tiled_copy_Oaccum, tOsOaccum, tOrOaccum); + + Tensor caccO = make_identity_tensor(Shape, Int>{}); // (BLK_M,BLK_K) -> (blk_m,blk_k) + Tensor taccOcO = thr_mma.partition_C(caccO); // (MMA,MMA_M,MMA_K) + static_assert(decltype(size<0>(taccOcO))::value == 4); + // Convert to ((2, 2), MMA_M, MMA_K) then take only the row indices. + Tensor taccOcO_row = logical_divide(taccOcO, Shape<_2>{})(make_coord(0, _), _, 0); + CUTE_STATIC_ASSERT_V(size(lse) == size(taccOcO_row)); // MMA_M + if (get<1>(taccOcO_row(0)) == 0) { + #pragma unroll + for (int mi = 0; mi < size(lse); ++mi) { + const int row = get<0>(taccOcO_row(mi)); + if (row < binfo.actual_seqlen_q - m_block * kBlockM) { gLSEaccum(row) = lse(mi); } + } + } + + // Construct identity layout for sO + Tensor cO = make_identity_tensor(make_shape(size<0>(sOaccum), size<1>(sOaccum))); // (BLK_M,BLK_K) -> (blk_m,blk_k) + // Repeat the partitioning with identity layouts + Tensor tOcO = gmem_thr_copy_Oaccum.partition_D(cO); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tOpO = make_tensor(make_shape(size<2>(tOgOaccum))); + if (!Is_even_K) { + #pragma unroll + for (int k = 0; k < size(tOpO); ++k) { tOpO(k) = get<1>(tOcO(0, 0, k)) < params.d; } + } + // Clear_OOB_K must be false since we don't want to write zeros to gmem + flash::copy( + gmem_tiled_copy_Oaccum, tOrOaccum, tOgOaccum, tOcO, tOpO, binfo.actual_seqlen_q - m_block * kBlockM + ); + +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + + template inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, const int bidb, const int bidh, const int m_block, const int n_split_idx, const int num_n_splits) { @@ -81,9 +576,12 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons constexpr int kBlockM = Kernel_traits::kBlockM; constexpr int kBlockN = Kernel_traits::kBlockN; + constexpr int kBlockN_residual = Kernel_traits::kBlockN_residual; constexpr int kBlockN_pack = Kernel_traits::kBlockN_pack; constexpr int kBlockP = Kernel_traits::kBlockP; + constexpr int kBlockP_new_pack = Kernel_traits::kBlockP_new_pack; constexpr int kBlockK_params = Kernel_traits::kBlockK_params; + constexpr int kBlockK_params_new = Kernel_traits::kBlockK_params_new; constexpr int kHeadDim = Kernel_traits::kHeadDim; constexpr int kHeadDim_pack = Kernel_traits::kHeadDim_pack; constexpr int kHeadDim_k = Kernel_traits::kHeadDim_k; @@ -94,10 +592,14 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons constexpr int tile_paramsk_k = Kernel_traits::tile_paramsk_k; constexpr int tile_paramsk_m = Kernel_traits::tile_paramsk_m; constexpr int tile_paramsk_g = Kernel_traits::tile_paramsk_g; + constexpr int tile_paramsk_g_r = Kernel_traits::tile_paramsk_g_r; constexpr int tile_paramsv_k = Kernel_traits::tile_paramsv_k; + constexpr int tile_paramsv_k_r = Kernel_traits::tile_paramsv_k_r; constexpr int num_params = Kernel_traits::num_params; constexpr int num_bits = Kernel_traits::num_bits; constexpr int group_size = Kernel_traits::group_size; + constexpr int k_pack_div = Kernel_traits::k_pack_div; + constexpr int residual_block_size = Kernel_traits::residual_block_size; using GmemTiledCopyO = std::conditional_t< !Split, @@ -109,7 +611,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons const BlockInfo binfo(params, bidb); if (m_block * kBlockM >= binfo.actual_seqlen_q) return; - const int n_blocks_per_split = ((binfo.actual_seqlen_k + kBlockN - 1) / kBlockN + num_n_splits - 1) / num_n_splits; + const int n_blocks_per_split = ((binfo.seqlen_k_cache + kBlockN - 1) / kBlockN + num_n_splits - 1) / num_n_splits; const int n_block_min = n_split_idx * n_blocks_per_split; int n_block_max = std::min(cute::ceil_div(binfo.actual_seqlen_k, kBlockN), (n_split_idx + 1) * n_blocks_per_split); @@ -154,11 +656,15 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons return; } - // We iterate over the blocks in reverse order. This is because the last block is the only one - // that needs masking when we read K and V from global memory. Moreover, iterating in reverse - // might save us 1 register (we just need n_block instead of both n_block and n_block_max). + #if DEBUG + if (threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { + printf("splitkv kernel \n"); + printf("params.new_lens = %d, binfo.actual_seqlen_k = %d, params.seqlen_k = %d, binfo.seqlen_k_cache = %d\n", params.new_lens, binfo.actual_seqlen_k, params.seqlen_k, binfo.seqlen_k_cache); + printf("bidb = %d, bidh = %d, m_block = %d, n_split_idx = %d, num_n_splits = %d\n", bidb, bidh, m_block, n_split_idx, num_n_splits); + printf("n_blocks_per_split = %d, n_block_min = %d, n_block_max = %d\n", n_blocks_per_split, n_block_min, n_block_max); + } + #endif - // We move K and V to the last block. const int bidb_cache = bidb; const int *block_table = !Paged_KV ? nullptr : params.block_table + bidb * params.block_table_batch_stride; const int block_table_idx = !Paged_KV ? 0 : (n_block_max - 1) * kBlockN / params.page_block_size; @@ -187,9 +693,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons const index_t row_offset_v_params = binfo.k_offset(params.v_params_batch_stride, params.v_params_row_stride, bidb) + (n_block_max - 1) * kBlockN * params.v_params_row_stride + (bidh / params.h_h_k_ratio) * params.v_params_head_stride; - // // Tensor, global memory - // // Q Tensor mQ = make_tensor(make_gmem_ptr(reinterpret_cast(params.q_ptr) + binfo.q_offset(params.q_batch_stride, params.q_row_stride, bidb)), @@ -198,15 +702,13 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons Tensor gQ = local_tile(mQ(_, bidh, _), Shape, Int>{}, make_coord(m_block, 0)); // (kBlockM, kHeadDim) - // K Tensor gK_pack = make_tensor(make_gmem_ptr(reinterpret_cast(params.K_pack_ptr) + row_offset_k_pack), - Shape, Int>{}, - make_stride(params.K_pack_row_stride, _1{})); + Shape, Int>{}, + make_stride(params.K_pack_row_stride, _1{})); Tensor gK_params = make_tensor(make_gmem_ptr(reinterpret_cast<__half2*>(params.k_params_ptr) + row_offset_k_params), Shape, Int>{}, make_stride(params.k_params_row_stride, _1{})); - // V Tensor gV_pack = make_tensor(make_gmem_ptr(reinterpret_cast(params.v_pack_ptr) + row_offset_v_pack), Shape, Int>{}, make_stride(params.v_pack_row_stride, _1{})); @@ -214,15 +716,6 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons Shape, Int>{}, make_stride(_1{}, params.v_params_dim_stride)); - #if DEBUG - if (threadIdx.x == 0 && threadIdx.y == 0 && blockIdx.x == 0 && blockIdx.y == 0 && blockIdx.z == 0) { - printf("num_n_splits = %d, n_blocks_per_split = %d\n", num_n_splits, n_blocks_per_split); - printf("params.new_lens = %d, binfo.actual_seqlen_k = %d\n", params.new_lens, binfo.actual_seqlen_k); - printf("params.seqlen_k = %d, binfo.seqlen_k_cache = %d\n", params.seqlen_k, binfo.seqlen_k_cache); - printf("n_block_min = %d, n_block_max = %d\n", n_block_min, n_block_max); - } - #endif - // // Tensor, shared memory // @@ -231,12 +724,9 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons Tensor sQ = make_tensor(make_smem_ptr(shared_storage.smem_Q .data()), typename Kernel_traits::SmemLayoutQ{}); // K - Tensor sK = make_tensor(make_smem_ptr(shared_storage.smem_Q .data()), typename Kernel_traits::SmemLayoutKV{}); - Tensor sK_new_pack = make_tensor(make_smem_ptr(shared_storage.smem_Kpack .data()), typename Kernel_traits::SmemLayoutKNewPack{}); - Tensor sK_new_pack_transposed = make_tensor(sK_new_pack .data(), typename Kernel_traits::SmemLayoutKNewPacktransposed{}); Tensor sK_dequant = make_tensor(make_smem_ptr(shared_storage.smem_Q .data()), typename Kernel_traits::SmemLayoutKV{}); Tensor sK_pack = make_tensor(make_smem_ptr(shared_storage.smem_Kpack .data()), typename Kernel_traits::SmemLayoutKPack{}); - Tensor sK_pack_transposed = make_tensor(sK_pack .data(), typename Kernel_traits::SmemLayoutKPacktransposed{}); + Tensor sK_pack_transposed = make_tensor(sK_pack .data(), typename Kernel_traits::SmemLayoutKPacktransposed{}); Tensor sK_params = make_tensor(make_smem_ptr(shared_storage.smem_Kparams .data()), typename Kernel_traits::SmemLayoutKParams{}); // Acc @@ -266,13 +756,13 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons auto gmem_thr_copy_k_params = gmem_tiled_copy_k_params.get_thread_slice(tidx); auto gmem_thr_copy_v_params = gmem_tiled_copy_v_params.get_thread_slice(tidx); - Tensor tQgQ = gmem_thr_copy_QKV.partition_S(gQ); - Tensor tQsQ = gmem_thr_copy_QKV.partition_D(sQ); + Tensor tQgQ = gmem_thr_copy_QKV.partition_S(gQ); + Tensor tQsQ = gmem_thr_copy_QKV.partition_D(sQ); - Tensor tKgK_pack = gmem_thr_copy_k_pack.partition_S(gK_pack); - Tensor tKsK_pack = gmem_thr_copy_k_pack.partition_D(sK_pack); - Tensor tKgK_params = gmem_thr_copy_k_params.partition_S(gK_params); - Tensor tKsK_params = gmem_thr_copy_k_params.partition_D(sK_params); + Tensor tKgK_pack = gmem_thr_copy_k_pack.partition_S(gK_pack); + Tensor tKsK_pack = gmem_thr_copy_k_pack.partition_D(sK_pack); + Tensor tKgK_params = gmem_thr_copy_k_params.partition_S(gK_params); + Tensor tKsK_params = gmem_thr_copy_k_params.partition_D(sK_params); Tensor tVgV_pack = gmem_thr_copy_v_pack.partition_S(gV_pack); Tensor tVsV_pack = gmem_thr_copy_v_pack.partition_D(sV_pack); @@ -283,20 +773,33 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons // Tensor, quantization params // - using TensorParamsKC = decltype(make_tensor(make_shape(Int<4 * num_params>{}, Int{}, Int{}))); - using TensorParamsVG = decltype(make_tensor(make_shape(Int{}, Int{}))); - using TensorParamsG = decltype(make_tensor(make_shape(Int{}))); + using TensorParamsKC = decltype(make_tensor(make_shape(Int<4 * num_params>{}, Int{}, Int{}))); + using TensorParamsVG = decltype(make_tensor(make_shape(Int{}, Int{}))); + using TensorParamsG = decltype(make_tensor(make_shape(Int{}))); + using TensorParamsKC_residual = decltype(make_tensor(make_shape(Int<4 * num_params>{}, Int{}))); + using TensorParamsKT_residual = decltype(make_tensor(make_shape(Int{}))); + using TensorParamsVT_residual = decltype(make_tensor(make_shape(Int{}, Int{}))); TensorParamsKC tScales_k_c, tZeros_k_c; TensorParamsVG tScales_v_c, tZeros_v_c; TensorParamsG tScales_k_g, tZeros_k_g; + TensorParamsKC_residual tScales_k_cr, tZeros_k_cr; + TensorParamsKT_residual tScales_k_tr, tZeros_k_tr; + TensorParamsVT_residual tScales_v_tr, tZeros_v_tr; + auto tScales_k_h2_c = cute::recast<__half2>(tScales_k_c); auto tZeros_k_h2_c = cute::recast<__half2>(tZeros_k_c); auto tScales_k_h2_g = cute::recast<__half2>(tScales_k_g); auto tZeros_k_h2_g = cute::recast<__half2>(tZeros_k_g); + auto tScales_k_h2_cr = cute::recast<__half2>(tScales_k_cr); + auto tZeros_k_h2_cr = cute::recast<__half2>(tZeros_k_cr); + auto tScales_k_h2_tr = cute::recast<__half2>(tScales_k_tr); + auto tZeros_k_h2_tr = cute::recast<__half2>(tZeros_k_tr); auto tScales_v_h2 = cute::recast<__half2>(tScales_v_c); auto tZeros_v_h2 = cute::recast<__half2>(tZeros_v_c); + auto tScales_v_h2_tr = cute::recast<__half2>(tScales_v_tr); + auto tZeros_v_h2_tr = cute::recast<__half2>(tZeros_v_tr); // // MMA Atom partitioning @@ -314,11 +817,11 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons Tensor tSrK_pack_tmp = thr_mma_KV_i4.partition_fragment_B(sK_pack_transposed); Tensor tSrK_pack = make_fragment_like(tSrK_pack_tmp); - Tensor tOrVt_dequant = thr_mma.partition_fragment_B(sVtNoSwizzle_dequant); - Tensor tOrVt_pack_tmp = thr_mma_KV_i4.partition_fragment_B(sVtNoSwizzle_pack); - Tensor tOrVt_pack = make_fragment_like(tOrVt_pack_tmp); + Tensor tOrVt_dequant = thr_mma.partition_fragment_B(sVtNoSwizzle_dequant); + Tensor tOrVt_pack_tmp = thr_mma_KV_i4.partition_fragment_B(sVtNoSwizzle_pack); + Tensor tOrVt_pack = make_fragment_like(tOrVt_pack_tmp); - Tensor acc_o = partition_fragment_C(tiled_mma, Shape, Int>{}); // MMA, MMA_M, MMA_K + Tensor acc_o = partition_fragment_C(tiled_mma, Shape, Int>{}); // MMA, MMA_M, MMA_K // // Tensor, ACC @@ -332,7 +835,6 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons auto smem_tiled_copy_Acc = make_tiled_copy_A(typename Kernel_traits::SmemCopyAtom{}, tiled_mma); auto smem_thr_copy_Acc = smem_tiled_copy_Acc.get_thread_slice(tidx); Tensor tSsAcc_view = smem_thr_copy_Acc.partition_S(sAcc); - Tensor tSrAcc_view = smem_thr_copy_Acc.retile_D(tSrAcc); // // Copy, Shared memory to register @@ -348,6 +850,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons auto smem_tiled_copy_K_pack = make_tiled_copy_B(typename Kernel_traits::S2RCopyAtomK_i4{}, tiled_mma_KV_i4); auto smem_thr_copy_K = smem_tiled_copy_K.get_thread_slice(tidx); auto smem_thr_copy_K_pack = smem_tiled_copy_K_pack.get_thread_slice(tidx); + Tensor tSsK_pack = smem_thr_copy_K_pack.partition_S(sK_pack); Tensor tSrK_pack_view = smem_thr_copy_K_pack.retile_D(tSrK_pack); @@ -365,9 +868,23 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons // // Construct identity layout for sQ and sK - Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_N,BLK_K) -> (blk_n,blk_k) - Tensor tQcQ = gmem_thr_copy_QKV.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) - Tensor tQpQ = make_tensor(make_shape(size<2>(tQsQ))); + Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + Tensor cK_pack = make_identity_tensor(make_shape(size<0>(sK_pack), size<1>(sK_pack))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + Tensor cV_pack = make_identity_tensor(make_shape(size<0>(sV_pack), size<1>(sV_pack))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + Tensor cK_params = make_identity_tensor(make_shape(size<0>(sK_params), size<1>(sK_params))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + Tensor cV_params = make_identity_tensor(make_shape(size<0>(sV_params), size<1>(sV_params))); // (BLK_N,BLK_K) -> (blk_n,blk_k) + + // Repeat the partitioning with identity layouts + Tensor tQcQ = gmem_thr_copy_QKV.partition_S(cQ); // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k) + Tensor tKcK_pack = gmem_thr_copy_k_pack.partition_S(cK_pack); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + Tensor tKcK_params = gmem_thr_copy_k_params.partition_S(cK_params); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + Tensor tVcV_pack = gmem_thr_copy_v_pack.partition_S(cV_pack); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + Tensor tVcV_params = gmem_thr_copy_v_params.partition_S(cV_params); // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k) + + // Allocate predicate tensors for k we assume is_even_k = true + Tensor tQpQ = make_tensor(make_shape(size<2>(tQsQ))); + Tensor tKVpKV_pack = make_tensor(make_shape(size<2>(tKsK_pack))); + Tensor tKVpKV_params = make_tensor(make_shape(size<2>(tKsK_params))); // // Start of the main loop @@ -397,10 +914,16 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons clear(tVsV_params); clear(tKsK_pack); clear(tKsK_params); - - // Advance gK - cute::copy(gmem_tiled_copy_k_pack, tKgK_pack, tKsK_pack); - cute::copy(gmem_tiled_copy_k_params, tKgK_params, tKsK_params); + + const bool is_short_block = num_bits == 4 && (binfo.actual_seqlen_k - n_block * kBlockN) % kBlockN < residual_block_size; + const int k_params_div = is_short_block && Kernel_traits::quant_mode == 1 ? group_size : 1; + + flash::copy( + gmem_tiled_copy_k_pack, tKgK_pack, tKsK_pack, tKcK_pack, tKVpKV_pack, (binfo.seqlen_k_cache - n_block * kBlockN) / k_pack_div + ); + flash::copy( + gmem_tiled_copy_k_params, tKgK_params, tKsK_params, tKcK_params, tKVpKV_params, (binfo.seqlen_k_cache - n_block * kBlockN) / k_params_div + ); cute::cp_async_fence(); const int n_masking_steps = 1; @@ -413,7 +936,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons __syncthreads(); // Advance gV - if ((masking_step > 0)) { + if (masking_step > 0) { if (block_table == nullptr) { tVgV_pack.data() = tVgV_pack.data() + (-int(kBlockN * params.v_pack_row_stride)); tVgV_params.data() = tVgV_params.data() + (-int(kBlockN * params.v_params_row_stride)); @@ -426,11 +949,22 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons tVgV_pack.data() = tVgV_pack.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.v_pack_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_pack_row_stride; tVgV_params.data() = tVgV_params.data() + (-int(kBlockN * params.v_params_row_stride)); } - cute::copy(gmem_tiled_copy_v_pack, tVgV_pack, tVsV_pack); - cute::copy(gmem_tiled_copy_v_params, tVgV_params, tVsV_params); + // cute::copy(gmem_tiled_copy_QKV, tVgV, tVsV); + flash::copy( + gmem_tiled_copy_v_pack, tVgV_pack, tVsV_pack, tVcV_pack, tKVpKV_pack, binfo.seqlen_k_cache - n_block * kBlockN + ); + flash::copy( + gmem_tiled_copy_v_params, tVgV_params, tVsV_params, tVcV_params, tKVpKV_params, binfo.seqlen_k_cache - n_block * kBlockN + ); + } else { - cute::copy(gmem_tiled_copy_v_pack, tVgV_pack, tVsV_pack); - cute::copy(gmem_tiled_copy_v_params, tVgV_params, tVsV_params); + // cute::copy(gmem_tiled_copy_QKV, tVgV, tVsV); + flash::copy( + gmem_tiled_copy_v_pack, tVgV_pack, tVsV_pack, tVcV_pack, tKVpKV_pack, binfo.seqlen_k_cache - n_block * kBlockN + ); + flash::copy( + gmem_tiled_copy_v_params, tVgV_params, tVsV_params, tVcV_params, tKVpKV_params, binfo.seqlen_k_cache - n_block * kBlockN + ); } cute::cp_async_fence(); @@ -483,6 +1017,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons if (n_block > n_block_min) { // Advance gK if (block_table == nullptr) { + // tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); tKgK_pack.data() = tKgK_pack.data() + (-int(kBlockP * params.K_pack_row_stride)); tKgK_params.data() = tKgK_params.data() + (-int(kBlockK_params * params.k_params_row_stride)); } else { @@ -492,9 +1027,11 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons const int block_table_idx_next = (n_block - 1) * kBlockN / params.page_block_size; const int block_table_offset_next = (n_block - 1) * kBlockN - block_table_idx_next * params.page_block_size; const int block_table_offset_pack_next = (n_block - 1) * kBlockP - block_table_idx_next * params.page_block_size_pack; + // tKgK.data() = tKgK.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.k_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.k_row_stride; tKgK_pack.data() = tKgK_pack.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.K_pack_batch_stride + (block_table_offset_pack_next - block_table_offset_pack_cur) * params.K_pack_row_stride; tKgK_params.data() = tKgK_params.data() + (-int(kBlockK_params * params.k_params_row_stride)); } + // cute::copy(gmem_tiled_copy_QKV, tKgK, tKsK); cute::copy(gmem_tiled_copy_k_pack, tKgK_pack, tKsK_pack); cute::copy(gmem_tiled_copy_k_params, tKgK_params, tKsK_params); // This cp_async_fence needs to be in the if block, otherwise the synchronization @@ -530,7 +1067,6 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons smem_thr_copy_V_pack, num_params ); - } // These are the iterations where we don't need masking on S @@ -549,10 +1085,11 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons const int block_table_offset_cur = (n_block + 1) * kBlockN - block_table_idx_cur * params.page_block_size; const int block_table_idx_next = n_block * kBlockN / params.page_block_size; const int block_table_offset_next = n_block * kBlockN - block_table_idx_next * params.page_block_size; + // tVgV.data() = tVgV.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.v_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_row_stride; tVgV_pack.data() = tVgV_pack.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.v_pack_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.v_pack_row_stride; tVgV_params.data() = tVgV_params.data() + (-int(kBlockN * params.v_params_row_stride)); } - + // cute::copy(gmem_tiled_copy_QKV, tVgV, tVsV); cute::copy(gmem_tiled_copy_v_pack, tVgV_pack, tVsV_pack); cute::copy(gmem_tiled_copy_v_params, tVgV_params, tVsV_params); cute::cp_async_fence(); @@ -601,6 +1138,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons if (n_block > n_block_min) { // Advance gK if (block_table == nullptr) { + // tKgK.data() = tKgK.data() + (-int(kBlockN * params.k_row_stride)); tKgK_pack.data() = tKgK_pack.data() + (-int(kBlockP * params.K_pack_row_stride)); tKgK_params.data() = tKgK_params.data() + (-int(kBlockK_params * params.k_params_row_stride)); @@ -611,9 +1149,11 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons const int block_table_idx_next = (n_block - 1) * kBlockN / params.page_block_size; const int block_table_offset_next = (n_block - 1) * kBlockN - block_table_idx_next * params.page_block_size; const int block_table_offset_pack_next = (n_block - 1) * kBlockP - block_table_idx_next * params.page_block_size_pack; + // tKgK.data() = tKgK.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.k_batch_stride + (block_table_offset_next - block_table_offset_cur) * params.k_row_stride; tKgK_pack.data() = tKgK_pack.data() + (block_table[block_table_idx_next] - block_table[block_table_idx_cur]) * params.K_pack_batch_stride + (block_table_offset_pack_next - block_table_offset_pack_cur) * params.K_pack_row_stride; tKgK_params.data() = tKgK_params.data() + (-int(kBlockK_params * params.k_params_row_stride)); } + // cute::copy(gmem_tiled_copy_QKV, tKgK, tKsK); cute::copy(gmem_tiled_copy_k_pack, tKgK_pack, tKsK_pack); cute::copy(gmem_tiled_copy_k_params, tKgK_params, tKsK_params); // This cp_async_fence needs to be in the if block, otherwise the synchronization @@ -652,6 +1192,7 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons __syncthreads(); Tensor lse = softmax.template normalize_softmax_lse(acc_o, sReduce_tmp, params.scale_softmax); __syncthreads(); + // if (cute::thread0()) { print(lse); } Tensor sOaccum = make_tensor(make_smem_ptr(reinterpret_cast(smem_)), typename Kernel_traits::SmemLayoutO{}); // (SMEM_M,SMEM_N) // Partition sO to match the accumulator partitioning @@ -724,7 +1265,6 @@ inline __device__ void compute_attn_1rowblock_splitkv(const Params ¶ms, cons flash::copy( gmem_tiled_copy_Oaccum, tOrOaccum, tOgOaccum, tOcO, tOpO, binfo.actual_seqlen_q - m_block * kBlockM ); - } //////////////////////////////////////////////////////////////////////////////////////////////////// @@ -893,7 +1433,7 @@ inline __device__ void compute_qpack_1rowblock(const Params ¶ms, const int b Tensor tSrV_view = smem_thr_copy_V.retile_D(tSrV); Tensor tSsV_pack_s2r = smem_thr_copy_V_pack.partition_S(sVt_pack); Tensor tSrV_pack_s2r_view = smem_thr_copy_V_pack.retile_D(tSrV_pack); - + // Advance gK cute::copy(gmem_tiled_copy_QKV, tKgK, tKsK); @@ -914,7 +1454,7 @@ inline __device__ void compute_qpack_1rowblock(const Params ¶ms, const int b TensorParamsKC tScales_k_c, tZeros_k_c; TensorParamsVG tScales_v_c, tZeros_v_c; - TensorParamsG tScales_k_g, tZeros_k_g; + TensorParamsG tScales_k_g, tZeros_k_g; if (Kernel_traits::quant_mode == 1) { quant::qpack_Kchannel_Vtensor(tSrK, tSrK_pack, tScales_k_c, tZeros_k_c, sReduce_tmp, num_params); @@ -979,6 +1519,7 @@ inline __device__ void compute_qpack_1rowblock(const Params ¶ms, const int b cute::copy(gmem_tiled_copy_k_pack, tKsK_pack_s2g, tKgK_pack_s2g); __syncthreads(); cute::copy(gmem_tiled_copy_v_pack, tVsV_pack_s2g, tVgV_pack_s2g); + __syncthreads(); // ////////////////////////////////////////////////////////////////////////////// // // verify the quantize @@ -1018,7 +1559,7 @@ inline __device__ void compute_qpack_1rowblock(const Params ¶ms, const int b // quant::dequant_Kchannel_Vtensor(tSrV_pack(_,_,i), tSrV_dequant(_,_,i), tScales_v_c(_,i), tZeros_v_c(_,i), num_params); // } - // if (Kernel_traits::quant_mode == 1) { + if (Kernel_traits::quant_mode == 1) { // CUTE_UNROLL // for (int i = 0; i < size<1>(tScales_k_h2_c); ++i) { // CUTE_UNROLL @@ -1032,7 +1573,7 @@ inline __device__ void compute_qpack_1rowblock(const Params ¶ms, const int b // for (int i = 0; i < size<2>(tSrK_pack); ++i) { // quant::dequant_Kchannel_Vtensor(tSrK_pack(_,_,i), tSrK_dequant(_,_,i), tScales_k_c(_,i), tZeros_k_c(_,i), num_params); // } - // } else { + } else { // CUTE_UNROLL // for (int j = 0; j < size<0>(tScales_k_h2_g); ++j) { // tScales_k_h2_g(j) = gK_params(0 + 32*j + tidx/4, 0); @@ -1042,11 +1583,11 @@ inline __device__ void compute_qpack_1rowblock(const Params ¶ms, const int b // auto tScales_k_h1_g = cute::recast<__half>(tScales_k_h2_g); // auto tZeros_k_h1_g = cute::recast<__half>(tZeros_k_h2_g); - // CUTE_UNROLL - // for (int i = 0; i < size<2>(tSrK_pack); ++i) { - // quant::dequantize_Ktensor(tSrK_pack, tSrK_dequant, tScales_k_h2_g, tZeros_k_h2_g, 4, group_size, i); - // } - // } + CUTE_UNROLL + for (int i = 0; i < size<2>(tSrK_pack); ++i) { + quant::dequantize_Ktensor(tSrK_pack, tSrK_dequant, tScales_k_h2_g, tZeros_k_h2_g, 4, group_size, i); + } + } // // ////////////////////////////////////////////////////////////////////////////// #if DEBUG2 @@ -1127,13 +1668,27 @@ inline __device__ void compute_attn(const Params ¶ms) { //////////////////////////////////////////////////////////////////////////////////////////////////// +template +inline __device__ void compute_attn_residualkv(const Params ¶ms) { + const int m_block = blockIdx.x; + // The block index for the batch. + const int bidb = blockIdx.y; + // The block index for the head. + const int bidh = blockIdx.z; + const int n_split_idx = params.num_splits - 1; + const int num_n_splits = 1; + flash::compute_attn_1rowblock_residualkv(params, bidb, bidh, m_block, n_split_idx, num_n_splits); +} + +//////////////////////////////////////////////////////////////////////////////////////////////////// + template inline __device__ void compute_attn_splitkv(const Params ¶ms) { const int m_block = blockIdx.x; // The block index for the batch. - const int bidb = Split ? blockIdx.z / params.h : blockIdx.y; + const int bidb = Split ? blockIdx.z / params.h : blockIdx.y; // The block index for the head. - const int bidh = Split ? blockIdx.z - bidb * params.h : blockIdx.z; + const int bidh = Split ? blockIdx.z - bidb * params.h : blockIdx.z; const int n_split_idx = Split ? blockIdx.y : 0; const int num_n_splits = Split ? gridDim.y : 1; flash::compute_attn_1rowblock_splitkv(params, bidb, bidh, m_block, n_split_idx, num_n_splits); diff --git a/csrc/bit_decode/src/flash_fwd_launch_template.h b/csrc/bit_decode/src/flash_fwd_launch_template.h index a90b4d9..6cf77ee 100644 --- a/csrc/bit_decode/src/flash_fwd_launch_template.h +++ b/csrc/bit_decode/src/flash_fwd_launch_template.h @@ -47,6 +47,14 @@ DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_kernel, bool Is_dropout, bool Is_causal, b #endif } +DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_residual_kernel, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K, bool Is_softcap, bool Split, bool Append_KV, bool Paged_KV) { + #if defined(ARCH_SUPPORTS_FLASH) + flash::compute_attn_residualkv(params); + #else + FLASH_UNSUPPORTED_ARCH + #endif +} + DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_splitkv_kernel, bool Is_causal, bool Is_local, bool Has_alibi, bool Is_even_MN, bool Is_even_K, bool Is_softcap, bool Split, bool Append_KV, bool Paged_KV) { #if defined(ARCH_SUPPORTS_FLASH) flash::compute_attn_splitkv(params); @@ -62,92 +70,40 @@ DEFINE_FLASH_FORWARD_KERNEL(flash_fwd_splitkv_combine_kernel, int kBlockM, int L template void run_flash_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { - constexpr size_t smem_size = Kernel_traits::kSmemSize; - // printf("smem_size = %d\n", smem_size); - - // Work-around for gcc 7. It doesn't like nested BOOL_SWITCH. - // https://github.com/kokkos/kokkos-kernels/issues/349 - // https://github.com/HazyResearch/flash-attention/issues/21 - const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM; - dim3 grid(num_m_block, params.b, params.h); - const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr && params.seqlen_k % Kernel_traits::kBlockN == 0 && params.seqlen_q % Kernel_traits::kBlockM == 0; - const bool is_even_K = params.d == Kernel_traits::kHeadDim; - const bool return_softmax = params.p_ptr != nullptr; - // BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] { - // EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { - - // Will only return softmax if dropout, to reduce compilation time. - // If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates. - // If return_softmax, set IsEvenMNConst to false to reduce number of templates - // If head dim > 128, set IsEvenMNConst to false to reduce number of templates - // If Is_local, set Is_causal to false - auto kernel = &flash_fwd_kernel; - // auto kernel = &flash_fwd_kernel; - // printf("IsEvenMNConst = %d, IsEvenKConst = %d, Is_local = %d, Is_causal = %d, ReturnSoftmaxConst = %d, Is_dropout = %d\n", int(IsEvenMNConst), int(IsEvenKConst), int(Is_local), int(Is_causal), int(ReturnSoftmaxConst), int(Is_dropout)); - // auto kernel = &flash_fwd_kernel; - if (smem_size >= 48 * 1024) { - C10_CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - // int ctas_per_sm; - // cudaError status_ = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - // &ctas_per_sm, kernel, Kernel_traits::kNThreads, smem_size); - // printf("smem_size = %d, CTAs per SM = %d\n", int(smem_size), ctas_per_sm); - kernel<<>>(params); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - // }); - // }); } template void run_flash_splitkv_fwd(Flash_fwd_params ¶ms, cudaStream_t stream) { static_assert(!Kernel_traits::Is_Q_in_regs, "SplitKV implementation does not support Is_Q_in_regs"); static_assert(!Kernel_traits::Share_Q_K_smem, "SplitKV implementation does not support Share_Q_K_smem"); - constexpr size_t smem_size = Kernel_traits::kSmemSize; - const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM; - - dim3 grid(num_m_block, params.num_splits > 1 ? params.num_splits : params.b, params.num_splits > 1 ? params.b * params.h : params.h); + constexpr size_t smem_size = Kernel_traits::kSmemSize; + constexpr size_t smem_size_res = Kernel_traits::kSmemSize_res; + const int num_m_block = (params.seqlen_q + Kernel_traits::kBlockM - 1) / Kernel_traits::kBlockM; + const int num_splits_ = params.num_splits - 1; + + dim3 grid_res(num_m_block, params.b, params.h); + dim3 grid(num_m_block, num_splits_, params.b * params.h); const bool is_even_MN = params.cu_seqlens_q == nullptr && params.cu_seqlens_k == nullptr && params.seqlen_k % Kernel_traits::kBlockN == 0 && params.seqlen_q % Kernel_traits::kBlockM == 0; const bool is_even_K = params.d == Kernel_traits::kHeadDim; - // BOOL_SWITCH(is_even_MN, IsEvenMNConst, [&] { - // EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { - // LOCAL_SWITCH((params.window_size_left >= 0 || params.window_size_right >= 0) && !Is_causal, Is_local, [&] { - BOOL_SWITCH(params.num_splits > 1, Split, [&] { - // BOOL_SWITCH(params.knew_ptr != nullptr, Append_KV, [&] { - BOOL_SWITCH(params.block_table != nullptr, Paged_KV, [&] { - // ALIBI_SWITCH(params.alibi_slopes_ptr != nullptr, Has_alibi, [&] { - // SOFTCAP_SWITCH(params.softcap > 0.0, Is_softcap, [&] { - // If Append_KV, then we must have seqlen_offsets, which means cu_seqlens_k != nullptr. - // If not IsEvenKConst, we also set IsEvenMNConst to false to reduce number of templates. - // If Is_local, set Is_causal to false - // IsEvenMNConst: 0 - // IsEvenKConst: 1 - // Is_local: 0 - // Split: 1 - // Append_KV: - // Has_alibi: 0 - // Is_softcap: 0 - auto kernel = &flash_fwd_splitkv_kernel; - // auto kernel = &flash_fwd_splitkv_kernel; - // auto kernel = &flash_fwd_splitkv_kernel; - if (smem_size >= 48 * 1024) { - C10_CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - } - kernel<<>>(params); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - // }); - }); - // }); - }); - // }); - // }); - // }); + + auto kernel_res = &flash_fwd_residual_kernel; + if (smem_size_res >= 48 * 1024) { + C10_CUDA_CHECK(cudaFuncSetAttribute( + kernel_res, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_res)); + } + kernel_res<<>>(params); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + auto kernel = &flash_fwd_splitkv_kernel; + if (smem_size >= 48 * 1024) { + C10_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } + kernel<<>>(params); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + if (params.num_splits > 1) { - // We want kBlockM to be as small as possible for more parallelism. - // With 128 threads we can load 512 elements at a time, so if headdim is divisible by 128, kBlockM = 4. - // If headdim is divisible by 64, then we set kBlockM = 8, etc. constexpr static int kBlockM = Kernel_traits::kHeadDim % 128 == 0 ? 4 : (Kernel_traits::kHeadDim % 64 == 0 ? 8 : 16); dim3 grid_combine((params.b * params.h * params.seqlen_q + kBlockM - 1) / kBlockM); EVENK_SWITCH(is_even_K, IsEvenKConst, [&] { diff --git a/csrc/bit_decode/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu b/csrc/bit_decode/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu index dfdd014..bd792d8 100644 --- a/csrc/bit_decode/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu +++ b/csrc/bit_decode/src/genfile/flash_fwd_split_hdim128_fp16_sm80_4bit.cu @@ -11,4 +11,4 @@ template void run_mha_fwd_splitkv_dispatch(Flash_fwd_params ¶ms, cudaStream_t stream); // template void run_mha_fwd_splitkv_dispatch(Flash_fwd_params ¶ms, cudaStream_t stream); -// template void run_mha_fwd_splitkv_dispatch(Flash_fwd_params ¶ms, cudaStream_t stream); +template void run_mha_fwd_splitkv_dispatch(Flash_fwd_params ¶ms, cudaStream_t stream); diff --git a/csrc/bit_decode/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu b/csrc/bit_decode/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu index 0fa4011..2132d45 100644 --- a/csrc/bit_decode/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu +++ b/csrc/bit_decode/src/genfile/flash_qpack_hdim128_fp16_sm80_4bit.cu @@ -4,10 +4,10 @@ #include "../flash_fwd_launch_template.h" -// template<> -// void run_kvcache_qpack_(Flash_fwd_params ¶ms, cudaStream_t stream) { -// run_kvcache_qpack_hdim128(params, stream); -// } +template<> +void run_kvcache_qpack_(Flash_fwd_params ¶ms, cudaStream_t stream) { + run_kvcache_qpack_hdim128(params, stream); +} // template<> // void run_kvcache_qpack_(Flash_fwd_params ¶ms, cudaStream_t stream) { // run_kvcache_qpack_hdim128(params, stream); diff --git a/csrc/bit_decode/src/include/kernel_traits.h b/csrc/bit_decode/src/include/kernel_traits.h index 394dac1..2c6ab3d 100644 --- a/csrc/bit_decode/src/include/kernel_traits.h +++ b/csrc/bit_decode/src/include/kernel_traits.h @@ -2,567 +2,575 @@ * Copyright (c) 2024, Tri Dao. ******************************************************************************/ - #pragma once - - #include "cute/tensor.hpp" - - #include "cutlass/cutlass.h" - #include "cutlass/layout/layout.h" - #include - - #define DEBUG 0 - #define DEBUG1 0 - #define DEBUG2 0 - - using namespace cute; - - template - struct Flash_kernel_traits { - - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - using Element = elem_type; - static constexpr bool Has_cp_async = true; - #else - using Element = cutlass::half_t; - static constexpr bool Has_cp_async = false; - #endif - - using ElementAccum = float; - using index_t = int64_t; - - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 - using MMA_Atom_Arch = std::conditional_t< - std::is_same_v, - MMA_Atom, - MMA_Atom - >; - #else - using MMA_Atom_Arch = MMA_Atom; - #endif - - #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750 - using SmemCopyAtom = Copy_Atom; - using SmemCopyAtomTransposed = Copy_Atom; - #else - using SmemCopyAtom = Copy_Atom; - using SmemCopyAtomTransposed = Copy_Atom; - #endif - }; - - // If Share_Q_K_smem is true, that forces Is_Q_in_regs to be true - template > - struct Flash_fwd_kernel_traits : public Base { - static constexpr bool Has_cp_async = Base::Has_cp_async; - - // TODO - using Element = typename Base::Element; - using ElementKVPack = cute::uint16_t; - using ElementAccum = typename Base::ElementAccum; - using index_t = typename Base::index_t; - - using SmemCopyAtom = typename Base::SmemCopyAtom; - using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed; - - static constexpr bool Share_Q_K_smem = Share_Q_K_smem_; - static constexpr bool Is_Q_in_regs = Is_Q_in_regs_ || Share_Q_K_smem; - - static constexpr int quant_mode = quant_mode_; - static constexpr int group_size = group_size_; - static constexpr int full_len = 32; - static constexpr int num_bits = num_bits_; - static constexpr int pack_num = 16 / num_bits; - - static constexpr int residual_block_size = num_bits == 4 ? 128 : 256; - - // The number of threads. - static constexpr int kNWarps = kNWarps_; - static constexpr int kNThreads = kNWarps * 32; - - static constexpr int kBlockM = kBlockM_; - static constexpr int kBlockN = kBlockN_; - static constexpr int kBlockN_pack = num_bits == 4 ? 128 : 256; - static constexpr int kBlockN_residual = kBlockN_pack; - static constexpr int kBlockP = quant_mode == 1 ? kBlockN / pack_num : kBlockN; - static constexpr int kBlockP_new_pack = quant_mode == 1 ? kBlockN_pack / pack_num : kBlockN_pack; - static constexpr int kBlockK_params = quant_mode == 1 ? kBlockN / group_size : kBlockN; - static constexpr int kBlockK_params_new = quant_mode == 1 ? kBlockN_pack / group_size : kBlockN_pack; - static constexpr int kHeadDim = kHeadDim_; - static constexpr int kHeadDim_pack = kHeadDim / pack_num; - static constexpr int kHeadDim_k = quant_mode == 1 ? kHeadDim : kHeadDim_pack; - static constexpr int kHeadDim_k_params = quant_mode == 1 ? kHeadDim : kHeadDim / group_size; - static constexpr int kHeadDim_v_params = kHeadDim / group_size; - - static constexpr int k_pack_div = quant_mode == 1 ? pack_num : 1; - // static constexpr int k_params_div = quant_mode == 1 ? group_size : 1; - - static_assert(kHeadDim % 32 == 0); - static constexpr int kBlockKSmem = kHeadDim % 64 == 0 ? 64 : 32; - static constexpr int kBlockKGmem = kHeadDim % 128 == 0 ? 128 : (kHeadDim % 64 == 0 ? 64 : 32); - static constexpr int kSwizzle = kBlockKSmem == 32 ? 2 : 3; - - static constexpr int tile_paramsk_g = kBlockN / 32 * (kBlockN / group_size); // TODO: check - static constexpr int tile_paramsk_g_r = kBlockN_residual / 32 * (kBlockN_residual / group_size); // TODO: check - static constexpr int tile_paramsk_j = kBlockN / group_size; - static constexpr int tile_paramsk_m = kBlockN / kBlockN_pack; - static constexpr int tile_paramsk_k = kHeadDim / 16; - static constexpr int tile_paramsv_k = kBlockN / 16; - static constexpr int tile_paramsv_k_r = kBlockN_residual / 16; - - static constexpr int num_params = kBlockN_pack / group_size; // TODO: check 128 - - // - // Tiled MMA - // - - using TiledMma = TiledMMA< - typename Base::MMA_Atom_Arch, - Layout,_4,_1>>, - Tile, Int<128>, _16>>; - - using TiledMmaKV_i4 = TiledMMA< - typename Base::MMA_Atom_Arch, - Layout,_4,_1>>, - Tile, Int<32>, _16>>; - - using TiledMma_residual = TiledMmaKV_i4; - - // - // Shared memory layout - // - - // Q - using SmemLayoutAtomQ = decltype( - composition(Swizzle{}, - // This has to be kBlockKSmem, using kHeadDim gives wrong results for d=128 - Layout>, - Stride, _1>>{})); - using SmemLayoutQ = decltype(tile_to_shape( - SmemLayoutAtomQ{}, - Shape, Int>{})); - - // K - using SmemLayoutKSize = decltype( - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutAtomK = decltype( - make_layout(make_shape(Int<8>{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutKV = decltype(tile_to_shape( - SmemLayoutAtomK{}, - Shape, Int>{})); - using SmemLayoutAtomK_SW = decltype( - composition(Swizzle{}, - // TODO: 32 - Layout>, - Stride, _1>>{})); - using SmemLayoutKPack = decltype(tile_to_shape( - SmemLayoutAtomK_SW{}, - Shape, Int>{})); - using SmemLayoutKPacktransposed_ = decltype( - composition(SmemLayoutKPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - using SmemLayoutKPacktransposed = std::conditional_t< - quant_mode == 1, - SmemLayoutKPack, - SmemLayoutKPacktransposed_ - >; - using SmemLayoutKResidual = decltype( - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutKNewPack = decltype( - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutKNewPacktransposed_ = decltype( - composition(SmemLayoutKNewPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - using SmemLayoutKNewPacktransposed = std::conditional_t< - quant_mode == 1, - SmemLayoutKNewPack, - SmemLayoutKNewPacktransposed_ - >; - - // K params - // TODO: clear - using SmemLayoutKParams_channel = decltype( - composition(Swizzle<2, 2, 3>{}, - Layout, Int>, - Stride, Int>>{})); - // using SmemLayoutKParams_channel = decltype( - // make_layout(make_shape(Int{}, Int{}), - // make_stride(Int<1>{}, Int{}))); - - using SmemLayoutAtomKParams_group = decltype( - make_layout(make_shape(Int<32>{}, Int<1>{}), - make_stride(Int<1>{}, Int<1>{}))); - using SmemLayoutKParams_group = decltype(tile_to_shape( - SmemLayoutAtomKParams_group{}, - Shape, Int>{})); - using SmemLayoutKParams = std::conditional_t< - quant_mode == 1, - SmemLayoutKParams_channel, - SmemLayoutKParams_group - >; - - // V - using SmemLayoutVSize = decltype( - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutVtransposed = decltype( - composition(SmemLayoutKV{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{})); - - using SmemLayoutAtomV = decltype( - make_layout(make_shape(Int<8>{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutAtomV_SW = decltype( - composition(Swizzle<3, 3, 3>{}, - Layout>, - Stride, _1>>{})); - - using SmemLayoutVPack = decltype(tile_to_shape( - SmemLayoutAtomV_SW{}, - Shape, Int>{})); - using SmemLayoutVPacktransposed = decltype( - composition(SmemLayoutVPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - using SmemLayoutVPacktransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVPacktransposed{})); - - using SmemLayoutVResidual = decltype( - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutVResidualtransposed = decltype( - composition(SmemLayoutVResidual{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - using SmemLayoutVResidualtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVResidualtransposed{})); - - using SmemLayoutVNewPack = decltype( - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutVNewPacktransposed = decltype( - composition(SmemLayoutVNewPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - - // using SmemLayoutVNewPack = SmemLayoutKNewPack; - // using SmemLayoutVNewPacktransposed = SmemLayoutKNewPacktransposed_; - - // V params - using SmemLayoutAtomVParams = decltype( - make_layout(make_shape(Int<32>{}, Int<1>{}), - make_stride(Int<1>{}, Int<1>{}))); - using SmemLayoutVParams = decltype(tile_to_shape( - SmemLayoutAtomVParams{}, - Shape, Int>{})); - - // acc - using SmemLayoutAtomACC = decltype(composition( - Swizzle<3, 3, 3>{}, make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{})))); - using SmemLayoutAcc = decltype(tile_to_shape( - SmemLayoutAtomACC{}, - Shape, Int>{})); - using SmemLayoutAtomACC_residual = decltype(composition( - Swizzle<3, 3, 3>{}, make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{})))); - using SmemLayoutAcc_residual = decltype( - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using R2SCopyAtomAcc = Copy_Atom, Element>; - - // O - using SmemLayoutAtomO = decltype( - composition(Swizzle{}, - Layout, Int>, - Stride, _1>>{})); - using SmemLayoutO = decltype(tile_to_shape( - SmemLayoutAtomO{}, - Shape, Int>{})); - using SmemCopyAtomO = Copy_Atom; - using SmemCopyAtomOaccum = Copy_Atom; - - // Reduce tmp - // TODO: check 32 - using SmemLayoutReduce_tmp = decltype( - composition(Swizzle<1, 3, 3>{}, - Layout>, - Stride, _1>>{})); - - struct SharedStorage - { - array_aligned> smem_Q; - // array_aligned> smem_K; - array_aligned> smem_Kpack; - // array_aligned> smem_Kpack; - array_aligned<__half2, cosize_v> smem_Kparams; - // array_aligned> smem_V; - array_aligned> smem_Vpack; - // array_aligned> smem_Vpack; - array_aligned<__half2, cosize_v> smem_Vparams; - array_aligned> smem_acc; - }; - static constexpr int kSmemSize = int(sizeof(SharedStorage)); - - // - // Copy Atom, global to shared memory - // - - static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); - static_assert(kHeadDim % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad"); - static constexpr int kGmemThreadsPerRow = kBlockKSmem / kGmemElemsPerLoad; - static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow"); - using GmemLayoutAtom = Layout, Int>, - Stride, _1>>; // (16, 8) - - using Gmem_copy_struct = std::conditional_t< - Has_cp_async, - SM80_CP_ASYNC_CACHEGLOBAL, - DefaultCopy - >; - using GmemTiledCopyQKV = decltype( - make_tiled_copy(Copy_Atom{}, - GmemLayoutAtom{}, - Layout>{})); // Val layout, 8 vals per read - - using GmemTileCopyKV_Residual = decltype( - make_tiled_copy(Copy_Atom{}, - GmemLayoutAtom{}, - Layout>{})); // Val layout, 8 vals per read - - // TODO: check 32, 4 - // KV - using GmemTileCopyK_Pack = decltype( - make_tiled_copy(Copy_Atom{}, - make_layout(make_shape(_32{}, _4{}), make_stride(_4{}, _1{})), - Layout>{})); // Val layout, 8 vals per store - using GmemTileCopyV_Pack_2bit = decltype( - make_tiled_copy(Copy_Atom{}, - make_layout(make_shape(_64{}, _2{}), make_stride(_2{}, _1{})), - Layout>{})); // Val layout, 8 vals per store - using GmemTileCopyV_Pack_4bit = decltype( - make_tiled_copy(Copy_Atom{}, - make_layout(make_shape(_32{}, _4{}), make_stride(_4{}, _1{})), - Layout>{})); // Val layout, 8 vals per store - using GmemTileCopyV_Pack = std::conditional_t< - num_bits == 2, - GmemTileCopyV_Pack_2bit, - GmemTileCopyV_Pack_4bit - >; - using GmemTileCopyKV_NewPack = decltype( - make_tiled_copy(Copy_Atom{}, - make_layout(make_shape(_32{}, _4{}), make_stride(_4{}, _1{})), - Layout>{})); // Val layout, 8 vals per read - - // KV params - using GmemTileCopyKParams_BN128 = decltype( - make_tiled_copy(Copy_Atom, __half2>{}, - make_layout(make_shape(_1{}, _128{}), make_stride(_1{}, _1{})), - Layout>{})); // Val layout, 4 vals per store - using GmemTileCopyKParams_BN256 = decltype( - make_tiled_copy(Copy_Atom, __half2>{}, - make_layout(make_shape(_1{}, _128{}), make_stride(_1{}, _1{})), - Layout>{})); // Val layout, 4 vals per store - using GmemTileCopyKParams_channel = std::conditional_t< - kBlockN == 256, - GmemTileCopyKParams_BN128, // TODO: check - GmemTileCopyKParams_BN256 - >; - - - - using GmemTileCopyVParams_BN128 = decltype( - make_tiled_copy(Copy_Atom, __half2>{}, - make_layout(make_shape(_128{}, _1{}), make_stride(_1{}, _1{})), - Layout>{})); // Val layout, 4 vals per store - using GmemTileCopyVParams_BN256 = decltype( - make_tiled_copy(Copy_Atom, __half2>{}, - make_layout(make_shape(_128{}, _1{}), make_stride(_1{}, _1{})), - Layout>{})); // Val layout, 4 vals per store - using GmemTileCopyVParams = std::conditional_t< - kBlockN == 256, - GmemTileCopyVParams_BN256, - GmemTileCopyVParams_BN128 - >; - - using GmemTileCopyKParams = std::conditional_t< - quant_mode == 1, - GmemTileCopyKParams_channel, - GmemTileCopyVParams - >; - - // using GmemTileCopyKParams = GmemTileCopyVParams_BN128; - - - // O - using GmemTiledCopyO = decltype( - make_tiled_copy(Copy_Atom{}, - GmemLayoutAtom{}, - Layout>{})); // Val layout, 8 vals per store - using GmemLayoutAtomOaccum = std::conditional_t< - kBlockKSmem == 32, - Layout, // Thread layout, 8 threads per row - Stride< _8, _1>>, - Layout, // Thread layout, 16 threads per row - Stride< _16, _1>> - >; - using GmemTiledCopyOaccum = decltype( - make_tiled_copy(Copy_Atom{}, - GmemLayoutAtomOaccum{}, - Layout>{})); // Val layout, 4 vals per store - using GmemLayoutAtomRotcossin = GmemLayoutAtom; - using GmemTiledCopyRotcossin = decltype( - make_tiled_copy(Copy_Atom, Element>{}, - GmemLayoutAtomRotcossin{}, - Layout>{})); // Val layout, 4 vals per load - using GmemTiledCopyRotcossinCont = decltype( - make_tiled_copy(Copy_Atom{}, - GmemLayoutAtomRotcossin{}, - Layout>{})); // Val layout, 8 vals per load - - // - // Copy Atom, shared to register - // - - using S2RCopyOpQ = SM75_U32x2_LDSM_N; - using S2RCopyTraitsQ = Copy_Traits; - - using S2RCopyAtomK = Copy_Atom; - using S2RCopyAtomK_i4 = Copy_Atom; - - using S2RCopyAtomV = SmemCopyAtomTransposed; - using S2RCopyAtomV_i4 = Copy_Atom; - - using SmemCopyAtomTransposed_i4 = Copy_Atom; - // using SmemCopyAtomTransposed_i4 = Copy_Atom; - using SmemCopyAtomTransposed_residual = SmemCopyAtomTransposed_i4; - - using R2SCopyAtomPack = Copy_Atom; - - // using R2SCopyAtomPack = Copy_Atom; - // using ElementKVPack = cute::uint16_t; - - }; - - //////////////////////////////////////////////////////////////////////////////////////////////////// - - template > - struct Flash_qpack_traits : public Base { - using Element = typename Base::Element; - using ElementKVPack = cute::uint16_t; - using index_t = typename Base::index_t; - using SmemCopyAtom = typename Base::SmemCopyAtom; - using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed; - - static constexpr int quant_mode = quantmode_; - static constexpr int group_size = group_size_; - static constexpr int num_bits = num_bits_; - static constexpr int pack_num = 16 / num_bits; - - // The number of threads. - static constexpr int kNWarps = kNWarps_; - static constexpr int kNThreads = kNWarps * 32; - - static constexpr int kBlockN = kBlockN_; - static constexpr int kBlockN_pack = num_bits == 4 ? 128 : 256; - static constexpr int kBlockP = quant_mode == 1 ? kBlockN / pack_num : kBlockN; - static constexpr int kBlockK_params = quant_mode == 1 ? kBlockN / group_size : kBlockN; - static constexpr int kHeadDim = kHeadDim_; - static constexpr int kHeadDim_pack = kHeadDim / pack_num; // TODO - static constexpr int kHeadDim_k = quant_mode == 1 ? kHeadDim : kHeadDim_pack; - static constexpr int kHeadDim_k_params = quant_mode == 1 ? kHeadDim : kHeadDim / group_size; - static constexpr int kHeadDim_v_params = kHeadDim / group_size; - static_assert(kHeadDim % 32 == 0); - static constexpr int kBlockKSmem = kHeadDim % 64 == 0 ? 64 : 32; - static constexpr int kBlockKGmem = kHeadDim % 128 == 0 ? 128 : (kHeadDim % 64 == 0 ? 64 : 32); - static constexpr int kSwizzle = kBlockKSmem == 32 ? 2 : 3; - - static constexpr int tile_paramsk_g = kBlockN / 32 * (kBlockN / group_size); // TODO: check - static constexpr int tile_paramsk_j = kBlockN / group_size; - static constexpr int tile_paramsk_k = kHeadDim / 16; - static constexpr int tile_paramsv_k = kBlockN / 16; // TODO: check 128 - - static constexpr int num_params = kBlockN_pack / group_size; - - using TiledMma = TiledMMA< - typename Base::MMA_Atom_Arch, - Layout,_4,_1>>, - Tile, _128, _16>>; - - using TiledMmaK_i4 = TiledMMA< - typename Base::MMA_Atom_Arch, - Layout,_4,_1>>, - Tile, Int<32>, _16>>; - - using SmemLayoutAtomKV_SW = decltype( - composition(Swizzle{}, - // This has to be kBlockKSmem, using kHeadDim gives wrong results for d=128 - Layout>, - Stride, _1>>{})); - using SmemLayoutAtomK_tiled = decltype( - make_layout(make_shape(Int<8>{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutKV = decltype(tile_to_shape( - SmemLayoutAtomK_tiled{}, - Shape, Int>{})); - using SmemLayoutVtransposed = decltype( - composition(SmemLayoutKV{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{})); - - using SmemLayoutKPack = decltype( - make_layout(make_shape(Int{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutKPacktransposed_ = decltype( - composition(SmemLayoutKPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - using SmemLayoutKPacktransposed = std::conditional_t< - quant_mode == 1, - SmemLayoutKPack, - SmemLayoutKPacktransposed_ - >; - - using SmemLayoutAtomV = decltype( - make_layout(make_shape(Int<8>{}, Int{}), - make_stride(Int{}, Int<1>{}))); - using SmemLayoutVPack = decltype(tile_to_shape( - SmemLayoutAtomV{}, - Shape, Int>{})); - using SmemLayoutVPacktransposed = decltype( - composition(SmemLayoutVPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); - using SmemLayoutVPacktransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVPacktransposed{})); - - // TODO: 32 of x can be determined. - using SmemLayoutReduce_tmp = decltype( - make_layout(make_shape(Int<32>{}, Int<32>{}), - make_stride(Int<32>{}, Int<1>{}))); - - using R2SCopyAtom = Copy_Atom; - using R2SCopyAtomPack = Copy_Atom; - - struct SharedStorage - { - array_aligned> smem_K; - array_aligned> smem_V; - array_aligned> smem_Kpack; - array_aligned> smem_Vpack; - array_aligned> smem_reduce_tmp; - }; - static constexpr int kSmemSize = int(sizeof(SharedStorage)); - - static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); - static_assert(kHeadDim % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad"); - static constexpr int kGmemThreadsPerRow = kBlockKSmem / kGmemElemsPerLoad; - static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow"); - using GmemLayoutAtom = Layout, Int>, - Stride, _1>>; - - // We use CACHEGLOBAL instead of CACHEALWAYS for both Q and K/V, since we won't be reading - // from the same address by the same threadblock. This is slightly faster. - using GmemTiledCopyQKV = decltype( - make_tiled_copy(Copy_Atom, Element>{}, - GmemLayoutAtom{}, - Layout>{})); // Val layout, 8 vals per read - using GmemTileCopyK_Pack = decltype( - make_tiled_copy(Copy_Atom{}, - make_layout(make_shape(_32{}, _4{}), make_stride(_4{}, _1{})), - Layout>{})); // Val layout, 8 vals per store - using GmemTileCopyV_Pack = decltype( - make_tiled_copy(Copy_Atom{}, - make_layout(make_shape(_64{}, _2{}), make_stride(_2{}, _1{})), - Layout>{})); // Val layout, 8 vals per store - }; - - //////////////////////////////////////////////////////////////////////////////////////////////////// \ No newline at end of file +#pragma once + +#include "cute/tensor.hpp" + +#include "cutlass/cutlass.h" +#include "cutlass/layout/layout.h" +#include + +#define DEBUG 0 +#define DEBUG1 0 +#define DEBUG2 0 + +using namespace cute; + +template +struct Flash_kernel_traits { + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + using Element = elem_type; + static constexpr bool Has_cp_async = true; +#else + using Element = cutlass::half_t; + static constexpr bool Has_cp_async = false; +#endif + + using ElementAccum = float; + using index_t = int64_t; + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + using MMA_Atom_Arch = std::conditional_t< + std::is_same_v, + MMA_Atom, + MMA_Atom + >; +#else + using MMA_Atom_Arch = MMA_Atom; +#endif + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750 + using SmemCopyAtom = Copy_Atom; + using SmemCopyAtomTransposed = Copy_Atom; +#else + using SmemCopyAtom = Copy_Atom; + using SmemCopyAtomTransposed = Copy_Atom; +#endif +}; + +// If Share_Q_K_smem is true, that forces Is_Q_in_regs to be true +template > +struct Flash_fwd_kernel_traits : public Base { + static constexpr bool Has_cp_async = Base::Has_cp_async; + + // TODO + using Element = typename Base::Element; + using ElementKVPack = cute::uint16_t; + using ElementAccum = typename Base::ElementAccum; + using index_t = typename Base::index_t; + + using SmemCopyAtom = typename Base::SmemCopyAtom; + using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed; + + static constexpr bool Share_Q_K_smem = Share_Q_K_smem_; + static constexpr bool Is_Q_in_regs = Is_Q_in_regs_ || Share_Q_K_smem; + + static constexpr int quant_mode = quant_mode_; + static constexpr int group_size = group_size_; + static constexpr int num_bits = num_bits_; + static constexpr int pack_num = 16 / num_bits; + + static constexpr int residual_block_size = num_bits == 4 ? 128 : 256; + + // The number of threads. + static constexpr int kNWarps = kNWarps_; + static constexpr int kNThreads = kNWarps * 32; + + static constexpr int kBlockM = kBlockM_; + static constexpr int kBlockN = kBlockN_; + static constexpr int kBlockN_pack = num_bits == 4 ? 128 : 256; + static constexpr int kBlockN_residual = kBlockN_pack; + static constexpr int kBlockP = quant_mode == 1 ? kBlockN / pack_num : kBlockN; + static constexpr int kBlockP_new_pack = quant_mode == 1 ? kBlockN_pack / pack_num : kBlockN_pack; + static constexpr int kBlockK_params = quant_mode == 1 ? kBlockN / group_size : kBlockN; + static constexpr int kBlockK_params_new = quant_mode == 1 ? kBlockN_pack / group_size : kBlockN_pack; + static constexpr int kHeadDim = kHeadDim_; + static constexpr int kHeadDim_pack = kHeadDim / pack_num; + static constexpr int kHeadDim_k = quant_mode == 1 ? kHeadDim : kHeadDim_pack; + static constexpr int kHeadDim_k_params = quant_mode == 1 ? kHeadDim : kHeadDim / group_size; + static constexpr int kHeadDim_v_params = kHeadDim / group_size; + + static constexpr int k_pack_div = quant_mode == 1 ? pack_num : 1; + // static constexpr int k_params_div = quant_mode == 1 ? group_size : 1; + + static_assert(kHeadDim % 32 == 0); + static constexpr int kBlockKSmem = kHeadDim % 64 == 0 ? 64 : 32; + static constexpr int kBlockKGmem = kHeadDim % 128 == 0 ? 128 : (kHeadDim % 64 == 0 ? 64 : 32); + static constexpr int kSwizzle = kBlockKSmem == 32 ? 2 : 3; + + static constexpr int tile_paramsk_g = kBlockN / 32 * (kBlockN / group_size); // TODO: check + static constexpr int tile_paramsk_g_r = kBlockN_residual / 32 * (kBlockN_residual / group_size); // TODO: check + static constexpr int tile_paramsk_j = kBlockN / group_size; + static constexpr int tile_paramsk_m = kBlockN / kBlockN_pack; + static constexpr int tile_paramsk_k = kHeadDim / 16; + static constexpr int tile_paramsv_k = kBlockN / 16; + static constexpr int tile_paramsv_k_r = kBlockN_residual / 16; + + static constexpr int num_params = kBlockN_pack / group_size; // TODO: check 128 + + // + // Tiled MMA + // + + using TiledMma = TiledMMA< + typename Base::MMA_Atom_Arch, + Layout,_4,_1>>, + Tile, Int<128>, _16>>; + + using TiledMmaKV_i4 = TiledMMA< + typename Base::MMA_Atom_Arch, + Layout,_4,_1>>, + Tile, Int<32>, _16>>; + + using TiledMma_residual = TiledMmaKV_i4; + + // + // Shared memory layout + // + + // Q + using SmemLayoutAtomQ = decltype( + composition(Swizzle{}, + // This has to be kBlockKSmem, using kHeadDim gives wrong results for d=128 + Layout>, + Stride, _1>>{})); + using SmemLayoutQ = decltype(tile_to_shape( + SmemLayoutAtomQ{}, + Shape, Int>{})); + + // K + using SmemLayoutKSize = decltype( + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutAtomK = decltype( + make_layout(make_shape(Int<8>{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutKV = decltype(tile_to_shape( + SmemLayoutAtomK{}, + Shape, Int>{})); + using SmemLayoutAtomK_SW = decltype( + composition(Swizzle{}, + // TODO: 32 + Layout>, + Stride, _1>>{})); + using SmemLayoutKPack = decltype(tile_to_shape( + SmemLayoutAtomK_SW{}, + Shape, Int>{})); + using SmemLayoutKPacktransposed_ = decltype( + composition(SmemLayoutKPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutKPacktransposed = std::conditional_t< + quant_mode == 1, + SmemLayoutKPack, + SmemLayoutKPacktransposed_ + >; + using SmemLayoutKResidual = decltype( + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutKNewPack = decltype( + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutKNewPacktransposed_ = decltype( + composition(SmemLayoutKNewPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutKNewPacktransposed = std::conditional_t< + quant_mode == 1, + SmemLayoutKNewPack, + SmemLayoutKNewPacktransposed_ + >; + + // K params + // TODO: clear + using SmemLayoutKParams_channel = decltype( + composition(Swizzle<2, 2, 3>{}, + Layout, Int>, + Stride, Int>>{})); + // using SmemLayoutKParams_channel = decltype( + // make_layout(make_shape(Int{}, Int{}), + // make_stride(Int<1>{}, Int{}))); + + using SmemLayoutAtomKParams_group = decltype( + make_layout(make_shape(Int<32>{}, Int<1>{}), + make_stride(Int<1>{}, Int<1>{}))); + using SmemLayoutKParams_group = decltype(tile_to_shape( + SmemLayoutAtomKParams_group{}, + Shape, Int>{})); + using SmemLayoutKParams = std::conditional_t< + quant_mode == 1, + SmemLayoutKParams_channel, + SmemLayoutKParams_group + >; + + // V + using SmemLayoutVSize = decltype( + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutVtransposed = decltype( + composition(SmemLayoutKV{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{})); + + using SmemLayoutAtomV = decltype( + make_layout(make_shape(Int<8>{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutAtomV_SW = decltype( + composition(Swizzle<3, 3, 3>{}, + Layout>, + Stride, _1>>{})); + + using SmemLayoutVPack = decltype(tile_to_shape( + SmemLayoutAtomV_SW{}, + Shape, Int>{})); + using SmemLayoutVPacktransposed = decltype( + composition(SmemLayoutVPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutVPacktransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVPacktransposed{})); + + using SmemLayoutVResidual = decltype( + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutVResidualtransposed = decltype( + composition(SmemLayoutVResidual{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutVResidualtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVResidualtransposed{})); + + using SmemLayoutVNewPack = decltype( + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutVNewPacktransposed = decltype( + composition(SmemLayoutVNewPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + + // using SmemLayoutVNewPack = SmemLayoutKNewPack; + // using SmemLayoutVNewPacktransposed = SmemLayoutKNewPacktransposed_; + + // V params + using SmemLayoutAtomVParams = decltype( + make_layout(make_shape(Int<32>{}, Int<1>{}), + make_stride(Int<1>{}, Int<1>{}))); + using SmemLayoutVParams = decltype(tile_to_shape( + SmemLayoutAtomVParams{}, + Shape, Int>{})); + + // acc + using SmemLayoutAtomACC = decltype(composition( + Swizzle<3, 3, 3>{}, make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})))); + using SmemLayoutAcc = decltype(tile_to_shape( + SmemLayoutAtomACC{}, + Shape, Int>{})); + using SmemLayoutAtomACC_residual = decltype(composition( + Swizzle<3, 3, 3>{}, make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{})))); + using SmemLayoutAcc_residual = decltype( + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using R2SCopyAtomAcc = Copy_Atom, Element>; + + // O + using SmemLayoutAtomO = decltype( + composition(Swizzle{}, + Layout, Int>, + Stride, _1>>{})); + using SmemLayoutO = decltype(tile_to_shape( + SmemLayoutAtomO{}, + Shape, Int>{})); + using SmemCopyAtomO = Copy_Atom; + using SmemCopyAtomOaccum = Copy_Atom; + + // Reduce tmp + // TODO: check 32 + using SmemLayoutReduce_tmp = decltype( + composition(Swizzle<1, 3, 3>{}, + Layout>, + Stride, _1>>{})); + + struct SharedStorage + { + array_aligned> smem_Q; + // array_aligned> smem_K; + // array_aligned> smem_Kpack; + array_aligned> smem_Kpack; + array_aligned<__half2, cosize_v> smem_Kparams; + // array_aligned> smem_V; + // array_aligned> smem_Vpack; + array_aligned> smem_Vpack; + array_aligned<__half2, cosize_v> smem_Vparams; + array_aligned> smem_acc; + }; + static constexpr int kSmemSize = int(sizeof(SharedStorage)); + + struct SharedStorage_residual + { + array_aligned> smem_Q; + + array_aligned> smem_Kpack; + array_aligned> smem_Vpack; + array_aligned> smem_acc; + array_aligned<__half2, cosize_v> smem_Kparams; + }; + static constexpr int kSmemSize_res = int(sizeof(SharedStorage_residual)); + + // + // Copy Atom, global to shared memory + // + + static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); + static_assert(kHeadDim % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad"); + static constexpr int kGmemThreadsPerRow = kBlockKSmem / kGmemElemsPerLoad; + static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow"); + using GmemLayoutAtom = Layout, Int>, + Stride, _1>>; // (16, 8) + + using Gmem_copy_struct = std::conditional_t< + Has_cp_async, + SM80_CP_ASYNC_CACHEGLOBAL, + DefaultCopy + >; + using GmemTiledCopyQKV = decltype( + make_tiled_copy(Copy_Atom{}, + GmemLayoutAtom{}, + Layout>{})); // Val layout, 8 vals per read + + using GmemTileCopyKV_Residual = decltype( + make_tiled_copy(Copy_Atom{}, + GmemLayoutAtom{}, + Layout>{})); // Val layout, 8 vals per read + + // TODO: check 32, 4 + // KV + using GmemTileCopyK_Pack = decltype( + make_tiled_copy(Copy_Atom{}, + make_layout(make_shape(_32{}, _4{}), make_stride(_4{}, _1{})), + Layout>{})); // Val layout, 8 vals per store + using GmemTileCopyV_Pack_2bit = decltype( + make_tiled_copy(Copy_Atom{}, + make_layout(make_shape(_64{}, _2{}), make_stride(_2{}, _1{})), + Layout>{})); // Val layout, 8 vals per store + using GmemTileCopyV_Pack_4bit = decltype( + make_tiled_copy(Copy_Atom{}, + make_layout(make_shape(_32{}, _4{}), make_stride(_4{}, _1{})), + Layout>{})); // Val layout, 8 vals per store + using GmemTileCopyV_Pack = std::conditional_t< + num_bits == 2, + GmemTileCopyV_Pack_2bit, + GmemTileCopyV_Pack_4bit + >; + using GmemTileCopyKV_NewPack = decltype( + make_tiled_copy(Copy_Atom{}, + make_layout(make_shape(_32{}, _4{}), make_stride(_4{}, _1{})), + Layout>{})); // Val layout, 8 vals per read + + // KV params + using GmemTileCopyKParams_BN128 = decltype( + make_tiled_copy(Copy_Atom, __half2>{}, + make_layout(make_shape(_1{}, _128{}), make_stride(_1{}, _1{})), + Layout>{})); // Val layout, 4 vals per store + using GmemTileCopyKParams_BN256 = decltype( + make_tiled_copy(Copy_Atom, __half2>{}, + make_layout(make_shape(_1{}, _128{}), make_stride(_1{}, _1{})), + Layout>{})); // Val layout, 4 vals per store + using GmemTileCopyKParams_channel = std::conditional_t< + kBlockN == 256, + GmemTileCopyKParams_BN128, // TODO: check + GmemTileCopyKParams_BN128 + >; + + using GmemTileCopyVParams_BN128 = decltype( + make_tiled_copy(Copy_Atom, __half2>{}, + make_layout(make_shape(_128{}, _1{}), make_stride(_1{}, _1{})), + Layout>{})); // Val layout, 4 vals per store + using GmemTileCopyVParams_BN256 = decltype( + make_tiled_copy(Copy_Atom, __half2>{}, + make_layout(make_shape(_128{}, _1{}), make_stride(_1{}, _1{})), + Layout>{})); // Val layout, 4 vals per store + using GmemTileCopyVParams = std::conditional_t< + kBlockN == 256, + GmemTileCopyVParams_BN256, + GmemTileCopyVParams_BN128 + >; + + using GmemTileCopyKParams = std::conditional_t< + quant_mode == 1, + GmemTileCopyKParams_channel, + GmemTileCopyVParams + >; + + // using GmemTileCopyKParams = GmemTileCopyVParams_BN128; + + + // O + using GmemTiledCopyO = decltype( + make_tiled_copy(Copy_Atom{}, + GmemLayoutAtom{}, + Layout>{})); // Val layout, 8 vals per store + using GmemLayoutAtomOaccum = std::conditional_t< + kBlockKSmem == 32, + Layout, // Thread layout, 8 threads per row + Stride< _8, _1>>, + Layout, // Thread layout, 16 threads per row + Stride< _16, _1>> + >; + using GmemTiledCopyOaccum = decltype( + make_tiled_copy(Copy_Atom{}, + GmemLayoutAtomOaccum{}, + Layout>{})); // Val layout, 4 vals per store + using GmemLayoutAtomRotcossin = GmemLayoutAtom; + using GmemTiledCopyRotcossin = decltype( + make_tiled_copy(Copy_Atom, Element>{}, + GmemLayoutAtomRotcossin{}, + Layout>{})); // Val layout, 4 vals per load + using GmemTiledCopyRotcossinCont = decltype( + make_tiled_copy(Copy_Atom{}, + GmemLayoutAtomRotcossin{}, + Layout>{})); // Val layout, 8 vals per load + + // + // Copy Atom, shared to register + // + + using S2RCopyOpQ = SM75_U32x2_LDSM_N; + using S2RCopyTraitsQ = Copy_Traits; + + using S2RCopyAtomK = Copy_Atom; + using S2RCopyAtomK_i4 = Copy_Atom; + + using S2RCopyAtomV = SmemCopyAtomTransposed; + using S2RCopyAtomV_i4 = Copy_Atom; + + using SmemCopyAtomTransposed_i4 = Copy_Atom; + // using SmemCopyAtomTransposed_i4 = Copy_Atom; + using SmemCopyAtomTransposed_residual = SmemCopyAtomTransposed_i4; + + using R2SCopyAtomPack = Copy_Atom; + + // using R2SCopyAtomPack = Copy_Atom; + // using ElementKVPack = cute::uint16_t; + +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +template > +struct Flash_qpack_traits : public Base { + using Element = typename Base::Element; + using ElementKVPack = cute::uint16_t; + using index_t = typename Base::index_t; + using SmemCopyAtom = typename Base::SmemCopyAtom; + using SmemCopyAtomTransposed = typename Base::SmemCopyAtomTransposed; + + static constexpr int quant_mode = quantmode_; + static constexpr int group_size = group_size_; + static constexpr int num_bits = num_bits_; + static constexpr int pack_num = 16 / num_bits; + + // The number of threads. + static constexpr int kNWarps = kNWarps_; + static constexpr int kNThreads = kNWarps * 32; + + static constexpr int kBlockN = kBlockN_; + static constexpr int kBlockN_pack = num_bits == 4 ? 128 : 256; + static constexpr int kBlockP = quant_mode == 1 ? kBlockN / pack_num : kBlockN; + static constexpr int kBlockK_params = quant_mode == 1 ? kBlockN / group_size : kBlockN; + static constexpr int kHeadDim = kHeadDim_; + static constexpr int kHeadDim_pack = kHeadDim / pack_num; // TODO + static constexpr int kHeadDim_k = quant_mode == 1 ? kHeadDim : kHeadDim_pack; + static constexpr int kHeadDim_k_params = quant_mode == 1 ? kHeadDim : kHeadDim / group_size; + static constexpr int kHeadDim_v_params = kHeadDim / group_size; + static_assert(kHeadDim % 32 == 0); + static constexpr int kBlockKSmem = kHeadDim % 64 == 0 ? 64 : 32; + static constexpr int kBlockKGmem = kHeadDim % 128 == 0 ? 128 : (kHeadDim % 64 == 0 ? 64 : 32); + static constexpr int kSwizzle = kBlockKSmem == 32 ? 2 : 3; + + static constexpr int tile_paramsk_g = kBlockN / 32 * (kBlockN / group_size); // TODO: check + static constexpr int tile_paramsk_j = kBlockN / group_size; + static constexpr int tile_paramsk_k = kHeadDim / 16; + static constexpr int tile_paramsv_k = kBlockN / 16; // TODO: check 128 + + static constexpr int num_params = kBlockN_pack / group_size; + + using TiledMma = TiledMMA< + typename Base::MMA_Atom_Arch, + Layout,_4,_1>>, + Tile, _128, _16>>; + + using TiledMmaK_i4 = TiledMMA< + typename Base::MMA_Atom_Arch, + Layout,_4,_1>>, + Tile, Int<32>, _16>>; + + using SmemLayoutAtomKV_SW = decltype( + composition(Swizzle{}, + // This has to be kBlockKSmem, using kHeadDim gives wrong results for d=128 + Layout>, + Stride, _1>>{})); + using SmemLayoutAtomK_tiled = decltype( + make_layout(make_shape(Int<8>{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutKV = decltype(tile_to_shape( + SmemLayoutAtomK_tiled{}, + Shape, Int>{})); + using SmemLayoutVtransposed = decltype( + composition(SmemLayoutKV{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutVtransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVtransposed{})); + + using SmemLayoutKPack = decltype( + make_layout(make_shape(Int{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutKPacktransposed_ = decltype( + composition(SmemLayoutKPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutKPacktransposed = std::conditional_t< + quant_mode == 1, + SmemLayoutKPack, + SmemLayoutKPacktransposed_ + >; + + using SmemLayoutAtomV = decltype( + make_layout(make_shape(Int<8>{}, Int{}), + make_stride(Int{}, Int<1>{}))); + using SmemLayoutVPack = decltype(tile_to_shape( + SmemLayoutAtomV{}, + Shape, Int>{})); + using SmemLayoutVPacktransposed = decltype( + composition(SmemLayoutVPack{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using SmemLayoutVPacktransposedNoSwizzle = decltype(get_nonswizzle_portion(SmemLayoutVPacktransposed{})); + + // TODO: 32 of x can be determined. + using SmemLayoutReduce_tmp = decltype( + make_layout(make_shape(Int<32>{}, Int<32>{}), + make_stride(Int<32>{}, Int<1>{}))); + + using R2SCopyAtom = Copy_Atom; + using R2SCopyAtomPack = Copy_Atom; + + struct SharedStorage + { + array_aligned> smem_K; + array_aligned> smem_V; + array_aligned> smem_Kpack; + array_aligned> smem_Vpack; + array_aligned> smem_reduce_tmp; + }; + static constexpr int kSmemSize = int(sizeof(SharedStorage)); + + static constexpr int kGmemElemsPerLoad = sizeof(cute::uint128_t) / sizeof(Element); + static_assert(kHeadDim % kGmemElemsPerLoad == 0, "kHeadDim must be a multiple of kGmemElemsPerLoad"); + static constexpr int kGmemThreadsPerRow = kBlockKSmem / kGmemElemsPerLoad; + static_assert(kNThreads % kGmemThreadsPerRow == 0, "kNThreads must be a multiple of kGmemThreadsPerRow"); + using GmemLayoutAtom = Layout, Int>, + Stride, _1>>; + + // We use CACHEGLOBAL instead of CACHEALWAYS for both Q and K/V, since we won't be reading + // from the same address by the same threadblock. This is slightly faster. + using GmemTiledCopyQKV = decltype( + make_tiled_copy(Copy_Atom, Element>{}, + GmemLayoutAtom{}, + Layout>{})); // Val layout, 8 vals per read + using GmemTileCopyK_Pack = decltype( + make_tiled_copy(Copy_Atom{}, + make_layout(make_shape(_32{}, _4{}), make_stride(_4{}, _1{})), + Layout>{})); // Val layout, 8 vals per store + using GmemTileCopyV_Pack = decltype( + make_tiled_copy(Copy_Atom{}, + make_layout(make_shape(_64{}, _2{}), make_stride(_2{}, _1{})), + Layout>{})); // Val layout, 8 vals per store +}; + +//////////////////////////////////////////////////////////////////////////////////////////////////// \ No newline at end of file diff --git a/csrc/bit_decode/src/test_single_residual.cu b/csrc/bit_decode/src/test_single_residual.cu new file mode 100644 index 0000000..cba1d21 --- /dev/null +++ b/csrc/bit_decode/src/test_single_residual.cu @@ -0,0 +1,306 @@ +#include "flash_api.h" +#include +#include + +torch::Tensor single_mha(torch::Tensor& q, torch::Tensor& k, torch::Tensor& v, int head_dim) { + // Optional RoPE + const float sm_scale = 1.f / std::sqrt(float(head_dim)); + auto scaled_q = q * sm_scale; + + auto scores = torch::einsum("bthd,bshd->bhts", {scaled_q, k}); + auto attention = torch::softmax(scores, -1).to(v.dtype()); + auto output = torch::einsum("bhts,bshd->bthd", {attention, v}); + return output; +} + +template +void TestDecodingKernelCorrectness(const int bs, int seqlen_kv, const std::string& quant_mode, const int group_size) { + // Set the random seed for reproducibility + torch::manual_seed(42); + + const int seqlen_q = 1; + const int pack_nums = 16 / num_bits; + const int residual_block_size = num_bits == 4 ? 128 : 256; + int residual_len = seqlen_kv % residual_block_size == 0 ? residual_block_size : seqlen_kv % residual_block_size; + seqlen_kv = seqlen_kv - residual_len; + + printf("\n\n################## Round 0 ##################\n\n"); + + torch::Tensor Q_host = torch::rand({bs, seqlen_q, num_heads, head_dim}, torch::dtype(torch::kHalf)); + + torch::Tensor K_host = torch::randn({bs, seqlen_kv, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + torch::Tensor V_host = torch::randn({bs, seqlen_kv, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + + torch::Tensor Q_device = Q_host.to(torch::kCUDA); + torch::Tensor K_device = K_host.to(torch::kCUDA); + torch::Tensor V_device = V_host.to(torch::kCUDA); + + at::Tensor k_pack, k_params, v_pack, v_params, k_pack_new, k_params_new, v_pack_new, v_params_new; + if (quant_mode == "k-channel") { + k_pack = torch::empty({bs, seqlen_kv / pack_nums, num_heads_kv, head_dim}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + k_params = torch::empty({bs, seqlen_kv / group_size, num_heads_kv, head_dim}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + k_pack_new = torch::empty({bs, residual_block_size / pack_nums, num_heads_kv, k_pack.size(-1)}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + k_params_new = torch::empty({bs, residual_block_size / group_size, num_heads_kv, k_params.size(-1)}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + } else { + k_pack = torch::empty({bs, seqlen_kv, num_heads_kv, head_dim / pack_nums}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + k_params = torch::empty({bs, head_dim / group_size, num_heads_kv, seqlen_kv}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + k_pack_new = torch::empty({bs, residual_block_size, num_heads_kv, k_pack.size(-1)}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + k_params_new = torch::empty({bs, k_params.size(1), num_heads_kv, residual_block_size}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + } + + v_pack = torch::empty({bs, seqlen_kv, num_heads_kv, head_dim / pack_nums}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + v_params = torch::empty({bs, head_dim / group_size, num_heads_kv, seqlen_kv}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + v_pack_new = torch::empty({bs, residual_block_size, num_heads_kv, v_pack.size(-1)}, torch::dtype(torch::kUInt16)).to(torch::kCUDA); + v_params_new = torch::empty({bs, v_params.size(1), num_heads_kv, residual_block_size}, torch::dtype(torch::kFloat32)).to(torch::kCUDA); + + // Convert K, V to unpadded format + torch::Tensor K_unpad = K_device.reshape({bs * seqlen_kv, num_heads_kv, head_dim}); + torch::Tensor V_unpad = V_device.reshape({bs * seqlen_kv, num_heads_kv, head_dim}); + + auto cu_seqlens_k = torch::arange(0, (bs + 1) * seqlen_kv, seqlen_kv, torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA)); + // auto cu_seqlens_k = std::nullopt; + + std::optional opt_block_table = std::nullopt; + + kvcache_qpack( + K_unpad, k_pack, k_params, + V_unpad, v_pack, v_params, + opt_block_table, + cu_seqlens_k, + seqlen_kv, + quant_mode, + group_size + ); + + at::Tensor K_residual_host, V_residual_host, K_new_host, V_new_host, K_new_device, V_new_device, seqlens_k; + int new_lens = 0; + + seqlens_k = torch::full({bs}, seqlen_kv, torch::dtype(torch::kInt32).device(torch::kCUDA)); + + K_residual_host = torch::zeros({bs, residual_block_size, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + V_residual_host = torch::zeros({bs, residual_block_size, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + + K_new_host = torch::randn({bs, residual_len, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + V_new_host = torch::randn({bs, residual_len, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + + new_lens = residual_len; + + // Copy data from K_new_host to K_residual_host + K_residual_host.slice(1, 0, residual_len).copy_(K_new_host); + V_residual_host.slice(1, 0, residual_len).copy_(V_new_host); + + K_new_device = K_residual_host.to(torch::kCUDA); + V_new_device = V_residual_host.to(torch::kCUDA); + + // mha_fwd_kvcache + const float sm_scale = 1 / std::sqrt(float(head_dim)); + std::optional opt_K_new_device = std::make_optional(K_new_device); + std::optional opt_V_new_device = std::make_optional(V_new_device); + std::optional opt_seqlens_k = std::make_optional(seqlens_k); + + auto [out, k_pack_new_1, k_params_new_1, v_pack_new_1, v_params_new_1] + = mha_fwd_kvcache(Q_device, + k_pack, k_params, + v_pack, v_params, + opt_K_new_device, opt_V_new_device, opt_seqlens_k, + k_pack_new, k_params_new, v_pack_new, v_params_new, + opt_block_table, + sm_scale, + quant_mode, + group_size, + residual_block_size, + new_lens); + + torch::Tensor out_cpu = out.to(torch::kCPU); + + // CPU reference + torch::Tensor K_host_cat, V_host_cat; + + K_host_cat = torch::cat({K_host, K_new_host}, 1); + V_host_cat = torch::cat({V_host, V_new_host}, 1); + + torch::Tensor out_ref = single_mha(Q_host, K_host_cat, V_host_cat, head_dim); + + // Compute the difference + torch::Tensor diff = out_cpu - out_ref; + float mean_absolute_error = diff.abs().mean().item(); + float mean_squared_error = diff.pow(2).mean().item(); + float max_error = diff.abs().max().item(); + + printf("\nnum_bits: %d num_heads_kv: %d seqlen_kv: %d head_dim: %d Quant_mode: %s, Group_size: %d\n", num_bits, num_heads_kv, seqlen_kv + new_lens, head_dim, quant_mode.c_str(), group_size); + if (mean_absolute_error < 1e-1 && mean_squared_error < 1e-1) { + printf("test pass ! \n"); + printf("max_error: %f, mean_absolute_error: %f, mean_squared_error: %f\n", max_error, mean_absolute_error, mean_squared_error); + } else { + printf("test fail ! \n"); + printf("max_error: %f, mean_absolute_error: %f, mean_squared_error: %f\n", max_error, mean_absolute_error, mean_squared_error); + } + + printf("\nFirst head output (out_cpu[0,0,0,:]):\n"); + auto out_cpu_accessor = out_cpu.index({0,0,1}).data_ptr(); + for (int i = 0; i < head_dim; i++) { + printf("%.6f ", static_cast(out_cpu_accessor[i])); + } + printf("\n\nFirst head output (out_ref[0,0,0,:]):\n"); + auto out_ref_accessor = out_ref.index({0,0,1}).data_ptr(); + for (int i = 0; i < head_dim; i++) { + printf("%.6f ", static_cast(out_ref_accessor[i])); + } + printf("\n\n"); + + // + // Next round + // + + printf("\n\n################## Round 1 ##################\n\n"); + + auto seqlen_new = 1; + + K_new_host = torch::randn({bs, seqlen_new, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + V_new_host = torch::randn({bs, seqlen_new, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + + if (new_lens == residual_block_size) { + k_pack = torch::cat({k_pack, k_pack_new}, 1); + k_params = torch::cat({k_params, k_params_new}, 1); + v_pack = torch::cat({v_pack, v_pack_new}, 1); + v_params = torch::cat({v_params, v_params_new}, -1); + seqlens_k = torch::full({bs}, seqlen_kv + new_lens, torch::dtype(torch::kInt32).device(torch::kCUDA)); + + K_residual_host = torch::zeros({bs, residual_block_size, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + V_residual_host = torch::zeros({bs, residual_block_size, num_heads_kv, head_dim}, torch::dtype(torch::kHalf)); + K_residual_host.slice(1, 0, 0 + seqlen_new).copy_(K_new_host); + V_residual_host.slice(1, 0, 0 + seqlen_new).copy_(V_new_host); + } else { + seqlens_k = torch::full({bs}, seqlen_kv, torch::dtype(torch::kInt32).device(torch::kCUDA)); + + K_residual_host.slice(1, new_lens, new_lens + seqlen_new).copy_(K_new_host); + V_residual_host.slice(1, new_lens, new_lens + seqlen_new).copy_(V_new_host); + } + + seqlen_kv = seqlen_kv + new_lens + seqlen_new; + residual_len = seqlen_kv % residual_block_size == 0 ? residual_block_size : seqlen_kv % residual_block_size; + new_lens = residual_len; + + K_new_device = K_residual_host.to(torch::kCUDA); + V_new_device = V_residual_host.to(torch::kCUDA); + + std::optional opt_K_new_device_1 = std::make_optional(K_new_device); + std::optional opt_V_new_device_1 = std::make_optional(V_new_device); + std::optional opt_seqlens_k_1 = std::make_optional(seqlens_k); + + auto [out_2, k_pack_new_2, k_params_new_2, v_pack_new_2, v_params_new_2] + = mha_fwd_kvcache(Q_device, + k_pack, k_params, + v_pack, v_params, + opt_K_new_device_1, opt_V_new_device_1, opt_seqlens_k_1, + k_pack_new, k_params_new, v_pack_new, v_params_new, + opt_block_table, + sm_scale, + quant_mode, + group_size, + residual_block_size, + new_lens); + + torch::Tensor out_2_cpu = out_2.to(torch::kCPU); + + K_host_cat = torch::cat({K_host_cat, K_new_host}, 1); + V_host_cat = torch::cat({V_host_cat, V_new_host}, 1); + + torch::Tensor out_ref_2 = single_mha(Q_host, K_host_cat, V_host_cat, head_dim); + + diff = out_2_cpu - out_ref_2; + mean_absolute_error = diff.abs().mean().item(); + mean_squared_error = diff.pow(2).mean().item(); + max_error = diff.abs().max().item(); + + printf("\nnum_bits: %d num_heads_kv: %d seqlen_kv: %d head_dim: %d Quant_mode: %s, Group_size: %d\n", num_bits, num_heads_kv, seqlen_kv, head_dim, quant_mode.c_str(), group_size); + if (mean_absolute_error < 1e-1 && mean_squared_error < 1e-1) { + printf("test pass ! \n"); + printf("max_error: %f, mean_absolute_error: %f, mean_squared_error: %f\n", max_error, mean_absolute_error, mean_squared_error); + } else { + printf("test fail ! \n"); + printf("max_error: %f, mean_absolute_error: %f, mean_squared_error: %f\n", max_error, mean_absolute_error, mean_squared_error); + } + + printf("\nFirst head output (out_cpu[0,0,0,:]):\n"); + out_cpu_accessor = out_2_cpu.index({0,0,0}).data_ptr(); + for (int i = 0; i < head_dim; i++) { + printf("%.6f ", static_cast(out_cpu_accessor[i])); + } + printf("\n\nFirst head output (out_ref[0,0,0,:]):\n"); + out_ref_accessor = out_ref_2.index({0,0,0}).data_ptr(); + for (int i = 0; i < head_dim; i++) { + printf("%.6f ", static_cast(out_ref_accessor[i])); + } + printf("\n\n"); +} + + +int main() { + const int num_heads = 32; + const int num_heads_kv = 32; + const int head_dim = 128; + + const int batch_size = 1; + + const std::string quant_mode = "k-channel"; + const int num_bits = 4; + const int group_size = 128; + const int seqlen_kvs[] = {1024}; + + for (int base_seqlen : seqlen_kvs) { + int seqlen_kv = base_seqlen - 33; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen - 32; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen - 10; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen - 1; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 31; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 32; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 128 - 40; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 128 - 23; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 128 - 1; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 128 + 10; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 128 + 31; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 128 + 32; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 128 + 50; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 128 + 130; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 256 - 1; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 256; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + + seqlen_kv = base_seqlen + 256 + 1; + TestDecodingKernelCorrectness(batch_size, seqlen_kv, quant_mode, group_size); + } + + + return 0; +} \ No newline at end of file diff --git a/evaluation/ablation/script/test_bitblas.sh b/evaluation/ablation/script/test_bitblas.sh new file mode 100644 index 0000000..6b1528b --- /dev/null +++ b/evaluation/ablation/script/test_bitblas.sh @@ -0,0 +1 @@ +CUDA_DEVICE_ORDER=PCI_BUS_ID CUDA_VISIBLE_DEVICES=0 python test_bitblas.py \ No newline at end of file diff --git a/evaluation/ablation/test_bitblas.py b/evaluation/ablation/test_bitblas.py new file mode 100644 index 0000000..8d577d7 --- /dev/null +++ b/evaluation/ablation/test_bitblas.py @@ -0,0 +1,66 @@ +import bitblas +import torch +import time +import numpy as np + +# uncomment to enable debug output +# bitblas.set_log_level("Debug") + +# Prefill +n_heads = 1 +seq_len = 128 +dim = 128 +matmul_config = bitblas.MatmulConfig( + M=1, # M dimension + N=n_heads*seq_len, # N dimension + K=dim, # K dimension + A_dtype="float16", # activation A dtype + W_dtype="int4", # weight W dtype + accum_dtype="float16", # accumulation dtype + out_dtype="float16", # output dtype + layout="nt", # matrix layout, "nt" indicates the layout of A is non-transpose and the layout of W is transpose + with_bias=False, # bias + # configs for weight only quantization + group_size=None, # setting for grouped quantization + with_scaling=False, # setting for scaling factor + with_zeros=False, # setting for zeros + zeros_mode=None, # setting for how to calculating zeros +) + +matmul = bitblas.Matmul(config=matmul_config) + +# Create input matrices +# input_tensor = torch.rand((1, dim), dtype=torch.float16).cuda() +weight_tensor = torch.randint(0, 7, (n_heads*seq_len, dim), dtype=torch.int8).cuda() + +# Warmup runs +print("\nWarming up...") +for _ in range(5): + _ = matmul.transform_weight(weight_tensor) + torch.cuda.synchronize() + +# Timing runs +num_runs = 10 +times = [] + +print(f"\nRunning {num_runs} timing iterations...") + +for i in range(num_runs): + torch.cuda.synchronize() + start_time = time.perf_counter() + + weight_tensor_int4 = matmul.transform_weight(weight_tensor) + + torch.cuda.synchronize() + end_time = time.perf_counter() + + elapsed_time = (end_time - start_time) * 1000 # Convert to milliseconds + times.append(elapsed_time) + + if (i + 1) % 20 == 0: + print(f" Completed {i + 1}/{num_runs} runs") + +times = np.array(times) +mean_time = np.mean(times) + +print(f"Mean time: {mean_time} ms") \ No newline at end of file diff --git a/evaluation/ablation/test_marlin.py b/evaluation/ablation/test_marlin.py new file mode 100644 index 0000000..d5e0053 --- /dev/null +++ b/evaluation/ablation/test_marlin.py @@ -0,0 +1,183 @@ +import torch +import torch.nn as nn +import numpy as np +import time + +# Define the missing constants and functions for Marlin Layer +# These would normally come from marlin-specific modules +_perm = torch.randperm(128) # Placeholder permutation +_scale_perm = torch.randperm(4) # Placeholder scale permutation +_scale_perm_single = torch.randperm(2) # Placeholder single scale permutation + +def mul(A, B, C, s, workspace): + """Placeholder implementation of marlin mul function""" + # This is a simplified version - actual implementation would use CUDA kernels + A_flat = A.view(-1, A.shape[-1]) + C_flat = C.view(-1, C.shape[-1]) + + # Simulated quantized matrix multiplication + # In real implementation, this would dequantize B using s and perform actual GEMM + result = torch.matmul(A_flat.half(), torch.randn(A.shape[-1], C.shape[-1], device=A.device, dtype=torch.half)) + C_flat.copy_(result) + +class Layer(nn.Module): + """PyTorch compatible Marlin layer; 4-bit (symmetric grouped) linear layer without bias.""" + + def __init__(self, infeatures, outfeatures, groupsize=-1): + """Create an empty Marlin layer. + @infeatures: number of input features (must be divisible by 128) + @outfeatures: number of output features (must be divisible by 256) + @groupsize: quantization groupsize (must be -1 or 128) + """ + super().__init__() + if groupsize not in [-1, 128]: + raise ValueError('Only groupsize -1 and 128 are supported.') + if infeatures % 128 != 0 or outfeatures % 256 != 0: + raise ValueError('`infeatures` must be divisible by 128 and `outfeatures` by 256.') + if groupsize == -1: + groupsize = infeatures + if infeatures % groupsize != 0: + raise ValueError('`infeatures` must be divisible by `groupsize`.') + self.k = infeatures + self.n = outfeatures + self.groupsize = groupsize + self.register_buffer('B', torch.empty((self.k // 16, self.n * 16 // 8), dtype=torch.int)) + self.register_buffer('s', torch.empty((self.k // groupsize, self.n), dtype=torch.half)) + # 128 is currently the minimum `tile_n`, hence it gives the maximum workspace size; 16 is the default `max_par` + self.register_buffer('workspace', torch.zeros(self.n // 128 * 16, dtype=torch.int), persistent=False) + + def forward(self, A): + C = torch.empty(A.shape[:-1] + (self.s.shape[1],), dtype=A.dtype, device=A.device) + mul(A.view((-1, A.shape[-1])), self.B, C.view((-1, C.shape[-1])), self.s, self.workspace) + return C + + def pack(self, linear, scales): + """Pack a fake-quantized linear layer into this actual Marlin representation. + @linear: fake-quantized `torch.nn.Linear` layer to convert (must be of type `torch.half`) + @scales: corresponding quantization scales of shape `(infeatures, groups)` + """ + if linear.weight.dtype != torch.half: + raise ValueError('Only `torch.half` weights are supported.') + tile = 16 + maxq = 2 ** 4 - 1 + s = scales.t() + w = linear.weight.data.t() + if self.groupsize != self.k: + w = w.reshape((-1, self.groupsize, self.n)) + w = w.permute(1, 0, 2) + w = w.reshape((self.groupsize, -1)) + s = s.reshape((1, -1)) + w = torch.round(w / s).int() + w += (maxq + 1) // 2 + w = torch.clamp(w, 0, maxq) + if self.groupsize != self.k: + w = w.reshape((self.groupsize, -1, self.n)) + w = w.permute(1, 0, 2) + w = w.reshape((self.k, self.n)).contiguous() + s = s.reshape((-1, len(_scale_perm)))[:, _scale_perm] + else: + s = s.reshape((-1, len(_scale_perm_single)))[:, _scale_perm_single] + s = s.reshape((-1, self.n)).contiguous() + w = w.reshape((self.k // tile, tile, self.n // tile, tile)) + w = w.permute((0, 2, 1, 3)) + w = w.reshape((self.k // tile, self.n * tile)) + res = w + res = res.reshape((-1, _perm.numel()))[:, _perm].reshape(res.shape) + q = np.zeros((res.shape[0], res.shape[1] // 8), dtype=np.uint32) + res = res.cpu().numpy().astype(np.uint32) + for i in range(8): + q |= res[:, i::8] << 4 * i + q = torch.from_numpy(q.astype(np.int32)).to(w.device) + self.B[:, :] = q.to(self.B.device) + self.s[:, :] = s.to(self.s.device) + + +def test_marlin_pack_latency(): + """Test the Marlin layer pack function latency""" + print("Testing Marlin Layer pack function with weight dimensions (1024, 128) and group_size=128") + + # Based on user requirements: weight (1024, 128) means out_features=1024, in_features=128 + # After transpose in pack method: (128, 1024) -> infeatures=128, outfeatures=1024 + infeatures = 128 + outfeatures = 1024 + groupsize = 128 + + # Validate constraints + print(f"infeatures: {infeatures}, outfeatures: {outfeatures}, groupsize: {groupsize}") + print(f"infeatures % 128 = {infeatures % 128}") + print(f"outfeatures % 256 = {outfeatures % 256}") + print(f"infeatures % groupsize = {infeatures % groupsize}") + + # Create Marlin layer + marlin_layer = Layer(infeatures=infeatures, outfeatures=outfeatures, groupsize=groupsize) + + # Create a fake-quantized linear layer to pack + linear = nn.Linear(in_features=outfeatures, out_features=infeatures, bias=False) + linear.weight.data = torch.randn(infeatures, outfeatures, dtype=torch.half) + + # Create random scales with proper shape + # scales shape should be (infeatures, groups) = (128, 1) since groupsize=128=infeatures + num_groups = infeatures // groupsize + scales = torch.randn(infeatures, num_groups, dtype=torch.half) * 0.1 + 1.0 # scales around 1.0 + + print(f"Linear layer weight shape: {linear.weight.shape}") + print(f"Scales shape: {scales.shape}") + + # Move to GPU if available + if torch.cuda.is_available(): + marlin_layer = marlin_layer.cuda() + linear = linear.cuda() + scales = scales.cuda() + print("Using GPU for testing") + else: + print("Using CPU for testing") + + # Test pack function latency + print("\nTesting pack function latency...") + + # Warm up + print("Warming up...") + for _ in range(5): + marlin_layer.pack(linear, scales) + + # Measure latency + num_runs = 100 + print(f"Running {num_runs} iterations...") + + if torch.cuda.is_available(): + torch.cuda.synchronize() + + start_time = time.time() + + for _ in range(num_runs): + marlin_layer.pack(linear, scales) + + if torch.cuda.is_available(): + torch.cuda.synchronize() + + end_time = time.time() + + avg_latency = (end_time - start_time) / num_runs * 1000 # Convert to milliseconds + total_time = (end_time - start_time) * 1000 # Convert to milliseconds + + print(f"\nResults:") + print(f"Average pack function latency: {avg_latency:.4f} ms") + print(f"Total time for {num_runs} runs: {total_time:.2f} ms") + print(f"Throughput: {num_runs / (total_time / 1000):.2f} packs/sec") + + +if __name__ == "__main__": + # Set device + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Using device: {device}") + + # Set random seed for reproducibility + torch.manual_seed(42) + np.random.seed(42) + + try: + test_marlin_pack_latency() + except Exception as e: + print(f"Error during testing: {e}") + import traceback + traceback.print_exc() \ No newline at end of file diff --git a/evaluation/bench_throughput.py b/evaluation/bench_throughput.py new file mode 100644 index 0000000..89b29f5 --- /dev/null +++ b/evaluation/bench_throughput.py @@ -0,0 +1,140 @@ +import argparse +import dataclasses +import time +import numpy as np +import torch +from tqdm.auto import tqdm +from llama import LlamaForCausalLM +from transformers import LlamaConfig, AutoTokenizer + +@dataclasses.dataclass +class ModelConfig: + model_path: str + dtype: str = dataclasses.field(default="float16") +# device: str = dataclasses.field(default="cuda:0") + + +def load_model(args): + # device = torch.device(args.device) + dtype = getattr(torch, args.dtype) + torch.set_default_dtype(dtype) + + config = LlamaConfig.from_pretrained(args.model_path) + config.attn_backend = args.attn_backend + config.num_bits = args.num_bits + config.quant_mode = args.quant_mode + config.group_size = args.group_size + config.residual_block_size = 128 if args.num_bits == 4 else 256 + + model = LlamaForCausalLM.from_pretrained( + args.model_path, + config=config, + device_map="auto", + torch_dtype=dtype + ) + return model + +@torch.inference_mode() +def benchmark_throughput(): + parser = argparse.ArgumentParser() + parser.add_argument("--model_path", default="llama3-8b-instruct") + parser.add_argument("--batch_size", type=int, default=1) + parser.add_argument("--context_len", type=int, default=2*1024) + parser.add_argument("--decode_len", type=int, default=256) + parser.add_argument("--iteration", type=int, default=10) + parser.add_argument("--device", type=str, default="cuda:0") + parser.add_argument("--dtype", type=str, default="float16") + parser.add_argument("--attn_backend", type=str, default="flash_attention_2") + parser.add_argument("--num_bits", type=int, default=4) + parser.add_argument("--quant_mode", type=str, default="k-channel") + parser.add_argument("--group_size", type=int, default=128) + + args = parser.parse_args() + + model = load_model(args) + + context_len = args.context_len + decode_len = args.decode_len + batch_size = args.batch_size + + dtype = getattr(torch, args.dtype) + device = torch.device(args.device) + hidden_size = model.config.hidden_size + + prefill_latency = [] + decode_latency = [] + + for iter_idx in tqdm(range(args.iteration)): + # clear cuda cache + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats(device) + + # Prefill Stage + ts = time.perf_counter() + hidden_states = torch.randn(batch_size, context_len, hidden_size, dtype=dtype, device=device) + out = model( + inputs_embeds=hidden_states, + use_cache=True + ) + torch.cuda.synchronize() + te = time.perf_counter() + prefill_latency.append(te - ts) + + # Memory stats after prefill + if iter_idx == 0: + print(f"GPU Memory Allocated: {torch.cuda.memory_allocated(device) / 1e6:.2f} MB") + print(f"Peak GPU Memory: {torch.cuda.max_memory_allocated(device) / 1e6:.2f} MB") + + # Warm up for decode + for _ in range(5): + hidden_states = torch.randn(batch_size, 1, hidden_size, dtype=dtype, device=device) + model( + inputs_embeds=hidden_states, + past_key_values=out.past_key_values, + use_cache=True, + ) + + # Decode Stage - measure total time for all tokens + ts_decode_total = time.perf_counter() + for _ in range(decode_len): + hidden_states = torch.randn(batch_size, 1, hidden_size, dtype=dtype, device=device) + out = model( + inputs_embeds=hidden_states, + past_key_values=out.past_key_values, + use_cache=True, + ) + torch.cuda.synchronize() + te_decode_total = time.perf_counter() + decode_latency.append(te_decode_total - ts_decode_total) + + # Calculate metrics + avg_prefill_latency = np.mean(prefill_latency) + avg_decode_latency = np.mean(decode_latency) + # avg_decode_latency -= 0.0019366741180 * 32 + + # Calculate throughput + prefill_throughput = (batch_size * context_len) / avg_prefill_latency + decode_throughput = (batch_size * decode_len) / avg_decode_latency + + # Print results in a table format + print("\n===== BENCHMARK RESULTS =====") + print(f"Model: {args.model_path}") + print(f"Batch Size: {batch_size}") + print(f"Context Length: {context_len}") + print(f"Decode Length: {decode_len}") + print(f"Quantization: {args.num_bits}-bit {args.quant_mode}") + print("\n--- Latency ---") + print(f"Avg Prefill Latency: {avg_prefill_latency:.4f} s") + print(f"Avg Decode Latency (total): {avg_decode_latency:.4f} s") + print(f"Avg Decode Latency (per token): {avg_decode_latency/decode_len:.4f} s") + print("\n--- Throughput ---") + print(f"Prefill Throughput: {prefill_throughput:.2f} tokens/s") + print(f"Decode Throughput: {decode_throughput:.2f} tokens/s") + + # CSV format for easy parsing + print("\n--- CSV Format ---") + print("batch_size,context_len,decode_len,prefill_latency,decode_latency,prefill_throughput,decode_throughput") + print(f"{batch_size},{context_len},{decode_len},{avg_prefill_latency:.4f},{avg_decode_latency:.4f},{prefill_throughput:.2f},{decode_throughput:.2f}") + +if __name__ == "__main__": + benchmark_throughput() \ No newline at end of file diff --git a/evaluation/example.py b/evaluation/example.py new file mode 100644 index 0000000..8a3672d --- /dev/null +++ b/evaluation/example.py @@ -0,0 +1,100 @@ +# LLaMA model with KIVI +import warnings +warnings.filterwarnings("ignore") +import torch +import random +import argparse + +from bit_decode import DynamicCache, StaticCache, Cache +import transformers.cache_utils +transformers.cache_utils.DynamicCache = DynamicCache +transformers.cache_utils.StaticCache = StaticCache +transformers.cache_utils.Cache = Cache + +from llama import LlamaForCausalLM +from qwen3 import Qwen3ForCausalLM +from transformers import LlamaConfig, Qwen3Config, AutoTokenizer +from datasets import load_dataset + +def main(): + # Parse command line arguments + parser = argparse.ArgumentParser(description='Run LLaMA model with KIVI') + parser.add_argument('--model_path', type=str, required=True, help='Path to the pretrained model') + parser.add_argument('--max_length', type=int, default=131072, help='Maximum length of the input sequence') + parser.add_argument('--num_bits', type=int, default=4, help='Number of bits for quantization') + parser.add_argument('--quant_mode', type=str, default='k-channel', help='Quantization mode') + parser.add_argument('--group_size', type=int, default=128, help='Group size for quantization') + parser.add_argument('--attn_backend', type=str, default='flash_attention_2', help='Attention implementation') + args = parser.parse_args() + + # For reproducibility + random.seed(0) + torch.manual_seed(0) + + if "Llama" in args.model_path: + config = LlamaConfig.from_pretrained(args.model_path) + elif "Qwen" in args.model_path: + config = Qwen3Config.from_pretrained(args.model_path) + + config._attn_implementation = "flash_attention_2" + config.attn_backend = args.attn_backend + config.num_bits = args.num_bits + config.quant_mode = args.quant_mode + config.group_size = args.group_size + config.residual_block_size = 128 if args.num_bits == 4 else 256 + + if "Llama" in args.model_path: + model = LlamaForCausalLM.from_pretrained( + pretrained_model_name_or_path=args.model_path, + config=config, + low_cpu_mem_usage=True, + torch_dtype=torch.float16, + device_map="auto" + ) + elif "Qwen" in args.model_path: + model = Qwen3ForCausalLM.from_pretrained( + pretrained_model_name_or_path=args.model_path, + config=config, + low_cpu_mem_usage=True, + torch_dtype=torch.float16, + device_map="auto" + ) + + enc = AutoTokenizer.from_pretrained( + args.model_path, + use_fast=False, + trust_remote_code=True, + padding_side='left', # Add this line + pad_token='' # Add this line + ) + + dataset = load_dataset('gsm8k', 'main') + + prompt = '' + for i in range(15): + prompt += 'Question: ' + dataset['train'][i]['question'] + '\nAnswer: ' + dataset['train'][i]['answer'] + '\n' + prompt += "Arnel had ten boxes of pencils with the same number of pencils in each box. He kept ten pencils and shared the remaining pencils equally with his five friends. If his friends got eight pencils each, how many pencils are in each box?" + + inputs = enc( + prompt, + return_tensors="pt", + padding=True, + truncation=True, + max_length=args.max_length, + return_attention_mask=True + ).to('cuda') + + output = model.generate( + inputs.input_ids, + attention_mask=inputs.attention_mask, + pad_token_id=enc.pad_token_id, + max_new_tokens=125 + ) + config_str = f"# prompt tokens: {inputs.input_ids.shape[1]}" + + # print(prompt + "\n" + "=" * 10 + f'\n{config_str}\n' + "=" * 10 + "\nOutput:") + # print("\n" + "=" * 10 + f'\n{config_str}\n' + "=" * 10 + "\nOutput:") + print(enc.decode(output[0].tolist()[inputs.input_ids.shape[1]:], skip_special_tokens=True)) + +if __name__ == "__main__": + main() diff --git a/evaluation/llama.py b/evaluation/llama.py new file mode 100644 index 0000000..3fe3241 --- /dev/null +++ b/evaluation/llama.py @@ -0,0 +1,1658 @@ +# coding=utf-8 +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import math +from typing import List, Optional, Tuple, Union + +import torch +import torch.utils.checkpoint +from torch import nn + +from transformers.activations import ACT2FN +# from transformers.cache_utils import Cache, StaticCache +from transformers.generation import GenerationMixin +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs, _flash_attention_forward +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + QuestionAnsweringModelOutput, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS +from transformers.modeling_utils import PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS +from transformers.utils import ( + LossKwargs, + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + is_flash_attn_greater_or_equal_2_10, + logging, + replace_return_docstrings, +) +from transformers.models.llama.configuration_llama import LlamaConfig + +from flash_attn import flash_attn_with_kvcache +from bit_decode import kvcache_pack_int, fwd_kvcache_int +from bit_decode import Cache, DynamicCache, StaticCache +# from transformers.cache_utils import Cache, DynamicCache, StaticCache + +logger = logging.get_logger(__name__) + +_CHECKPOINT_FOR_DOC = "meta-llama/Llama-2-7b-hf" +_CONFIG_FOR_DOC = "LlamaConfig" + + +class LlamaRMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + LlamaRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +ALL_LAYERNORM_LAYERS.append(LlamaRMSNorm) + + +class LlamaRotaryEmbedding(nn.Module): + def __init__( + self, + dim=None, + max_position_embeddings=2048, + base=10000, + device=None, + scaling_factor=1.0, + rope_type="default", + config: Optional[LlamaConfig] = None, + ): + super().__init__() + # TODO (joao): remove the `if` below, only used for BC + self.rope_kwargs = {} + if config is None: + logger.warning_once( + "`LlamaRotaryEmbedding` can now be fully parameterized by passing the model config through the " + "`config` argument. All other arguments will be removed in v4.46" + ) + self.rope_kwargs = { + "rope_type": rope_type, + "factor": scaling_factor, + "dim": dim, + "base": base, + "max_position_embeddings": max_position_embeddings, + } + self.rope_type = rope_type + self.max_seq_len_cached = max_position_embeddings + self.original_max_seq_len = max_position_embeddings + else: + # BC: "rope_type" was originally "type" + if config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, **self.rope_kwargs) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + def _dynamic_frequency_update(self, position_ids, device): + """ + dynamic RoPE layers should recompute `inv_freq` in the following situations: + 1 - growing beyond the cached sequence length (allow scaling) + 2 - the current sequence length is in the original scale (avoid losing precision with small sequences) + """ + seq_len = torch.max(position_ids) + 1 + if seq_len > self.max_seq_len_cached: # growth + inv_freq, self.attention_scaling = self.rope_init_fn( + self.config, device, seq_len=seq_len, **self.rope_kwargs + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation + self.max_seq_len_cached = seq_len + + if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False) + self.max_seq_len_cached = self.original_max_seq_len + + @torch.no_grad() + def forward(self, x, position_ids): + if "dynamic" in self.rope_type: + self._dynamic_frequency_update(position_ids, device=x.device) + + # Core RoPE block + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + # Force float32 (see https://github.com/huggingface/transformers/pull/29285) + device_type = x.device.type + device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention + cos = cos * self.attention_scaling + sin = sin * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): + """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev""" + + def __init__(self, *args, **kwargs): + logger.warning_once( + "`LlamaLinearScalingRotaryEmbedding` is deprecated an will be removed in v4.46. Please use " + "`LlamaRotaryEmbedding`, which now also does linear scaling (simply pass the model config to __init__)." + ) + kwargs["rope_type"] = "linear" + super().__init__(*args, **kwargs) + + +class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding): + """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla""" + + def __init__(self, *args, **kwargs): + logger.warning_once( + "`LlamaDynamicNTKScalingRotaryEmbedding` is deprecated an will be removed in v4.46. Please use " + "`LlamaRotaryEmbedding`, which now also does dynamic ntk scaling (simply pass the model config to " + "__init__)." + ) + kwargs["rope_type"] = "dynamic" + super().__init__(*args, **kwargs) + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class LlamaMLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +class LlamaAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: LlamaConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + if layer_idx is None: + logger.warning_once( + f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will " + "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` " + "when creating this class." + ) + + self.attention_dropout = config.attention_dropout + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) + self.num_key_value_heads = config.num_key_value_heads + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + self.is_causal = True + + self.num_bits = config.num_bits + self.pack_nums = 16 / self.num_bits + self.quant_mode = config.quant_mode + self.group_size = config.group_size + self.residual_block_size = config.residual_block_size + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) + + # TODO (joao): remove in v4.46 (RoPE is computed in the model, not in the decoder layers) + self.rotary_emb = LlamaRotaryEmbedding(config=self.config) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # use -1 to infer num_heads and num_key_value_heads as they may vary if tensor parallel is used + query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " + "removed and `position_embeddings` will be mandatory." + ) + cos, sin = self.rotary_emb(value_states, position_ids) + else: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + key_states = repeat_kv(key_states, self.num_key_value_groups) + value_states = repeat_kv(value_states, self.num_key_value_groups) + attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) + + if attention_mask is not None: # no matter the length, we just slice it + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + # upcast attention to fp32 + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + attn_output = torch.matmul(attn_weights, value_states) + + if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): + raise ValueError( + f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" + f" {attn_output.size()}" + ) + + attn_output = attn_output.transpose(1, 2).contiguous() + + attn_output = attn_output.reshape(bsz, q_len, -1) + + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +class LlamaFlashAttention2(LlamaAttention): + """ + Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if isinstance(past_key_value, StaticCache): + raise ValueError( + "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " + "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" + ) + + output_attentions = False + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " + "removed and `position_embeddings` will be mandatory." + ) + cos, sin = self.rotary_emb(value_states, position_ids) + else: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, dim=-2, cache_kwargs=cache_kwargs) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + dropout_rate = self.attention_dropout if self.training else 0.0 + + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + position_ids=position_ids, + dropout=dropout_rate, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + +class LlamaFlashDecodingAttention(LlamaAttention): + """ + Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if isinstance(past_key_value, StaticCache): + raise ValueError( + "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " + "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" + ) + + output_attentions = False + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " + "removed and `position_embeddings` will be mandatory." + ) + cos, sin = self.rotary_emb(value_states, position_ids) + else: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + if q_len == 1: + print("query_states1: ", query_states.shape) + print("key_states1: ", key_states.shape) + print("value_states1: ", value_states.shape) + + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx) + + print("query_states2: ", query_states.shape) + print("key_states2: ", key_states.shape) + print("value_states2: ", value_states.shape) + attn_output = flash_attn_with_kvcache( + query_states, + key_states, + value_states + ) + else: + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx) + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + position_ids=position_ids, + dropout=0.0, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + **kwargs, + ) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + +class LlamaBitDecoding(LlamaAttention): + """ + Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays + untouched. The only required change would be on the forward pass where it needs to correctly call the public API of + flash attention and deal with padding tokens in case the input contains any of them. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1. + # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0. + # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left). + self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10() + + + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: bool = False, + use_cache: bool = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + if isinstance(past_key_value, StaticCache): + raise ValueError( + "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` " + "make sure to use `sdpa` in the mean time, and open an issue at https://github.com/huggingface/transformers" + ) + + output_attentions = False + + bsz, q_len, _ = hidden_states.size() + + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # Flash attention requires the input to have the shape + # batch_size x seq_length x head_dim x hidden_dim + # therefore we just need to keep the original shape + query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) + + if position_embeddings is None: + logger.warning_once( + "The attention layers in this model are transitioning from computing the RoPE embeddings internally " + "through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed " + "`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be " + "removed and `position_embeddings` will be mandatory." + ) + cos, sin = self.rotary_emb(value_states, position_ids) + else: + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache + # to be able to avoid many of these transpose/reshape/view. + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + # start + batch_size = query_states.shape[0] + device = query_states.device + dtype = query_states.dtype + nheads_k = key_states.shape[2] + d = key_states.shape[3] + + if q_len == 1: + k_pack, k_params, v_pack, v_params = past_key_value.update_pack(None, None, None, None, self.layer_idx) + + seqlen_pack = v_pack.shape[1] + + seqlens_k = torch.full((batch_size,), seqlen_pack, dtype=torch.int32, device=device) + + # Get kv cache_residual and append new kv + k_residual = torch.zeros((batch_size, self.residual_block_size, nheads_k, d), device=device, dtype=dtype) + v_residual = torch.zeros((batch_size, self.residual_block_size, nheads_k, d), device=device, dtype=dtype) + k_residual_cache, v_residual_cache = past_key_value.update_residual(key_states, value_states, self.layer_idx) + + cur_residual_len = k_residual_cache.shape[1] + + k_residual[:, :cur_residual_len, :, :] = k_residual_cache + v_residual[:, :cur_residual_len, :, :] = v_residual_cache + + # Run attention with KV cache + attn_output, self.k_pack_new, self.k_params_new, self.v_pack_new, self.v_params_new = fwd_kvcache_int( + query_states, + k_pack, k_params, + v_pack, v_params, + k_residual, v_residual, seqlens_k, #seqlends_k + self.k_pack_new, self.k_params_new, self.v_pack_new, self.v_params_new, + None, # opt_block_table + 1.0 / math.sqrt(d), # sm_scale + self.quant_mode, + self.group_size, + self.residual_block_size, + cur_residual_len, # new_lens + self.num_bits + ) + + if cur_residual_len == self.residual_block_size: + k_pack, k_params, v_pack, v_params = past_key_value.update_pack(self.k_pack_new, self.k_params_new, self.v_pack_new, self.v_params_new, self.layer_idx) + past_key_value.clear_residual(self.layer_idx) + + # if self.layer_idx == 0: + # print(f"v_pack: {v_pack.shape}") + # print(f"cur_residual_len: {cur_residual_len}") + + else: + # key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx) + attn_output = _flash_attention_forward( + query_states, + key_states, + value_states, + attention_mask, + q_len, + position_ids=position_ids, + dropout=0.0, + sliding_window=getattr(self, "sliding_window", None), + use_top_left_mask=self._flash_attn_uses_top_left_mask, + is_causal=self.is_causal, + **kwargs, + ) + + seqlen_k = key_states.shape[1] + + cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, seqlen_k, dtype=torch.int32, device=device) + + residual_len = seqlen_k % self.residual_block_size + residual = residual_len > 0 + seqlen_k_pack = seqlen_k - residual_len + + # Initialize quantization tensors + if self.quant_mode == 'k-channel': + k_pack = torch.zeros((batch_size, int(seqlen_k_pack // self.pack_nums), nheads_k, d), dtype=torch.uint16, device=device) + k_params = torch.zeros((batch_size, int(seqlen_k_pack // self.group_size), nheads_k, d), dtype=torch.float32, device=device) + else: + k_pack = torch.zeros((batch_size, seqlen_k_pack, nheads_k, int(d // self.pack_nums)), dtype=torch.uint16, device=device) + k_params = torch.zeros((batch_size, int(d // self.group_size), nheads_k, seqlen_k_pack), dtype=torch.float32, device=device) + + v_pack = torch.zeros((batch_size, seqlen_k_pack, nheads_k, int(d // self.pack_nums)), dtype=torch.uint16, device=device) + v_params = torch.zeros((batch_size, int(d // self.group_size), nheads_k, seqlen_k_pack), dtype=torch.float32, device=device) + + if residual: + k_state_residual = key_states[:, -residual_len:, :, :] + v_state_residual = value_states[:, -residual_len:, :, :] + k_state_past = key_states[:, :-residual_len, :, :] + v_state_past = value_states[:, :-residual_len, :, :] + past_key_value.update_residual(k_state_residual, v_state_residual, self.layer_idx) + else: + k_state_past = key_states + v_state_past = value_states + + kvcache_pack_int( + k_state_past, k_pack, k_params, + v_state_past, v_pack, v_params, + None, # opt_block_table + cu_seqlens_k, + seqlen_k, + self.quant_mode, + self.group_size, + self.num_bits + ) + + past_key_value.update_pack(k_pack, k_params, v_pack, v_params, self.layer_idx) + + self.k_pack_new = torch.empty((batch_size, int(self.residual_block_size // self.pack_nums), nheads_k, k_pack.size(-1)), dtype=torch.uint16, device=device) + self.k_params_new = torch.empty((batch_size, int(self.residual_block_size // self.group_size), nheads_k, k_params.size(-1)), dtype=torch.float32, device=device) + self.v_pack_new = torch.empty((batch_size, self.residual_block_size, nheads_k, v_pack.size(-1)), dtype=torch.uint16, device=device) + self.v_params_new = torch.empty((batch_size, v_params.size(1), nheads_k, self.residual_block_size), dtype=torch.float32, device=device) + + attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() + attn_output = self.o_proj(attn_output) + + if not output_attentions: + attn_weights = None + + return attn_output, attn_weights, past_key_value + + +LLAMA_ATTENTION_CLASSES = { + "eager": LlamaAttention, + "flash_attention_2": LlamaFlashAttention2, + "flash_decoding": LlamaFlashDecodingAttention, + "bit_decoding": LlamaBitDecoding +} + + +class LlamaDecoderLayer(nn.Module): + def __init__(self, config: LlamaConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = LLAMA_ATTENTION_CLASSES[config.attn_backend](config=config, layer_idx=layer_idx) + # self.self_attn = LlamaFlashDecodingAttention(config=config, layer_idx=layer_idx) + # self.self_attn = LlamaFlashAttention2(config=config, layer_idx=layer_idx) + # self.self_attn = LlamaBitDecoding(config=config, layer_idx=layer_idx) + + self.mlp = LlamaMLP(config) + self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 + **kwargs, + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` + attention_mask (`torch.FloatTensor`, *optional*): + attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1, + query_sequence_length, key_sequence_length)` if default attention is used. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence + position_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*): + Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`, + with `head_dim` being the embedding dimension of each attention head. + kwargs (`dict`, *optional*): + Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code + into the model + """ + residual = hidden_states + + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs + + +LLAMA_START_DOCSTRING = r""" + This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the + library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads + etc.) + + This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. + Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage + and behavior. + + Parameters: + config ([`LlamaConfig`]): + Model configuration class with all the parameters of the model. Initializing with a config file does not + load the weights associated with the model, only the configuration. Check out the + [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + LLAMA_START_DOCSTRING, +) +class LlamaPreTrainedModel(PreTrainedModel): + config_class = LlamaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["LlamaDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + +LLAMA_INPUTS_DOCSTRING = r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide + it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + If `past_key_values` is used, optionally only the last `input_ids` have to be input (see + `past_key_values`). + + If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] + and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more + information on the default strategy. + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, + config.n_positions - 1]`. + + [What are position IDs?](../glossary#position-ids) + past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*): + Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention + blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values` + returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`. + + Two formats are allowed: + - a [`~cache_utils.Cache`] instance, see our + [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache); + - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of + shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy + cache format. + + The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the + legacy cache format will be returned. + + If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't + have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids` + of shape `(batch_size, sequence_length)`. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This + is useful if you want more control over how to convert `input_ids` indices into associated vectors than the + model's internal embedding lookup matrix. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see + `past_key_values`). + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): + Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, + this tensor is not affected by padding. It is used to update the cache in the correct position and to infer + the complete sequence length. +""" + + +@add_start_docstrings( + "The bare LLaMA Model outputting raw hidden-states without any specific head on top.", + LLAMA_START_DOCSTRING, +) +class LlamaModel(LlamaPreTrainedModel): + """ + Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] + + Args: + config: LlamaConfig + """ + + def __init__(self, config: LlamaConfig): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [LlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = LlamaRotaryEmbedding(config=config) + + self.gradient_checkpointing = False + if getattr(config, "pretraining_tp", 1) != 1: + logger.warn("`pretraining_tp` is deprecated, please use `model.tensor_parallel` instead.") + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ) -> Union[Tuple, BaseModelOutputWithPast]: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # kept for BC (non `Cache` `past_key_values` inputs) + return_legacy_cache = False + if use_cache and not isinstance(past_key_values, Cache): + # return_legacy_cache = False + if past_key_values is None: + past_key_values = DynamicCache() + else: + past_key_values = DynamicCache.from_legacy_cache(past_key_values) + logger.warning_once( + "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and " + "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class " + "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)" + ) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + next_decoder_cache = None + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + if self.gradient_checkpointing and self.training: + layer_outputs = self._gradient_checkpointing_func( + decoder_layer.__call__, + hidden_states, + causal_mask, + position_ids, + past_key_values, + output_attentions, + use_cache, + cache_position, + position_embeddings, + ) + else: + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + if use_cache: + next_decoder_cache = layer_outputs[2 if output_attentions else 1] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + next_cache = next_decoder_cache if use_cache else None + + if return_legacy_cache: + next_cache = next_cache.to_legacy_cache() + + if not return_dict: + return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=next_cache, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: torch.Tensor, + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool, + ): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions: + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + is_training=self.training, + ): + return None + + dtype, device = input_tensor.dtype, input_tensor.device + sequence_length = input_tensor.shape[1] + if using_static_cache: + target_length = past_key_values.get_max_cache_shape() + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + device=device, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + ) + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type == "cuda" + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + min_dtype = torch.finfo(dtype).min + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + @staticmethod + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + device: torch.device, + cache_position: torch.Tensor, + batch_size: int, + **kwargs, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape + `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, + to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + device (`torch.device`): + The device to plcae the 4D attention mask on. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device + ) + if sequence_length != 1: + causal_mask = torch.triu(causal_mask, diagonal=1) + causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1) + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :] + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + + return causal_mask + + +class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... + + +class LlamaForCausalLM(LlamaPreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + + def __init__(self, config): + super().__init__(config) + self.model = LlamaModel(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + num_logits_to_keep: int = 0, + **kwargs: Unpack[KwargsForCausalLM], + ) -> Union[Tuple, CausalLMOutputWithPast]: + r""" + Args: + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + num_logits_to_keep (`int`, *optional*): + Calculate logits for the last `num_logits_to_keep` tokens. If `0`, calculate logits for all + `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that + token can save memory, which becomes pretty significant for long sequences or large vocabulary size. + + Returns: + + Example: + + ```python + >>> from transformers import AutoTokenizer, LlamaForCausalLM + + >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = outputs[0] + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + The LLaMa Model transformer with a sequence classification head on top (linear layer). + + [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """, + LLAMA_START_DOCSTRING, +) +class LlamaForSequenceClassification(LlamaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = LlamaModel(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, SequenceClassifierOutputWithPast]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + transformer_outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + hidden_states = transformer_outputs[0] + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + sequence_lengths = -1 + else: + if input_ids is not None: + # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility + sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1 + sequence_lengths = sequence_lengths % input_ids.shape[-1] + sequence_lengths = sequence_lengths.to(logits.device) + else: + sequence_lengths = -1 + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + if not return_dict: + output = (pooled_logits,) + transformer_outputs[1:] + return ((loss,) + output) if loss is not None else output + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@add_start_docstrings( + """ +The Llama Model transformer with a span classification head on top for extractive question-answering tasks like +SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`). + """, + LLAMA_START_DOCSTRING, +) +class LlamaForQuestionAnswering(LlamaPreTrainedModel): + base_model_prefix = "transformer" + + # Copied from transformers.models.bloom.modeling_bloom.BloomForQuestionAnswering.__init__ with Bloom->Llama + def __init__(self, config): + super().__init__(config) + self.transformer = LlamaModel(config) + self.qa_outputs = nn.Linear(config.hidden_size, 2) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.transformer.embed_tokens + + def set_input_embeddings(self, value): + self.transformer.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.FloatTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + **kwargs, + ) -> Union[Tuple, QuestionAnsweringModelOutput]: + r""" + start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the start of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for position (index) of the end of the labelled span for computing the token classification loss. + Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence + are not taken into account for computing the loss. + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.transformer( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + loss = None + if start_positions is not None and end_positions is not None: + loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs) + + if not return_dict: + output = (start_logits, end_logits) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return QuestionAnsweringModelOutput( + loss=loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + The Llama Model transformer with a token classification head on top (a linear layer on top of the hidden-states + output) e.g. for Named-Entity-Recognition (NER) tasks. + """, + LLAMA_START_DOCSTRING, +) +class LlamaForTokenClassification(LlamaPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = LlamaModel(config) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=TokenClassifierOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, TokenClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.config) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/evaluation/print.log1 b/evaluation/print.log1 new file mode 100644 index 0000000..c378c01 --- /dev/null +++ b/evaluation/print.log1 @@ -0,0 +1,6 @@ + +Answer: Arnel shared 5 x 8 = <<5*8=40>>40 pencils with his friends. +He had 40 + 10 = <<40+10=50>>50 pencils in total. +Since there are ten boxes of pencils, there are 50 / 10 = <<50/10=5>>5 pencils in each box. +#### 5 +Question: A bakery has 600 cups of flour. If they use 12 cups of flour for every batch of bread and 8 cups of flour for every diff --git a/evaluation/qwen3.py b/evaluation/qwen3.py new file mode 100644 index 0000000..54d786d --- /dev/null +++ b/evaluation/qwen3.py @@ -0,0 +1,1248 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/qwen3/modular_qwen3.py. +# Do NOT edit this file manually as any edits will be overwritten by the generation of +# the file from the modular. If any change should be done, please apply the change to the +# modular_qwen3.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# coding=utf-8 +# Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import Callable, Optional, Tuple, Union + +import torch +from torch import nn +import math +from transformers.activations import ACT2FN +# from transformers.cache_utils import Cache, SlidingWindowCache, StaticCache +from transformers.generation import GenerationMixin +from transformers.integrations import use_kernel_forward_from_hub +from transformers.modeling_attn_mask_utils import AttentionMaskConverter +from transformers.modeling_flash_attention_utils import FlashAttentionKwargs +from transformers.modeling_layers import GradientCheckpointingLayer +from transformers.modeling_outputs import ( + BaseModelOutputWithPast, + CausalLMOutputWithPast, + QuestionAnsweringModelOutput, + SequenceClassifierOutputWithPast, + TokenClassifierOutput, +) +from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update +from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from transformers.processing_utils import Unpack +from transformers.utils import LossKwargs, auto_docstring, can_return_tuple, is_torch_flex_attn_available, logging +from transformers.models.qwen3.configuration_qwen3 import Qwen3Config + +from flash_attn import flash_attn_with_kvcache +from bit_decode import kvcache_pack_int, fwd_kvcache_int +from bit_decode import Cache, DynamicCache, StaticCache + +if is_torch_flex_attn_available(): + from torch.nn.attention.flex_attention import BlockMask + + from transformers.integrations.flex_attention import make_flex_block_causal_mask + + +logger = logging.get_logger(__name__) + + +@use_kernel_forward_from_hub("RMSNorm") +class Qwen3RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + Qwen3RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + def extra_repr(self): + return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}" + + +class Qwen3MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +def rotate_half(x): + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """ + This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, + num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) + """ + batch, num_key_value_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) + return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) + + +def eager_attention_forward( + module: nn.Module, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor], + scaling: float, + dropout: float = 0.0, + **kwargs, +): + key_states = repeat_kv(key, module.num_key_value_groups) + value_states = repeat_kv(value, module.num_key_value_groups) + + attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling + if attention_mask is not None: + causal_mask = attention_mask[:, :, :, : key_states.shape[-2]] + attn_weights = attn_weights + causal_mask + + attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype) + attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training) + attn_output = torch.matmul(attn_weights, value_states) + attn_output = attn_output.transpose(1, 2).contiguous() + + return attn_output, attn_weights + + +class Qwen3Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = config.attention_dropout + self.is_causal = True + + self.q_proj = nn.Linear( + config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias + ) + self.k_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.v_proj = nn.Linear( + config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias + ) + self.o_proj = nn.Linear( + config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias + ) + self.q_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim! + self.k_norm = Qwen3RMSNorm(self.head_dim, eps=config.rms_norm_eps) # thus post q_norm does not need reshape + self.sliding_window = config.sliding_window + if not ( + self.config.use_sliding_window + and getattr(self.config, "sliding_window", None) is not None + and self.layer_idx >= self.config.max_window_layers + ): + self.sliding_window = None + + self.num_bits = config.num_bits + self.pack_nums = 16 / self.num_bits + self.quant_mode = config.quant_mode + self.group_size = config.group_size + self.residual_block_size = config.residual_block_size + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config.attn_backend] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, # diff with Llama + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +class Qwen3FlashDecodingAttention(Qwen3Attention): + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__(config, layer_idx) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + q_len = input_shape[-1] + + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_value is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) + + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS["flash_attention_2"] + + if q_len == 1: + + query_states = query_states.transpose(1, 2) + key_states = key_states.transpose(1, 2) + value_states = value_states.transpose(1, 2) + + attn_output = flash_attn_with_kvcache( + query_states, + key_states, + value_states + ) + attn_weights = None + else: + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, # diff with Llama + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +class Qwen3BitDecoding(Qwen3Attention): + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__(config, layer_idx) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + attention_mask: Optional[torch.Tensor], + past_key_value: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: + input_shape = hidden_states.shape[:-1] + q_len = input_shape[-1] + + hidden_shape = (*input_shape, -1, self.head_dim) + + # Project queries, keys, values - keep in (batch, seq, num_heads, head_dim) format initially + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)) + value_states = self.v_proj(hidden_states).view(hidden_shape) + + # Apply RoPE - needs (batch, num_heads, seq, head_dim) format + cos, sin = position_embeddings + query_states_rope = query_states.transpose(1, 2) + key_states_rope = key_states.transpose(1, 2) + query_states_rope, key_states_rope = apply_rotary_pos_emb(query_states_rope, key_states_rope, cos, sin) + + batch_size = query_states.shape[0] + device = query_states.device + dtype = query_states.dtype + + if q_len == 1: + # Decoding path: use (batch, seq, num_heads, head_dim) format directly + # Transpose back from RoPE format to required format + query_states = query_states_rope.transpose(1, 2) + key_states = key_states_rope.transpose(1, 2) + # value_states already in correct format + + nheads_k = key_states.shape[2] + d = key_states.shape[3] + + k_pack, k_params, v_pack, v_params = past_key_value.update_pack(None, None, None, None, self.layer_idx) + + seqlen_pack = v_pack.shape[1] + + seqlens_k = torch.full((batch_size,), seqlen_pack, dtype=torch.int32, device=device) + + # Get kv cache_residual and append new kv + k_residual = torch.zeros((batch_size, self.residual_block_size, nheads_k, d), device=device, dtype=dtype) + v_residual = torch.zeros((batch_size, self.residual_block_size, nheads_k, d), device=device, dtype=dtype) + k_residual_cache, v_residual_cache = past_key_value.update_residual(key_states, value_states, self.layer_idx) + + cur_residual_len = k_residual_cache.shape[1] + + k_residual[:, :cur_residual_len, :, :] = k_residual_cache + v_residual[:, :cur_residual_len, :, :] = v_residual_cache + + # Run attention with KV cache + attn_output, self.k_pack_new, self.k_params_new, self.v_pack_new, self.v_params_new = fwd_kvcache_int( + query_states, + k_pack, k_params, + v_pack, v_params, + k_residual, v_residual, seqlens_k, #seqlends_k + self.k_pack_new, self.k_params_new, self.v_pack_new, self.v_params_new, + None, # opt_block_table + 1.0 / math.sqrt(d), # sm_scale + self.quant_mode, + self.group_size, + self.residual_block_size, + cur_residual_len, # new_lens + self.num_bits + ) + + if cur_residual_len == self.residual_block_size: + k_pack, k_params, v_pack, v_params = past_key_value.update_pack(self.k_pack_new, self.k_params_new, self.v_pack_new, self.v_params_new, self.layer_idx) + past_key_value.clear_residual(self.layer_idx) + + attn_weights = None + + else: + # Prefill path: use (batch, num_heads, seq, head_dim) format for attention + attention_interface: Callable = eager_attention_forward + if self.config._attn_implementation != "eager": + if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + logger.warning_once( + "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to " + 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.' + ) + else: + attention_interface = ALL_ATTENTION_FUNCTIONS["flash_attention_2"] + + # Use RoPE-applied states for attention (already in correct format) + attn_output, attn_weights = attention_interface( + self, + query_states_rope, + key_states_rope, + value_states.transpose(1, 2), # Transpose value_states for attention + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + sliding_window=self.sliding_window, # diff with Llama + **kwargs, + ) + + # Now transpose key/value states for packing (batch, seq, num_heads, head_dim) + key_states = key_states_rope.transpose(1, 2) + value_states = value_states # Already in correct format + + nheads_k = key_states.shape[2] + d = key_states.shape[3] + seqlen_k = key_states.shape[1] + + cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k, seqlen_k, dtype=torch.int32, device=device) + + residual_len = seqlen_k % self.residual_block_size + residual = residual_len > 0 + seqlen_k_pack = seqlen_k - residual_len + + if self.quant_mode == 'k-channel': + k_pack = torch.zeros((batch_size, int(seqlen_k_pack // self.pack_nums), nheads_k, d), dtype=torch.uint16, device=device) + k_params = torch.zeros((batch_size, int(seqlen_k_pack // self.group_size), nheads_k, d), dtype=torch.float32, device=device) + + v_pack = torch.zeros((batch_size, seqlen_k_pack, nheads_k, int(d // self.pack_nums)), dtype=torch.uint16, device=device) + v_params = torch.zeros((batch_size, int(d // self.group_size), nheads_k, seqlen_k_pack), dtype=torch.float32, device=device) + + if residual: + k_state_residual = key_states[:, -residual_len:, :, :] + v_state_residual = value_states[:, -residual_len:, :, :] + k_state_past = key_states[:, :-residual_len, :, :] + v_state_past = value_states[:, :-residual_len, :, :] + past_key_value.update_residual(k_state_residual, v_state_residual, self.layer_idx) + else: + k_state_past = key_states + v_state_past = value_states + + kvcache_pack_int( + k_state_past, k_pack, k_params, + v_state_past, v_pack, v_params, + None, # opt_block_table + cu_seqlens_k, + seqlen_k, + self.quant_mode, + self.group_size, + self.num_bits + ) + + past_key_value.update_pack(k_pack, k_params, v_pack, v_params, self.layer_idx) + + self.k_pack_new = torch.empty((batch_size, int(self.residual_block_size // self.pack_nums), nheads_k, k_pack.size(-1)), dtype=torch.uint16, device=device) + self.k_params_new = torch.empty((batch_size, int(self.residual_block_size // self.group_size), nheads_k, k_params.size(-1)), dtype=torch.float32, device=device) + self.v_pack_new = torch.empty((batch_size, self.residual_block_size, nheads_k, v_pack.size(-1)), dtype=torch.uint16, device=device) + self.v_params_new = torch.empty((batch_size, v_params.size(1), nheads_k, self.residual_block_size), dtype=torch.float32, device=device) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + +QWEN_ATTENTION_CLASSES = { + "eager": Qwen3Attention, + "flash_attention_2": Qwen3Attention, + "flash_decoding": Qwen3FlashDecodingAttention, + "bit_decoding": Qwen3BitDecoding +} + +class Qwen3DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Qwen3Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = QWEN_ATTENTION_CLASSES[config.attn_backend](config=config, layer_idx=layer_idx) + self.mlp = Qwen3MLP(config) + self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + if ( + config.sliding_window and config._attn_implementation != "flash_attention_2" + ): # diff with Llama is this warning + logger.warning_once( + f"Sliding Window Attention is enabled but not implemented for `{config._attn_implementation}`; " + "unexpected results may be encountered." + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Cache] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC + **kwargs: Unpack[FlashAttentionKwargs], + ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + # Self Attention + hidden_states, self_attn_weights = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + outputs = (hidden_states,) + if output_attentions: + outputs += (self_attn_weights,) + + return outputs + + +@auto_docstring +class Qwen3PreTrainedModel(PreTrainedModel): + config_class = Qwen3Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Qwen3DecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_cache_class = True + _supports_quantized_cache = True + _supports_static_cache = True + _supports_attention_backend = True + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + elif isinstance(module, Qwen3RMSNorm): + module.weight.data.fill_(1.0) + + +class Qwen3RotaryEmbedding(nn.Module): + def __init__(self, config: Qwen3Config, device=None): + super().__init__() + # BC: "rope_type" was originally "type" + if hasattr(config, "rope_scaling") and config.rope_scaling is not None: + self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type")) + else: + self.rope_type = "default" + self.max_seq_len_cached = config.max_position_embeddings + self.original_max_seq_len = config.max_position_embeddings + + self.config = config + self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type] + + inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device) + self.register_buffer("inv_freq", inv_freq, persistent=False) + self.original_inv_freq = self.inv_freq + + @torch.no_grad() + @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope) + def forward(self, x, position_ids): + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): # Force float32 + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +@auto_docstring +class Qwen3Model(Qwen3PreTrainedModel): + def __init__(self, config: Qwen3Config): + super().__init__(config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3RotaryEmbedding(config=config) + self.gradient_checkpointing = False + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + **flash_attn_kwargs: Unpack[FlashAttentionKwargs], + ) -> BaseModelOutputWithPast: + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if self.gradient_checkpointing and self.training and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." + ) + use_cache = False + + # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache + # if not isinstance(past_key_values, (type(None), Cache)): + # raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if use_cache and past_key_values is None: + print("use_cache and past_key_values is None") + past_key_values = DynamicCache() + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.unsqueeze(0) + + causal_mask = self._update_causal_mask( + attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions + ) + + hidden_states = inputs_embeds + + # create position embeddings to be shared across the decoder layers + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + + for decoder_layer in self.layers[: self.config.num_hidden_layers]: + if output_hidden_states: + all_hidden_states += (hidden_states,) + + layer_outputs = decoder_layer( + hidden_states, + attention_mask=causal_mask, + position_ids=position_ids, + past_key_value=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **flash_attn_kwargs, + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attns += (layer_outputs[1],) + + hidden_states = self.norm(hidden_states) + + # add hidden states from the last decoder layer + if output_hidden_states: + all_hidden_states += (hidden_states,) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + hidden_states=all_hidden_states, + attentions=all_self_attns, + ) + + def _update_causal_mask( + self, + attention_mask: Union[torch.Tensor, "BlockMask"], + input_tensor: torch.Tensor, + cache_position: torch.Tensor, + past_key_values: Cache, + output_attentions: bool = False, + ): + if self.config._attn_implementation == "flash_attention_2": + if attention_mask is not None and past_key_values is not None: + is_padding_right = attention_mask[:, -1].sum().item() != input_tensor.size()[0] + if is_padding_right: + raise ValueError( + "You are attempting to perform batched generation with padding_side='right'" + " this may lead to unexpected behaviour for Flash Attention version of Qwen3. Make sure to " + " call `tokenizer.padding_side = 'left'` before tokenizing the input. " + ) + if attention_mask is not None and 0.0 in attention_mask: + return attention_mask + return None + if self.config._attn_implementation == "flex_attention": + if isinstance(attention_mask, torch.Tensor): + attention_mask = make_flex_block_causal_mask(attention_mask) + return attention_mask + + # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in + # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail + # to infer the attention mask. + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + using_static_cache = isinstance(past_key_values, StaticCache) + using_sliding_window_cache = isinstance(past_key_values, SlidingWindowCache) + + # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward + if ( + self.config._attn_implementation == "sdpa" + and not (using_static_cache or using_sliding_window_cache) + and not output_attentions + ): + if AttentionMaskConverter._ignore_causal_mask_sdpa( + attention_mask, + inputs_embeds=input_tensor, + past_key_values_length=past_seen_tokens, + sliding_window=self.config.sliding_window, + is_training=self.training, + ): + return None + + dtype = input_tensor.dtype + min_dtype = torch.finfo(dtype).min + sequence_length = input_tensor.shape[1] + # SlidingWindowCache or StaticCache + if using_sliding_window_cache or using_static_cache: + target_length = past_key_values.get_max_cache_shape() + # DynamicCache or no cache + else: + target_length = ( + attention_mask.shape[-1] + if isinstance(attention_mask, torch.Tensor) + else past_seen_tokens + sequence_length + 1 + ) + + # In case the provided `attention` mask is 2D, we generate a causal mask here (4D). + print("prepare 4d causal mask with cache position") + causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position( + attention_mask, + sequence_length=sequence_length, + target_length=target_length, + dtype=dtype, + cache_position=cache_position, + batch_size=input_tensor.shape[0], + config=self.config, + past_key_values=past_key_values, + ) + + if ( + self.config._attn_implementation == "sdpa" + and attention_mask is not None + and attention_mask.device.type in ["cuda", "xpu", "npu"] + and not output_attentions + ): + # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when + # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. + # Details: https://github.com/pytorch/pytorch/issues/110213 + causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) + + return causal_mask + + @staticmethod + def _prepare_4d_causal_attention_mask_with_cache_position( + attention_mask: torch.Tensor, + sequence_length: int, + target_length: int, + dtype: torch.dtype, + cache_position: torch.Tensor, + batch_size: int, + config: Qwen3Config, + past_key_values: Cache, + ): + """ + Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape + `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. + + Args: + attention_mask (`torch.Tensor`): + A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. + sequence_length (`int`): + The sequence length being processed. + target_length (`int`): + The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. + dtype (`torch.dtype`): + The dtype to use for the 4D attention mask. + cache_position (`torch.Tensor`): + Indices depicting the position of the input sequence tokens in the sequence. + batch_size (`torch.Tensor`): + Batch size. + config (`Qwen3Config`): + The model's configuration class + past_key_values (`Cache`): + The cache class that is being used currently to generate + """ + if attention_mask is not None and attention_mask.dim() == 4: + # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. + causal_mask = attention_mask + else: + min_dtype = torch.finfo(dtype).min + causal_mask = torch.full( + (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device + ) + diagonal_attend_mask = torch.arange(target_length, device=cache_position.device) > cache_position.reshape( + -1, 1 + ) + text_config = config.get_text_config() + if getattr(text_config, "use_sliding_window", True) and text_config.sliding_window is not None: + # if we have sliding window, we should not attend to tokens beyond sliding window length, so we mask them out also + # the check is needed to verify is current checkpoint was trained with sliding window or not + if not isinstance(past_key_values, SlidingWindowCache) or sequence_length > target_length: + sliding_attend_mask = torch.arange(target_length, device=cache_position.device) <= ( + cache_position.reshape(-1, 1) - text_config.sliding_window + ) + diagonal_attend_mask.bitwise_or_(sliding_attend_mask) + causal_mask *= diagonal_attend_mask + causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) + if attention_mask is not None: + causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit + if attention_mask.shape[-1] > target_length: + attention_mask = attention_mask[:, :target_length] + mask_length = attention_mask.shape[-1] + padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( + causal_mask.device + ) + padding_mask = padding_mask == 0 + causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( + padding_mask, min_dtype + ) + return causal_mask + + +class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ... + + +@auto_docstring +class Qwen3ForCausalLM(Qwen3PreTrainedModel, GenerationMixin): + _tied_weights_keys = ["lm_head.weight"] + _tp_plan = {"lm_head": "colwise_rep"} + _pp_plan = {"lm_head": (["hidden_states"], ["logits"])} + + def __init__(self, config): + super().__init__(config) + self.model = Qwen3Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def set_decoder(self, decoder): + self.model = decoder + + def get_decoder(self): + return self.model + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + cache_position: Optional[torch.LongTensor] = None, + logits_to_keep: Union[int, torch.Tensor] = 0, + **kwargs: Unpack[KwargsForCausalLM], + ) -> CausalLMOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): + Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., + config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored + (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> from transformers import AutoTokenizer, Qwen3ForCausalLM + + >>> model = Qwen3ForCausalLM.from_pretrained("Qwen/Qwen3-8B") + >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-8B") + + >>> prompt = "Hey, are you conscious? Can you talk to me?" + >>> inputs = tokenizer(prompt, return_tensors="pt") + + >>> # Generate + >>> generate_ids = model.generate(inputs.input_ids, max_length=30) + >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] + "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you." + ```""" + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + # Only compute necessary logits, and do not upcast them to float if we are not computing the loss + slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + logits = self.lm_head(hidden_states[:, slice_indices, :]) + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs) + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring( + custom_intro=""" + The Qwen3 Model transformer with a sequence classification head on top (linear layer). + + [`Qwen3ForSequenceClassification`] uses the last token in order to do the classification, as other causal models + (e.g. GPT-2) do. + + Since it does classification on the last token, it requires to know the position of the last token. If a + `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If + no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the + padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in + each row of the batch). + """ +) +class Qwen3ForSequenceClassification(Qwen3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = Qwen3Model(config) + self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> SequenceClassifierOutputWithPast: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + + transformer_outputs: BaseModelOutputWithPast = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + hidden_states = transformer_outputs.last_hidden_state + logits = self.score(hidden_states) + + if input_ids is not None: + batch_size = input_ids.shape[0] + else: + batch_size = inputs_embeds.shape[0] + + if self.config.pad_token_id is None and batch_size != 1: + raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.") + if self.config.pad_token_id is None: + last_non_pad_token = -1 + elif input_ids is not None: + # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id + non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32) + token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32) + last_non_pad_token = (token_indices * non_pad_mask).argmax(-1) + else: + last_non_pad_token = -1 + logger.warning_once( + f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be " + "unexpected if using padding tokens in conjunction with `inputs_embeds.`" + ) + + pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token] + + loss = None + if labels is not None: + loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config) + + return SequenceClassifierOutputWithPast( + loss=loss, + logits=pooled_logits, + past_key_values=transformer_outputs.past_key_values, + hidden_states=transformer_outputs.hidden_states, + attentions=transformer_outputs.attentions, + ) + + +@auto_docstring +class Qwen3ForTokenClassification(Qwen3PreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.model = Qwen3Model(config) + if getattr(config, "classifier_dropout", None) is not None: + classifier_dropout = config.classifier_dropout + elif getattr(config, "hidden_dropout", None) is not None: + classifier_dropout = config.hidden_dropout + else: + classifier_dropout = 0.1 + self.dropout = nn.Dropout(classifier_dropout) + self.score = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + ) -> TokenClassifierOutput: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + + outputs: BaseModelOutputWithPast = self.model( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + sequence_output = outputs.last_hidden_state + sequence_output = self.dropout(sequence_output) + logits = self.score(sequence_output) + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.config) + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@auto_docstring +class Qwen3ForQuestionAnswering(Qwen3PreTrainedModel): + base_model_prefix = "transformer" + + def __init__(self, config): + super().__init__(config) + self.transformer = Qwen3Model(config) + self.qa_outputs = nn.Linear(config.hidden_size, 2) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.transformer.embed_tokens + + def set_input_embeddings(self, value): + self.transformer.embed_tokens = value + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Cache] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + start_positions: Optional[torch.LongTensor] = None, + end_positions: Optional[torch.LongTensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + **kwargs, + ) -> QuestionAnsweringModelOutput: + outputs: BaseModelOutputWithPast = self.transformer( + input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + ) + + sequence_output = outputs.last_hidden_state + + logits = self.qa_outputs(sequence_output) + start_logits, end_logits = logits.split(1, dim=-1) + start_logits = start_logits.squeeze(-1).contiguous() + end_logits = end_logits.squeeze(-1).contiguous() + + loss = None + if start_positions is not None and end_positions is not None: + loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs) + + return QuestionAnsweringModelOutput( + loss=loss, + start_logits=start_logits, + end_logits=end_logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +__all__ = [ + "Qwen3ForCausalLM", + "Qwen3ForQuestionAnswering", + "Qwen3Model", + "Qwen3PreTrainedModel", + "Qwen3ForSequenceClassification", + "Qwen3ForTokenClassification", +] diff --git a/evaluation/scripts/bench_throughput.sh b/evaluation/scripts/bench_throughput.sh new file mode 100644 index 0000000..dd59a0f --- /dev/null +++ b/evaluation/scripts/bench_throughput.sh @@ -0,0 +1,23 @@ +# BUDGET_POOL=('1024' '2048' '4096' '8192' '16384' '32768') +# BATCH_SIZE=('1' '2' '4' '8' '16' '32') + +BUDGET_POOL=('16384') +BATCH_SIZE=('1') + +for batch_size in ${BATCH_SIZE[@]}; do + for budget in ${BUDGET_POOL[@]}; do + python3 bench_throughput.py \ + --model_path meta-llama/Llama-3.1-70B-Instruct \ + --batch_size $batch_size \ + --context_len $budget \ + --decode_len 100 \ + --iteration 1 \ + --num_bits 4 \ + --quant_mode k-channel \ + --group_size 128 \ + --attn_backend flash_attention_2 + done +done + + +# flash_attention_2, flash_decoding, bit_decoding \ No newline at end of file diff --git a/evaluation/scripts/example.sh b/evaluation/scripts/example.sh new file mode 100644 index 0000000..fc9d5a6 --- /dev/null +++ b/evaluation/scripts/example.sh @@ -0,0 +1,11 @@ +python example.py \ + --model_path Qwen/Qwen3-8B \ + --max_length 131072 \ + --num_bits 4 \ + --quant_mode k-channel \ + --group_size 128 \ + --attn_backend bit_decoding # flash_attention_2, flash_decoding, bit_decoding + + +# meta-llama/Llama-3.1-8B-Instruct +# Qwen/Qwen3-8B \ No newline at end of file diff --git a/evaluation/test.py b/evaluation/test.py new file mode 100644 index 0000000..d3262d8 --- /dev/null +++ b/evaluation/test.py @@ -0,0 +1,160 @@ +import torch +import torch.nn as nn +import math + +import triton + +import numpy as np +import bit_decode_cuda as bit_decode_cuda +from bit_decode import kvcache_pack_int, fwd_kvcache_int +from bit_decode import DynamicCache + + +def attention_ref( + q, + k, + v, +): + """ + Arguments: + q: (batch_size, seqlen_q, nheads, head_dim) + k: (batch_size, seqlen_k, nheads_k, head_dim) + v: (batch_size, seqlen_k, nheads_k, head_dim) + Output: + output: (batch_size, seqlen_q, nheads, head_dim) + attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout + """ + dtype_og = q.dtype + + d = q.shape[-1] + + scores = torch.einsum("bthd,bshd->bhts", q / math.sqrt(d), k) + + attention = torch.softmax(scores, dim=-1).to(v.dtype) + + output = torch.einsum("bhts,bshd->bthd", attention, v) + + return output.to(dtype=dtype_og), attention.to(dtype=dtype_og) + + +# Quantization parameters +quant_mode = "k-channel" +num_bits = 4 +pack_nums = 16 / num_bits +group_size = 32 +residual_block_size = 128 + +device = "cuda" +dtype = torch.float16 + +layer_idx = 0 +batch_size = 1 +nheads = 32 +nheads_k = 32 +d = 128 + +seqlen_q = 1 +seqlen_k = 1024 +sm_scale = 1.0 / math.sqrt(d) + + +####### Round 1 : Prefill ####### +torch.manual_seed(42) + +q = torch.rand(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype) +k_state = torch.randn(batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype) +v_state = torch.randn(batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype) + +residual_len = seqlen_k % residual_block_size +residual = residual_len > 0 +seqlen_k_pack = seqlen_k - residual_len + +print(f"residual_len: {residual_len}, residual: {residual}, seqlen_k_pack: {seqlen_k_pack}") + +cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k_pack, seqlen_k_pack, + dtype=torch.int32, device=device) + +# Initialize quantization tensors +k_pack = torch.zeros((batch_size, int(seqlen_k_pack // pack_nums), nheads_k, d), dtype=torch.uint16, device=device) +k_params = torch.zeros((batch_size, int(seqlen_k_pack // group_size), nheads_k, d), dtype=torch.float32, device=device) + +v_pack = torch.zeros((batch_size, seqlen_k_pack, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device) +v_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seqlen_k_pack), dtype=torch.float32, device=device) + +# KV Cache Dynamic Cache +past_key_value = DynamicCache() + +if residual: + k_state_residual = k_state[:, -residual_len:, :, :] + v_state_residual = v_state[:, -residual_len:, :, :] + k_state_past = k_state[:, :-residual_len, :, :] + v_state_past = v_state[:, :-residual_len, :, :] + past_key_value.update_residual(k_state_residual, v_state_residual, layer_idx) +else: + k_state_past = k_state + v_state_past = v_state + +kvcache_pack_int( + k_state_past, k_pack, k_params, + v_state_past, v_pack, v_params, + None, # opt_block_table + cu_seqlens_k, + seqlen_k_pack, + quant_mode, + group_size, + num_bits +) +past_key_value.update_pack(k_pack, k_params, v_pack, v_params, layer_idx) + +# self +k_pack_new = torch.empty((batch_size, int(residual_block_size // pack_nums), nheads_k, k_pack.size(-1)), dtype=torch.uint16, device=device) +k_params_new = torch.empty((batch_size, int(residual_block_size // group_size), nheads_k, k_params.size(-1)), dtype=torch.float32, device=device) +v_pack_new = torch.empty((batch_size, residual_block_size, nheads_k, v_pack.size(-1)), dtype=torch.uint16, device=device) +v_params_new = torch.empty((batch_size, v_params.size(1), nheads_k, residual_block_size), dtype=torch.float32, device=device) + +####### Round 2-3 : Decode ####### +for round_idx in range(32): + k_new = torch.randn(batch_size, 1, nheads_k, d, device=device, dtype=dtype) + v_new = torch.randn(batch_size, 1, nheads_k, d, device=device, dtype=dtype) + + # Get kv cache_pack + k_pack, k_params, v_pack, v_params = past_key_value.update_pack(None, None, None, None, layer_idx) + + seqlen_pack = v_pack.shape[1] + seqlens_k = torch.full((batch_size,), seqlen_pack, dtype=torch.int32, device=device) + + # Get kv cache_residual and append new kv + k_residual = torch.zeros((batch_size, residual_block_size, nheads_k, d), device=device, dtype=dtype) + v_residual = torch.zeros((batch_size, residual_block_size, nheads_k, d), device=device, dtype=dtype) + k_residual_cache, v_residual_cache = past_key_value.update_residual(k_new, v_new, layer_idx) + + cur_residual_len = k_residual_cache.shape[1] + print(f"cur_residual_len: {cur_residual_len}") + + k_residual[:, :cur_residual_len, :, :] = k_residual_cache + v_residual[:, :cur_residual_len, :, :] = v_residual_cache + + out_bitdecode, k_pack_new, k_params_new, v_pack_new, v_params_new = fwd_kvcache_int( + q, + k_pack, k_params, + v_pack, v_params, + k_residual, v_residual, seqlens_k, #seqlens_k + k_pack_new, k_params_new, v_pack_new, v_params_new, + None, # opt_block_table + sm_scale, + quant_mode, + group_size, + residual_block_size, + cur_residual_len, # new_lens + num_bits + ) + + if cur_residual_len == residual_block_size: + past_key_value.update_pack(k_pack_new, k_params_new, v_pack_new, v_params_new, layer_idx) + past_key_value.clear_residual(layer_idx) + + k_state = torch.cat([k_state, k_new], dim=1) + v_state = torch.cat([v_state, v_new], dim=1) + + out_ref = attention_ref(q, k_state, v_state)[0] + print(f"Round {round_idx+2}: bitdecode vs pytorch: {(out_bitdecode - out_ref).abs().mean().item()}") \ No newline at end of file diff --git a/evaluation/test_bit_decode.ipynb b/evaluation/test_bit_decode.ipynb new file mode 100644 index 0000000..6f4c20c --- /dev/null +++ b/evaluation/test_bit_decode.ipynb @@ -0,0 +1,478 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/home/ddy/miniconda3/envs/bitdecode/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + } + ], + "source": [ + "import torch\n", + "import torch.nn as nn\n", + "import math\n", + "\n", + "import triton\n", + "\n", + "import numpy as np\n", + "import bit_decode_cuda as bit_decode_cuda\n", + "from bit_decode import kvcache_pack_int, fwd_kvcache_int\n", + "from bit_decode import DynamicCache" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "def attention_ref(\n", + " q,\n", + " k,\n", + " v,\n", + "):\n", + " \"\"\"\n", + " Arguments:\n", + " q: (batch_size, seqlen_q, nheads, head_dim)\n", + " k: (batch_size, seqlen_k, nheads_k, head_dim)\n", + " v: (batch_size, seqlen_k, nheads_k, head_dim)\n", + " Output:\n", + " output: (batch_size, seqlen_q, nheads, head_dim)\n", + " attention: (batch_size, nheads, seqlen_q, seqlen_k), softmax after dropout\n", + " \"\"\"\n", + " dtype_og = q.dtype\n", + "\n", + " d = q.shape[-1]\n", + "\n", + " scores = torch.einsum(\"bthd,bshd->bhts\", q / math.sqrt(d), k)\n", + " \n", + " attention = torch.softmax(scores, dim=-1).to(v.dtype)\n", + "\n", + " output = torch.einsum(\"bhts,bshd->bthd\", attention, v)\n", + "\n", + " return output.to(dtype=dtype_og), attention.to(dtype=dtype_og)" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "# Quantization parameters\n", + "quant_mode = \"k-channel\"\n", + "num_bits = 4\n", + "pack_nums = 16 / num_bits\n", + "group_size = 32\n", + "residual_block_size = 128\n", + "\n", + "device = \"cuda\"\n", + "dtype = torch.float16\n", + "\n", + "layer_idx = 0\n", + "batch_size = 1\n", + "nheads = 32\n", + "nheads_k = 32\n", + "d = 128\n", + "\n", + "seqlen_q = 1\n", + "seqlen_k = 1024\n", + "sm_scale = 1.0 / math.sqrt(d)" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Round 2: bitdecode vs pytorch: 0.0027313232421875\n", + "Round 3: bitdecode vs pytorch: 0.002727508544921875\n", + "Round 4: bitdecode vs pytorch: 0.0027256011962890625\n", + "Round 5: bitdecode vs pytorch: 0.0027217864990234375\n", + "Round 6: bitdecode vs pytorch: 0.002719879150390625\n", + "Round 7: bitdecode vs pytorch: 0.0027179718017578125\n", + "Round 8: bitdecode vs pytorch: 0.0027141571044921875\n", + "Round 9: bitdecode vs pytorch: 0.002712249755859375\n", + "Round 10: bitdecode vs pytorch: 0.00270843505859375\n", + "Round 11: bitdecode vs pytorch: 0.0027065277099609375\n", + "Round 12: bitdecode vs pytorch: 0.0027027130126953125\n", + "Round 13: bitdecode vs pytorch: 0.0027008056640625\n", + "Round 14: bitdecode vs pytorch: 0.002696990966796875\n", + "Round 15: bitdecode vs pytorch: 0.0026950836181640625\n", + "Round 16: bitdecode vs pytorch: 0.00269317626953125\n", + "Round 17: bitdecode vs pytorch: 0.0026912689208984375\n", + "Round 18: bitdecode vs pytorch: 0.0026874542236328125\n", + "Round 19: bitdecode vs pytorch: 0.002685546875\n", + "Round 20: bitdecode vs pytorch: 0.002681732177734375\n", + "Round 21: bitdecode vs pytorch: 0.0026798248291015625\n", + "Round 22: bitdecode vs pytorch: 0.00267791748046875\n", + "Round 23: bitdecode vs pytorch: 0.0026760101318359375\n", + "Round 24: bitdecode vs pytorch: 0.0026721954345703125\n", + "Round 25: bitdecode vs pytorch: 0.0026702880859375\n", + "Round 26: bitdecode vs pytorch: 0.0026683807373046875\n", + "Round 27: bitdecode vs pytorch: 0.0026645660400390625\n", + "Round 28: bitdecode vs pytorch: 0.00266265869140625\n", + "Round 29: bitdecode vs pytorch: 0.002658843994140625\n", + "Round 30: bitdecode vs pytorch: 0.0026569366455078125\n", + "Round 31: bitdecode vs pytorch: 0.002655029296875\n", + "Round 32: bitdecode vs pytorch: 0.0026531219482421875\n", + "Round 33: bitdecode vs pytorch: 0.0026493072509765625\n", + "Round 34: bitdecode vs pytorch: 0.00264739990234375\n", + "Round 35: bitdecode vs pytorch: 0.0026454925537109375\n", + "Round 36: bitdecode vs pytorch: 0.0026416778564453125\n", + "Round 37: bitdecode vs pytorch: 0.0026397705078125\n", + "Round 38: bitdecode vs pytorch: 0.0026378631591796875\n", + "Round 39: bitdecode vs pytorch: 0.002635955810546875\n", + "Round 40: bitdecode vs pytorch: 0.0026340484619140625\n", + "Round 41: bitdecode vs pytorch: 0.0026302337646484375\n", + "Round 42: bitdecode vs pytorch: 0.002628326416015625\n", + "Round 43: bitdecode vs pytorch: 0.0026264190673828125\n", + "Round 44: bitdecode vs pytorch: 0.00262451171875\n", + "Round 45: bitdecode vs pytorch: 0.002620697021484375\n", + "Round 46: bitdecode vs pytorch: 0.0026187896728515625\n", + "Round 47: bitdecode vs pytorch: 0.0026149749755859375\n", + "Round 48: bitdecode vs pytorch: 0.002613067626953125\n", + "Round 49: bitdecode vs pytorch: 0.0026092529296875\n", + "Round 50: bitdecode vs pytorch: 0.0026073455810546875\n", + "Round 51: bitdecode vs pytorch: 0.002605438232421875\n", + "Round 52: bitdecode vs pytorch: 0.0026035308837890625\n", + "Round 53: bitdecode vs pytorch: 0.0025997161865234375\n", + "Round 54: bitdecode vs pytorch: 0.002597808837890625\n", + "Round 55: bitdecode vs pytorch: 0.0025959014892578125\n", + "Round 56: bitdecode vs pytorch: 0.002593994140625\n", + "Round 57: bitdecode vs pytorch: 0.002590179443359375\n", + "Round 58: bitdecode vs pytorch: 0.0025882720947265625\n", + "Round 59: bitdecode vs pytorch: 0.00258636474609375\n", + "Round 60: bitdecode vs pytorch: 0.002582550048828125\n", + "Round 61: bitdecode vs pytorch: 0.0025806427001953125\n", + "Round 62: bitdecode vs pytorch: 0.0025787353515625\n", + "Round 63: bitdecode vs pytorch: 0.002574920654296875\n", + "Round 64: bitdecode vs pytorch: 0.0025730133056640625\n", + "Round 65: bitdecode vs pytorch: 0.00257110595703125\n", + "Round 66: bitdecode vs pytorch: 0.0025691986083984375\n", + "Round 67: bitdecode vs pytorch: 0.0025653839111328125\n", + "Round 68: bitdecode vs pytorch: 0.0025634765625\n", + "Round 69: bitdecode vs pytorch: 0.0025615692138671875\n", + "Round 70: bitdecode vs pytorch: 0.002559661865234375\n", + "Round 71: bitdecode vs pytorch: 0.0025577545166015625\n", + "Round 72: bitdecode vs pytorch: 0.0025539398193359375\n", + "Round 73: bitdecode vs pytorch: 0.002552032470703125\n", + "Round 74: bitdecode vs pytorch: 0.0025501251220703125\n", + "Round 75: bitdecode vs pytorch: 0.0025482177734375\n", + "Round 76: bitdecode vs pytorch: 0.002544403076171875\n", + "Round 77: bitdecode vs pytorch: 0.0025424957275390625\n", + "Round 78: bitdecode vs pytorch: 0.00254058837890625\n", + "Round 79: bitdecode vs pytorch: 0.0025386810302734375\n", + "Round 80: bitdecode vs pytorch: 0.002536773681640625\n", + "Round 81: bitdecode vs pytorch: 0.002532958984375\n", + "Round 82: bitdecode vs pytorch: 0.0025310516357421875\n", + "Round 83: bitdecode vs pytorch: 0.002529144287109375\n", + "Round 84: bitdecode vs pytorch: 0.0025272369384765625\n", + "Round 85: bitdecode vs pytorch: 0.00252532958984375\n", + "Round 86: bitdecode vs pytorch: 0.0025234222412109375\n", + "Round 87: bitdecode vs pytorch: 0.002521514892578125\n", + "Round 88: bitdecode vs pytorch: 0.0025196075439453125\n", + "Round 89: bitdecode vs pytorch: 0.0025177001953125\n", + "Round 90: bitdecode vs pytorch: 0.002513885498046875\n", + "Round 91: bitdecode vs pytorch: 0.0025119781494140625\n", + "Round 92: bitdecode vs pytorch: 0.00251007080078125\n", + "Round 93: bitdecode vs pytorch: 0.002506256103515625\n", + "Round 94: bitdecode vs pytorch: 0.0025043487548828125\n", + "Round 95: bitdecode vs pytorch: 0.00250244140625\n", + "Round 96: bitdecode vs pytorch: 0.0025005340576171875\n", + "Round 97: bitdecode vs pytorch: 0.002498626708984375\n", + "Round 98: bitdecode vs pytorch: 0.0024967193603515625\n", + "Round 99: bitdecode vs pytorch: 0.0024929046630859375\n", + "Round 100: bitdecode vs pytorch: 0.002490997314453125\n", + "Round 101: bitdecode vs pytorch: 0.0024890899658203125\n", + "Round 102: bitdecode vs pytorch: 0.0024871826171875\n", + "Round 103: bitdecode vs pytorch: 0.0024852752685546875\n", + "Round 104: bitdecode vs pytorch: 0.0024814605712890625\n", + "Round 105: bitdecode vs pytorch: 0.00247955322265625\n", + "Round 106: bitdecode vs pytorch: 0.0024776458740234375\n", + "Round 107: bitdecode vs pytorch: 0.002475738525390625\n", + "Round 108: bitdecode vs pytorch: 0.0024738311767578125\n", + "Round 109: bitdecode vs pytorch: 0.0024700164794921875\n", + "Round 110: bitdecode vs pytorch: 0.002468109130859375\n", + "Round 111: bitdecode vs pytorch: 0.0024662017822265625\n", + "Round 112: bitdecode vs pytorch: 0.00246429443359375\n", + "Round 113: bitdecode vs pytorch: 0.0024623870849609375\n", + "Round 114: bitdecode vs pytorch: 0.002460479736328125\n", + "Round 115: bitdecode vs pytorch: 0.0024566650390625\n", + "Round 116: bitdecode vs pytorch: 0.0024547576904296875\n", + "Round 117: bitdecode vs pytorch: 0.002452850341796875\n", + "Round 118: bitdecode vs pytorch: 0.0024509429931640625\n", + "Round 119: bitdecode vs pytorch: 0.00244903564453125\n", + "Round 120: bitdecode vs pytorch: 0.0024471282958984375\n", + "Round 121: bitdecode vs pytorch: 0.002445220947265625\n", + "Round 122: bitdecode vs pytorch: 0.00244140625\n", + "Round 123: bitdecode vs pytorch: 0.0024394989013671875\n", + "Round 124: bitdecode vs pytorch: 0.002437591552734375\n", + "Round 125: bitdecode vs pytorch: 0.0024356842041015625\n", + "Round 126: bitdecode vs pytorch: 0.00243377685546875\n", + "Round 127: bitdecode vs pytorch: 0.0024318695068359375\n", + "Round 128: bitdecode vs pytorch: 0.002429962158203125\n", + "Round 129: bitdecode vs pytorch: 0.0024280548095703125\n", + "Round 130: bitdecode vs pytorch: 0.0025577545166015625\n", + "Round 131: bitdecode vs pytorch: 0.00255584716796875\n", + "Round 132: bitdecode vs pytorch: 0.0025539398193359375\n", + "Round 133: bitdecode vs pytorch: 0.002552032470703125\n", + "Round 134: bitdecode vs pytorch: 0.0025501251220703125\n", + "Round 135: bitdecode vs pytorch: 0.0025463104248046875\n", + "Round 136: bitdecode vs pytorch: 0.002544403076171875\n", + "Round 137: bitdecode vs pytorch: 0.00254058837890625\n", + "Round 138: bitdecode vs pytorch: 0.0025386810302734375\n", + "Round 139: bitdecode vs pytorch: 0.0025386810302734375\n", + "Round 140: bitdecode vs pytorch: 0.0025348663330078125\n", + "Round 141: bitdecode vs pytorch: 0.002532958984375\n", + "Round 142: bitdecode vs pytorch: 0.0025310516357421875\n", + "Round 143: bitdecode vs pytorch: 0.002529144287109375\n", + "Round 144: bitdecode vs pytorch: 0.0025272369384765625\n", + "Round 145: bitdecode vs pytorch: 0.0025234222412109375\n", + "Round 146: bitdecode vs pytorch: 0.002521514892578125\n", + "Round 147: bitdecode vs pytorch: 0.0025196075439453125\n", + "Round 148: bitdecode vs pytorch: 0.0025177001953125\n", + "Round 149: bitdecode vs pytorch: 0.0025157928466796875\n", + "Round 150: bitdecode vs pytorch: 0.002513885498046875\n", + "Round 151: bitdecode vs pytorch: 0.0025119781494140625\n", + "Round 152: bitdecode vs pytorch: 0.00251007080078125\n", + "Round 153: bitdecode vs pytorch: 0.002506256103515625\n", + "Round 154: bitdecode vs pytorch: 0.0025043487548828125\n", + "Round 155: bitdecode vs pytorch: 0.00250244140625\n", + "Round 156: bitdecode vs pytorch: 0.0025005340576171875\n", + "Round 157: bitdecode vs pytorch: 0.002498626708984375\n", + "Round 158: bitdecode vs pytorch: 0.00249481201171875\n", + "Round 159: bitdecode vs pytorch: 0.0024929046630859375\n", + "Round 160: bitdecode vs pytorch: 0.002490997314453125\n", + "Round 161: bitdecode vs pytorch: 0.0024890899658203125\n", + "Round 162: bitdecode vs pytorch: 0.0024871826171875\n", + "Round 163: bitdecode vs pytorch: 0.0024852752685546875\n", + "Round 164: bitdecode vs pytorch: 0.002483367919921875\n", + "Round 165: bitdecode vs pytorch: 0.0024814605712890625\n", + "Round 166: bitdecode vs pytorch: 0.00247955322265625\n", + "Round 167: bitdecode vs pytorch: 0.0024776458740234375\n", + "Round 168: bitdecode vs pytorch: 0.002475738525390625\n", + "Round 169: bitdecode vs pytorch: 0.0024738311767578125\n", + "Round 170: bitdecode vs pytorch: 0.002471923828125\n", + "Round 171: bitdecode vs pytorch: 0.0024700164794921875\n", + "Round 172: bitdecode vs pytorch: 0.002468109130859375\n", + "Round 173: bitdecode vs pytorch: 0.0024662017822265625\n", + "Round 174: bitdecode vs pytorch: 0.00246429443359375\n", + "Round 175: bitdecode vs pytorch: 0.0024623870849609375\n", + "Round 176: bitdecode vs pytorch: 0.002460479736328125\n", + "Round 177: bitdecode vs pytorch: 0.0024585723876953125\n", + "Round 178: bitdecode vs pytorch: 0.0024566650390625\n", + "Round 179: bitdecode vs pytorch: 0.002452850341796875\n", + "Round 180: bitdecode vs pytorch: 0.0024509429931640625\n", + "Round 181: bitdecode vs pytorch: 0.00244903564453125\n", + "Round 182: bitdecode vs pytorch: 0.0024471282958984375\n", + "Round 183: bitdecode vs pytorch: 0.002445220947265625\n", + "Round 184: bitdecode vs pytorch: 0.0024433135986328125\n", + "Round 185: bitdecode vs pytorch: 0.00244140625\n", + "Round 186: bitdecode vs pytorch: 0.0024394989013671875\n", + "Round 187: bitdecode vs pytorch: 0.002437591552734375\n", + "Round 188: bitdecode vs pytorch: 0.0024356842041015625\n", + "Round 189: bitdecode vs pytorch: 0.00243377685546875\n", + "Round 190: bitdecode vs pytorch: 0.0024318695068359375\n", + "Round 191: bitdecode vs pytorch: 0.002429962158203125\n", + "Round 192: bitdecode vs pytorch: 0.0024261474609375\n", + "Round 193: bitdecode vs pytorch: 0.0024261474609375\n", + "Round 194: bitdecode vs pytorch: 0.0024242401123046875\n", + "Round 195: bitdecode vs pytorch: 0.002422332763671875\n", + "Round 196: bitdecode vs pytorch: 0.0024204254150390625\n", + "Round 197: bitdecode vs pytorch: 0.00241851806640625\n", + "Round 198: bitdecode vs pytorch: 0.0024166107177734375\n", + "Round 199: bitdecode vs pytorch: 0.002414703369140625\n", + "Round 200: bitdecode vs pytorch: 0.0024127960205078125\n", + "Round 201: bitdecode vs pytorch: 0.002410888671875\n", + "Round 202: bitdecode vs pytorch: 0.0024089813232421875\n", + "Round 203: bitdecode vs pytorch: 0.0024089813232421875\n", + "Round 204: bitdecode vs pytorch: 0.002407073974609375\n", + "Round 205: bitdecode vs pytorch: 0.00240325927734375\n", + "Round 206: bitdecode vs pytorch: 0.0024013519287109375\n", + "Round 207: bitdecode vs pytorch: 0.002399444580078125\n", + "Round 208: bitdecode vs pytorch: 0.0023975372314453125\n", + "Round 209: bitdecode vs pytorch: 0.0023956298828125\n", + "Round 210: bitdecode vs pytorch: 0.0023937225341796875\n", + "Round 211: bitdecode vs pytorch: 0.002391815185546875\n", + "Round 212: bitdecode vs pytorch: 0.00238800048828125\n", + "Round 213: bitdecode vs pytorch: 0.0023860931396484375\n", + "Round 214: bitdecode vs pytorch: 0.002384185791015625\n", + "Round 215: bitdecode vs pytorch: 0.0023822784423828125\n", + "Round 216: bitdecode vs pytorch: 0.0023822784423828125\n", + "Round 217: bitdecode vs pytorch: 0.0023784637451171875\n", + "Round 218: bitdecode vs pytorch: 0.002376556396484375\n", + "Round 219: bitdecode vs pytorch: 0.0023746490478515625\n", + "Round 220: bitdecode vs pytorch: 0.00237274169921875\n", + "Round 221: bitdecode vs pytorch: 0.0023708343505859375\n", + "Round 222: bitdecode vs pytorch: 0.002368927001953125\n", + "Round 223: bitdecode vs pytorch: 0.0023670196533203125\n", + "Round 224: bitdecode vs pytorch: 0.0023651123046875\n", + "Round 225: bitdecode vs pytorch: 0.0023632049560546875\n", + "Round 226: bitdecode vs pytorch: 0.002361297607421875\n", + "Round 227: bitdecode vs pytorch: 0.0023593902587890625\n", + "Round 228: bitdecode vs pytorch: 0.00235748291015625\n", + "Round 229: bitdecode vs pytorch: 0.0023555755615234375\n", + "Round 230: bitdecode vs pytorch: 0.002353668212890625\n", + "Round 231: bitdecode vs pytorch: 0.0023517608642578125\n", + "Round 232: bitdecode vs pytorch: 0.002349853515625\n", + "Round 233: bitdecode vs pytorch: 0.0023479461669921875\n", + "Round 234: bitdecode vs pytorch: 0.002346038818359375\n", + "Round 235: bitdecode vs pytorch: 0.0023441314697265625\n", + "Round 236: bitdecode vs pytorch: 0.00234222412109375\n", + "Round 237: bitdecode vs pytorch: 0.0023403167724609375\n", + "Round 238: bitdecode vs pytorch: 0.002338409423828125\n", + "Round 239: bitdecode vs pytorch: 0.0023365020751953125\n", + "Round 240: bitdecode vs pytorch: 0.0023345947265625\n", + "Round 241: bitdecode vs pytorch: 0.0023345947265625\n", + "Round 242: bitdecode vs pytorch: 0.0023326873779296875\n", + "Round 243: bitdecode vs pytorch: 0.0023288726806640625\n", + "Round 244: bitdecode vs pytorch: 0.0023288726806640625\n", + "Round 245: bitdecode vs pytorch: 0.0023250579833984375\n", + "Round 246: bitdecode vs pytorch: 0.002323150634765625\n", + "Round 247: bitdecode vs pytorch: 0.0023212432861328125\n", + "Round 248: bitdecode vs pytorch: 0.0023193359375\n", + "Round 249: bitdecode vs pytorch: 0.0023174285888671875\n", + "Round 250: bitdecode vs pytorch: 0.002315521240234375\n", + "Round 251: bitdecode vs pytorch: 0.0023136138916015625\n" + ] + } + ], + "source": [ + "####### Round 1 : Prefill #######\n", + "torch.manual_seed(42)\n", + "\n", + "q = torch.rand(batch_size, seqlen_q, nheads, d, device=device, dtype=dtype)\n", + "k_state = torch.randn(batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype)\n", + "v_state = torch.randn(batch_size, seqlen_k, nheads_k, d, device=device, dtype=dtype)\n", + "\n", + "residual_len = seqlen_k % residual_block_size\n", + "residual = residual_len > 0\n", + "seqlen_k_pack = seqlen_k - residual_len\n", + "\n", + "cu_seqlens_k = torch.arange(0, (batch_size + 1) * seqlen_k_pack, seqlen_k_pack, \n", + " dtype=torch.int32, device=device)\n", + "\n", + "# Initialize quantization tensors\n", + "if quant_mode == \"k-channel\":\n", + " k_pack = torch.zeros((batch_size, int(seqlen_k_pack // pack_nums), nheads_k, d), dtype=torch.uint16, device=device)\n", + " k_params = torch.zeros((batch_size, int(seqlen_k_pack // group_size), nheads_k, d), dtype=torch.float32, device=device)\n", + "else:\n", + " k_pack = torch.zeros((batch_size, seqlen_k_pack, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device)\n", + " k_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seqlen_k_pack), dtype=torch.float32, device=device)\n", + "\n", + "v_pack = torch.zeros((batch_size, seqlen_k_pack, nheads_k, int(d // pack_nums)), dtype=torch.uint16, device=device)\n", + "v_params = torch.zeros((batch_size, int(d // group_size), nheads_k, seqlen_k_pack), dtype=torch.float32, device=device)\n", + "\n", + "# KV Cache Dynamic Cache\n", + "past_key_value = DynamicCache()\n", + "\n", + "if residual:\n", + " k_state_residual = k_state[:, -residual_len:, :, :]\n", + " v_state_residual = v_state[:, -residual_len:, :, :]\n", + " k_state_past = k_state[:, :-residual_len, :, :]\n", + " v_state_past = v_state[:, :-residual_len, :, :]\n", + " past_key_value.update_residual(k_state_residual, v_state_residual, layer_idx)\n", + "else:\n", + " k_state_past = k_state\n", + " v_state_past = v_state\n", + "\n", + "kvcache_pack_int(\n", + " k_state_past, k_pack, k_params,\n", + " v_state_past, v_pack, v_params,\n", + " None, # opt_block_table\n", + " cu_seqlens_k, \n", + " seqlen_k_pack,\n", + " quant_mode,\n", + " group_size,\n", + " num_bits\n", + ")\n", + "past_key_value.update_pack(k_pack, k_params, v_pack, v_params, layer_idx)\n", + "\n", + "# self\n", + "k_pack_new = torch.empty((batch_size, int(residual_block_size // pack_nums), nheads_k, k_pack.size(-1)), dtype=torch.uint16, device=device)\n", + "k_params_new = torch.empty((batch_size, int(residual_block_size // group_size), nheads_k, k_params.size(-1)), dtype=torch.float32, device=device)\n", + "v_pack_new = torch.empty((batch_size, residual_block_size, nheads_k, v_pack.size(-1)), dtype=torch.uint16, device=device)\n", + "v_params_new = torch.empty((batch_size, v_params.size(1), nheads_k, residual_block_size), dtype=torch.float32, device=device)\n", + "\n", + "####### Round 2-3 : Decode #######\n", + "for round_idx in range(250):\n", + " k_new = torch.randn(batch_size, 1, nheads_k, d, device=device, dtype=dtype)\n", + " v_new = torch.randn(batch_size, 1, nheads_k, d, device=device, dtype=dtype)\n", + "\n", + " # Get kv cache_pack\n", + " k_pack, k_params, v_pack, v_params = past_key_value.update_pack(None, None, None, None, layer_idx)\n", + "\n", + " seqlen_pack = v_pack.shape[1]\n", + " seqlens_k = torch.full((batch_size,), seqlen_pack, dtype=torch.int32, device=device)\n", + "\n", + " # Get kv cache_residual and append new kv\n", + " k_residual = torch.zeros((batch_size, residual_block_size, nheads_k, d), device=device, dtype=dtype)\n", + " v_residual = torch.zeros((batch_size, residual_block_size, nheads_k, d), device=device, dtype=dtype)\n", + " k_residual_cache, v_residual_cache = past_key_value.update_residual(k_new, v_new, layer_idx)\n", + "\n", + " cur_residual_len = k_residual_cache.shape[1]\n", + "\n", + " k_residual[:, :cur_residual_len, :, :] = k_residual_cache\n", + " v_residual[:, :cur_residual_len, :, :] = v_residual_cache\n", + "\n", + " out_bitdecode, k_pack_new, k_params_new, v_pack_new, v_params_new = fwd_kvcache_int(\n", + " q,\n", + " k_pack, k_params, \n", + " v_pack, v_params,\n", + " k_residual, v_residual, seqlens_k, #seqlens_k\n", + " k_pack_new, k_params_new, v_pack_new, v_params_new,\n", + " None, # opt_block_table\n", + " sm_scale,\n", + " quant_mode, \n", + " group_size,\n", + " residual_block_size,\n", + " cur_residual_len, # new_lens\n", + " num_bits\n", + " )\n", + "\n", + " if cur_residual_len == residual_block_size:\n", + " past_key_value.update_pack(k_pack_new, k_params_new, v_pack_new, v_params_new, layer_idx)\n", + " past_key_value.clear_residual(layer_idx)\n", + "\n", + " k_state = torch.cat([k_state, k_new], dim=1)\n", + " v_state = torch.cat([v_state, v_new], dim=1)\n", + "\n", + " out_ref = attention_ref(q, k_state, v_state)[0]\n", + " print(f\"Round {round_idx+2}: bitdecode vs pytorch: {(out_bitdecode - out_ref).abs().mean().item()}\")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "bitdecode", + "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.10.18" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/install.sh b/install.sh new file mode 100644 index 0000000..7bdb4d4 --- /dev/null +++ b/install.sh @@ -0,0 +1,5 @@ +rm -rf bit_decode.egg-info build/ bit_decode_cuda.cpython-310-x86_64-linux-gnu.so dist/ + +python setup.py install + +# pip install -e . \ No newline at end of file