kvpress_decoding_aime25.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d5b82a94ac3477fd",
   "metadata": {},
   "source": [
    "# Using decoding compression on the AIME25 Math Dataset\n",
    "\n",
    "This notebook demonstrates how to compress during text generation.\n",
    "We use `nvidia/OpenMath-Nemotron-7B` to solve math problems from the AIME25 dataset. For each problem, the model generates an answer in a boxed format (e.g., `\\boxed{42}`).\n",
    "\n",
    "To optimize memory usage during long-context generation, the notebook applies key-value cache compression during decoding.\n",
    "Compression periodically reduces the cache size by keeping only the most relevant tokens, enabling efficient inference without sacrificing answer quality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cd1c9e43c5ca1bf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/data/projects/kvpress/.venv/lib/python3.12/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": [
    "from transformers import pipeline\n",
    "from datasets import load_dataset\n",
    "\n",
    "from kvpress import (\n",
    "    DecodingPress,\n",
    "    ExpectedAttentionPress,\n",
    "    KnormPress,\n",
    "    ObservedAttentionPress,\n",
    "    RandomPress,\n",
    "    SnapKVPress,\n",
    "    StreamingLLMPress,\n",
    "    TOVAPress,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "553f7d7f-7a2f-456d-a4b1-f38f3ead8767",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'problem': 'Find the sum of all integer bases $b>9$ for which $17_b$ is a divisor of $97_b.$',\n",
       " 'answer': '70',\n",
       " 'id': '0'}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load a sample\n",
    "dataset = load_dataset(\"math-ai/aime25\")\n",
    "sample = dataset[\"test\"][0]\n",
    "sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "02b0af83-e71a-4129-9692-921d2bc16cb8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████| 2/2 [00:00<00:00, 64.85it/s]\n",
      "Device set to use cuda:0\n"
     ]
    }
   ],
   "source": [
    "# Load the pipeline\n",
    "device = \"cuda:0\"\n",
    "ckpt = \"nvidia/OpenMath-Nemotron-7B\"\n",
    "attn_implementation = \"flash_attention_2\"\n",
    "pipe = pipeline(\"kv-press-text-generation\", model=ckpt, device=device, dtype=\"auto\", model_kwargs={\"attn_implementation\":attn_implementation})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "96e1f430-4ece-49dd-a5fb-68ae3c2ee8b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question:   Find the sum of all integer bases $b>9$ for which $17_b$ is a divisor of $97_b.$\n",
      "Answer:     70\n",
      "Prediction: <think>\n",
      "Okay, so I need to find all integer bases b greater than 9 where the number 17 in base b divides the number 97 in base b. Then, sum all those bases. Hmm, let me think step by step.\n",
      "\n",
      "First, I should recall how numbers are represented in different bases. The number 17 in base b is equal to 1 times b plus 7 times 1, right? So that's 1*b + 7. Similarly, 97 in base b would be 9*b + 7. So, translating both numbers to base 10, we have:\n",
      "\n",
      "17_b = b + 7\n",
      "\n",
      "97_b = 9b + 7\n",
      "\n",
      "The problem states that 17_b must divide 97_b. So, in mathematical terms, this means that (9b + 7) divided by (b + 7) should result in an integer. So, (9b + 7) must be divisible by (b + 7). \n",
      "\n",
      "Let me write that as an equation:\n",
      "\n",
      "(9b + 7) ÷ (b + 7) = integer.\n",
      "\n",
      "To find when this division results in an integer, maybe I can perform the division and see what the remainder is. If the remainder is zero, then it's divisible. Let's try polynomial division or maybe manipulate the expression.\n",
      "\n",
      "Let me rewrite 9b + 7. Let's see, if I write 9b + 7 as 9*(b + 7) minus something. Let's compute:\n",
      "\n",
      "9*(b + 7) = 9b + 63\n",
      "\n",
      "But 9b + 7 is the original numerator. So, subtracting 9*(b + 7) from 9b +7 gives:\n",
      "\n",
      "9b +7 - (9b +63) = 7 - 63 = -56\n",
      "\n",
      "Therefore, 9b +7 = 9*(b +7) -56\n",
      "\n",
      "So, (9b +7)/(b +7) = 9 - 56/(b +7)\n",
      "\n",
      "For this to be an integer, 56/(b +7) must be an integer. Therefore, (b +7) must be a divisor of 56. \n",
      "\n",
      "But since b is a base greater than 9, the digits in the numbers 17_b and 97_b must be valid. In base b, the digits can go from 0 to b-1. So, in 17_b, the digits are 1 and 7. Since the base is greater than 9, 7 is a valid digit. Similarly, in 97_b, the digits are 9 and 7. So, 9 must be less than b. Therefore, the base b must be greater than 9, which is already given. So, the constraints are:\n",
      "\n",
      "1. b > 9\n",
      "\n",
      "2. (b +7) divides 56.\n",
      "\n",
      "So, first, let's find all divisors of 56. The positive divisors of 56 are:\n",
      "\n",
      "1, 2, 4, 7, 8, 14, 28, 56.\n",
      "\n",
      "But since b +7 must be one of these divisors, and b >9, then b +7 must be a divisor of 56 greater than 9 +7 =16. Wait, because b >9, so b >=10, so b +7 >=17. Therefore, the divisors of 56 that are greater than or equal to 17 are 28 and 56. Wait, let's check:\n",
      "\n",
      "Divisors of 56: 1, 2, 4, 7, 8, 14, 28, 56.\n",
      "\n",
      "So, divisors >=17 are 28 and 56. Therefore, b +7 can be 28 or 56. Therefore, solving for b:\n",
      "\n",
      "If b +7 =28, then b=21\n",
      "\n",
      "If b +7=56, then b=49\n",
      "\n",
      "So, the possible bases are 21 and 49. Then, the sum of these bases is 21 +49=70.\n",
      "\n",
      "Wait, but let me check if there are any other divisors. Wait, 14 is a divisor, but 14 is less than 17 (since b >=10, so b +7 >=17). So, 14 is too small. Similarly, 8,7, etc. So, only 28 and 56. Therefore, the possible bases are 21 and 49. So, sum is 70.\n",
      "\n",
      "But wait, let me verify this. Let's check for b=21:\n",
      "\n",
      "17 in base 21 is 1*21 +7=28\n",
      "\n",
      "97 in base21 is 9*21 +7=196\n",
      "\n",
      "196 divided by28 is 7, which is an integer. So that works.\n",
      "\n",
      "For b=49:\n",
      "\n",
      "17 in base49 is 1*49 +7=56\n",
      "\n",
      "97 in base49 is 9*49 +7=448\n",
      "\n",
      "448 divided by56 is 8, which is an integer. So that works too.\n",
      "\n",
      "So, both bases are valid, and their sum is 21 +49=70.\n",
      "\n",
      "Wait, but hold on. Let me check if there are any negative divisors. Since 56 is positive, but divisors can be negative. However, since b is a base, it must be a positive integer greater than 9. So, b +7 must be positive. Therefore, we can ignore negative divisors. So, only positive divisors. So, 28 and 56 are the only ones. Therefore, the answer is 70.\n",
      "\n",
      "But let me just make sure I didn't miss any divisors. Let's list all divisors again:\n",
      "\n",
      "1, 2, 4, 7, 8, 14, 28, 56. So, yes, only 28 and 56 are >=17. So, that's correct.\n",
      "\n",
      "Therefore, the sum is 21 +49=70. So, the answer is 70.\n",
      "\n",
      "**Final Answer**\n",
      "\\boxed{70}\n",
      "</think>To solve the problem, we need to find all integer bases \\( b > 9 \\) for which \\( 17_b \\) is a divisor of \\( 97_b \\).\n",
      "\n",
      "First, we convert the numbers from base \\( b \\) to base 10:\n",
      "- \\( 17_b = 1 \\cdot b + 7 = b + 7 \\)\n",
      "- \\( 97_b = 9 \\cdot b + 7 = 9b + 7 \\)\n",
      "\n",
      "We need \\( 9b + 7 \\) to be divisible by \\( b + 7 \\). This can be expressed as:\n",
      "\\[\n",
      "\\frac{9b + 7}{b + 7} = 9 - \\frac{56}{b + 7}\n",
      "\\]\n",
      "For this to be an integer, \\( \\frac{56}{b + 7} \\) must be an integer. Therefore, \\( b + 7 \\) must be a divisor of 56.\n",
      "\n",
      "The positive divisors of 56 are:\n",
      "\\[\n",
      "1, 2, 4, 7, 8, 14, 28, 56\n",
      "\\]\n",
      "Since \\( b > 9 \\), we need \\( b + 7 \\geq 17 \\). The divisors of 56 that are greater than or equal to 17 are 28 and 56.\n",
      "\n",
      "Solving for \\( b \\):\n",
      "- If \\( b + 7 = 28 \\), then \\( b = 21 \\)\n",
      "- If \\( b + 7 = 56 \\), then \\( b = 49 \\)\n",
      "\n",
      "We verify these bases:\n",
      "- For \\( b = 21 \\):\n",
      "  \\[\n",
      "  17_{21} = 1 \\cdot 21 + 7 = 28\n",
      "  \\]\n",
      "  \\[\n",
      "  97_{21} = 9 \\cdot 21 + 7 = 196\n",
      "  \\]\n",
      "  \\[\n",
      "  196 \\div 28 = 7 \\quad \\text{(an integer)}\n",
      "  \\]\n",
      "\n",
      "- For \\( b = 49 \\):\n",
      "  \\[\n",
      "  17_{49} = 1 \\cdot 49 + 7 = 56\n",
      "  \\]\n",
      "  \\[\n",
      "  97_{49} = 9 \\cdot 49 + 7 = 448\n",
      "  \\]\n",
      "  \\[\n",
      "  448 \\div 56 = 8 \\quad \\text{(an integer)}\n",
      "  \\]\n",
      "\n",
      "Both bases are valid. The sum of these bases is:\n",
      "\\[\n",
      "21 + 49 = 70\n",
      "\\]\n",
      "\n",
      "Thus, the final answer is:\n",
      "\\[\n",
      "\\boxed{70}\n",
      "\\]\n"
     ]
    }
   ],
   "source": [
    "# Run the pipeline without compression\n",
    "\n",
    "question = sample[\"problem\"]\n",
    "true_answer = sample[\"answer\"]\n",
    "pred_answer = pipe(\" \", question=question, press=None, max_new_tokens=16_000)[\"answer\"]\n",
    "\n",
    "print(f\"Question:   {question}\")\n",
    "print(f\"Answer:     {true_answer}\")\n",
    "print(f\"Prediction: {pred_answer}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7f4fa366-7e62-443a-b5c8-274128fe6237",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question:   Find the sum of all integer bases $b>9$ for which $17_b$ is a divisor of $97_b.$\n",
      "Answer:     70\n",
      "Prediction: <think>\n",
      "Okay, so I need to find all integer bases b greater than 9 where the number 17 in base b divides the number 97 in base b. Then, sum all those bases. Hmm, let me think step by step.\n",
      "\n",
      "First, I should recall how numbers are represented in different bases. The number 17 in base b is equal to 1 times b plus 7 times 1, right? So that's 1*b + 7. Similarly, 97 in base b would be 9*b + 7. So, translating both numbers to base 10, we have:\n",
      "\n",
      "17_b = b + 7\n",
      "\n",
      "97_b = 9b + 7\n",
      "\n",
      "The problem states that 17_b must divide 97_b. So, in mathematical terms, this means that (9b + 7) divided by (b + 7) should result in an integer. So, (9b + 7) must be divisible by (b + 7). \n",
      "\n",
      "Let me write that as an equation:\n",
      "\n",
      "(9b + 7) ÷ (b + 7) = integer.\n",
      "\n",
      "To find when this division results in an integer, maybe I can perform the division and see what the remainder is. If the remainder is zero, then it's divisible. Let's try polynomial division or maybe manipulate the expression.\n",
      "\n",
      "Let me rewrite 9b + 7. Let's see, if I write 9b + 7 as 9*(b + 7) minus something. Let's compute:\n",
      "\n",
      "9*(b + 7) = 9b + 63\n",
      "\n",
      "But 9b + 7 is the original numerator. So, subtracting 9*(b + 7) from 9b +7 gives:\n",
      "\n",
      "9b +7 - (9b +63) = 7 - 63 = -56\n",
      "\n",
      "Therefore, 9b +7 = 9*(b +7) -56\n",
      "\n",
      "So, (9b +7)/(b +7) = 9 - 56/(b +7)\n",
      "\n",
      "For this to be an integer, 56/(b +7) must be an integer. Therefore, (b +7) must be a divisor of 56. \n",
      "\n",
      "But since b is a base greater than 9, the digits in the numbers 17_b and 97_b must be valid. In base b, the digits can go from 0 to b-1. So, in 17_b, the digits are 1 and 7. Since the base is greater than 9, 7 is a valid digit. Similarly, in 97_b, the digits are 9 and 7. So, 9 must be less than b. Therefore, the base b must be greater than 9, which is already given. So, the constraints are:\n",
      "\n",
      "1. b > 9\n",
      "\n",
      "2. (b +7) divides 56.\n",
      "\n",
      "So, first, let's find all divisors of 56. The positive divisors of 56 are:\n",
      "\n",
      "1, 2, 4, 7, 8, 14, 28, 56.\n",
      "\n",
      "But since b +7 must be one of these divisors, and b >9, then b +7 must be a divisor of 56 greater than 9 +7 =16. Wait, because b >9, so b >=10, so b +7 >=17. Therefore, the divisors of 56 that are greater than or equal to 17 are 28 and 56. Wait, let's check:\n",
      "\n",
      "Divisors of 56: 1, 2, 4, 7, 8, 14, 28, 56.\n",
      "\n",
      "So, divisors >=17 are 28 and 56. Therefore, b +7 can be 28 or 56. Therefore, solving for b:\n",
      "\n",
      "If b +7 =28, then b=21\n",
      "\n",
      "If b +7=56, then b=49\n",
      "\n",
      "So, the possible bases are 21 and 49. Then, the sum of these bases is 21 +49=70.\n",
      "\n",
      "Wait, but let me check if there are any other divisors. Wait, 14 is a divisor, but 14 is less than 17 (since b >=10, so b +7 >=17). So, 14 is too small. Similarly, 8,7, etc. So, only 28 and 56. Therefore, the possible bases are 21 and 49. So, sum is 70.\n",
      "\n",
      "But wait, let me verify this. Let's check for b=21:\n",
      "\n",
      "17 in base 21 is 2*21 +7=49. Wait, no, 17 in base 21 is 1*21 +7=28. 97 in base 21 is 9*21 +7=196. Then, 196 divided by 28 is 7, which is an integer. So that works.\n",
      "\n",
      "For b=49: 17 in base 49 is 1*49 +7=56. 97 in base 49 is 9*49 +7=448. 448 divided by 56 is 8, which is an integer. So that works too.\n",
      "\n",
      "So, both bases 21 and 49 are valid. Therefore, the sum is 21 +49=70.\n",
      "\n",
      "Wait, but hold on. Let me check if there are any other divisors. Wait, 56 is a divisor of 56, but 56 is 56. So, if b +7=56, then b=49. Which we already have. Similarly, 28 gives b=21. Are there any negative divisors? But since b is a base, it must be a positive integer greater than 9, so negative divisors don't make sense here. So, only 28 and 56. Therefore, the answer is 70.\n",
      "\n",
      "But wait, let me check if there's a mistake here. Let me think again. The problem says \"all bases b>9\". So, the possible divisors of 56 are 1,2,4,7,8,14,28,56. But since b +7 must be one of these, and b>9, so b +7 must be at least 17. So, the possible divisors are 28 and 56. Therefore, b=21 and 49. So, sum is 70. That seems correct.\n",
      "\n",
      "But let me check if there's another way to approach this. For example, maybe using modular arithmetic. Let's see.\n",
      "\n",
      "We have 9b +7 ≡0 mod (b +7). Let me express 9b +7 in terms of (b +7). Let's write 9b +7 =9*(b +7) - 9*7 +7=9*(b +7) -63 +7=9*(b +7) -56. Therefore, 9b +7 ≡ -56 mod (b +7). Therefore, -56 ≡0 mod (b +7), which implies that (b +7) divides 56. So, same result as before. Therefore, the same answer.\n",
      "\n",
      "Therefore, the possible values of b are 21 and 49, sum is 70. So, the answer is 70. Therefore, the sum is 70. So, I think that's correct.\n",
      "\n",
      "**Final Answer**\n",
      "The sum of all such bases is \\boxed{70}.\n",
      "</think>To solve the problem, we need to find all integer bases \\( b > 9 \\) such that \\( 17_b \\) divides \\( 97_b \\).\n",
      "\n",
      "First, we convert the numbers from base \\( b \\) to base 10:\n",
      "- \\( 17_b \\) in base 10 is \\( 1 \\cdot b + 7 = b + 7 \\).\n",
      "- \\( 97_b \\) in base 10 is \\( 9 \\cdot b + 7 = 9b + 7 \\).\n",
      "\n",
      "We need \\( 9b + 7 \\) to be divisible by \\( b + 7 \\). This can be expressed as:\n",
      "\\[\n",
      "9b + 7 \\equiv 0 \\pmod{b + 7}\n",
      "\\]\n",
      "\n",
      "Rewriting \\( 9b + 7 \\) in terms of \\( b + 7 \\):\n",
      "\\[\n",
      "9b + 7 = 9(b + 7) - 56\n",
      "\\]\n",
      "Thus, we have:\n",
      "\\[\n",
      "9b + 7 \\equiv -56 \\pmod{b + 7}\n",
      "\\]\n",
      "This implies that \\( b + 7 \\) must be a divisor of 56. The positive divisors of 56 are:\n",
      "\\[\n",
      "1, 2, 4, 7, 8, 14, 28, 56\n",
      "\\]\n",
      "\n",
      "Since \\( b > 9 \\), \\( b + 7 \\) must be at least 17. The divisors of 56 that are greater than or equal to 17 are:\n",
      "\\[\n",
      "28 \\quad \\text{and} \\quad 56\n",
      "\\]\n",
      "\n",
      "Solving for \\( b \\):\n",
      "- If \\( b + 7 = 28 \\), then \\( b = 21 \\).\n",
      "- If \\( b + 7 = 56 \\), then \\( b = 49 \\).\n",
      "\n",
      "Thus, the valid bases are \\( b = 21 \\) and \\( b = 49 \\).\n",
      "\n",
      "Summing these bases:\n",
      "\\[\n",
      "21 + 49 = 70\n",
      "\\]\n",
      "\n",
      "Therefore, the sum of all valid bases \\( b \\) is:\n",
      "\\[\n",
      "\\boxed{70}\n",
      "\\]\n"
     ]
    }
   ],
   "source": [
    "# Run the pipeline with compression\n",
    "\n",
    "compression_interval = 1024  # compress every compression_steps\n",
    "target_size = 512 # number of tokens to keep after compression. Note that actual cache size lies in [target_size, compression_interval]\n",
    "\n",
    "press = DecodingPress(base_press=ExpectedAttentionPress(), compression_interval=compression_interval, target_size=target_size)\n",
    "\n",
    "question = sample[\"problem\"]\n",
    "true_answer = sample[\"answer\"]\n",
    "pred_answer = pipe(\" \", question=question, press=press, max_new_tokens=16_000)[\"answer\"]\n",
    "\n",
    "print(f\"Question:   {question}\")\n",
    "print(f\"Answer:     {true_answer}\")\n",
    "print(f\"Prediction: {pred_answer}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}