Measure the KV cache and sampling with MiniMind's generate

Goal

With MiniMind's generate, check whether greedy generation gives the same tokens with and without the KV cache, match the tokens fed into the forward to the formula, and measure the time. Match the cache's actual bytes to the formula, and count with the SFT model what temperature and top-p change.

Why it matters

The cache decides generation cost. Without a cache, every new token recomputes everything so far, and with a cache only one token is computed. The difference grows with the square of the length. In exchange, the cache eats memory, and its size is computed exactly from the configuration (layers, KV heads, head_dim) — in serving, how many users one card can take is this number. How you pick is separate from the cache. The cache is "the same result cheaply," and sampling is "which result." If you mix the two in your thinking, you make wrong diagnoses such as "I turned on the cache and the answer changed."

Steps

1. With the reference pretrained model, generate 64 tokens greedily with and without the cache after a prompt made from the first 6 validation documents joined and cut to 200 tokens (with bos in front), and write whether they are the same to /root/mm/infer/same.json.
2. Count the tokens fed into the forward in the same two generations and write it to /root/mm/infer/count.json.
3. Write the generation time of the two methods (measure three times after one warm-up and take the median) to /root/mm/infer/time.json.
4. Measure the KV cache size after feeding 100 tokens, match it to the formula, and write it to /root/mm/infer/kvbytes.json.
5. Have the SFT model answer the question "What is Garam village's specialty?" (the question in Korean) 30 times each at temperatures 0.3, 1.0, and 1.5, and write the number of distinct answers to /root/mm/infer/temp.json.
6. From the next-token distribution after feeding [bos] 민수는 (the Korean word means "Minsu, as the topic") to the pretrained model, write the number of candidates that survive at p=0.5, 0.9, and 0.99 under MiniMind's top-p rule to /root/mm/infer/topp.json.
7. In /root/mm/infer/report.md, write the three sections ## KV 캐시, ## 온도, and ## top-p (the Korean headings mean "KV cache", "Temperature", and "top-p"), and include the speedup from step 3 and the cache bytes from step 4.

Notes

• model.generate(ids, max_new_tokens=…, do_sample=False, top_k=0, top_p=1.0, eos_token_id=None, use_cache=True|False) — MiniMind's generate has a top_k default of 50, so in the greedy and temperature experiments you must turn it off with top_k=0 to keep the conditions clean.
• You can count the tokens fed into the forward with a function that wraps model.forward (generate calls self.forward).
• Times vary with the node and how busy it is. Grading is by token counts, and for time it only checks that the cache side is faster.
• Common mistakes: measuring the time of the first run (warm-up) as it is, not resetting the random seed for each temperature, and computing the cache bytes at the size after repeat_kv.
• Sources: model_minimind.py — generate · eval_llm.py · nucleus sampling paper

The same answer with or without the cache

Write the script /root/mm/infer/cache.py, load /opt/mm/ref/pretrain.pth, take the first 200 tokens of the tokens cut from the text of the first 6 documents of /opt/mm/data/pretrain_val.jsonl joined with spaces, attach bos(1) to make the prompt, and generate 64 tokens greedily (do_sample=False, top_k=0, top_p=1.0, eos_token_id=None) twice, with use_cache=True and False. Write whether the new tokens are the same (same_ids) and the cache side's new tokens (ids) to /root/mm/infer/same.json together with prompt_len and new_tokens.

Thanks to the causal mask, the K and V of earlier tokens do not change when tokens are appended after them. So the result is the same when you pull them from the cache. If it differs here, it is not the cache but the conditions (sampling, random numbers) that differ.

Count the tokens fed into the forward

Add up the lengths of the input_ids fed into model.forward in the two generations of step 1, and write it to /root/mm/infer/count.json as prompt_len, new_tokens, fed_with_cache, and fed_without_cache.

With the cache, the first step feeds the whole prompt and from then on one token at a time. Without it, every step feeds everything so far again. The grader compares the two values you wrote with values computed from these two formulas.

In terms of time

Generate 64 tokens with and without the cache, measure each three times after one warm-up, and write the median to /root/mm/infer/time.json as with_cache_s, without_cache_s, and speedup (without ÷ with).

Measure with time.perf_counter(). Exclude the first run because it is slow due to memory allocation and preparation. The token counts differ by more than 50 times, but the time does not differ by that much — because in a small model the fixed cost per step (Python and kernel calls) is large.

How many bytes is the cache

Feed 100 random tokens (torch.randint(3, 어휘, (1, 100), generator=시드 0), where the Korean words stand for the vocabulary size and seed 0) into the reference pretrained model with use_cache=True, and write the byte sum of all tensors in the returned past_key_values (measured_bytes), the value computed by the formula (formula_bytes = 2 × layers × tokens × KV heads × head_dim × 4), and the shape of the first layer's K (k_shape) to /root/mm/infer/kvbytes.json together with tokens.

MiniMind's Attention creates past_kv = (xk, xv) before repeat_kv. So K's shape is (batch, length, 2 KV heads, 32), and the cache is as small as GQA reduced it.

When you raise the temperature

Feed the question "What is Garam village's specialty?" (the question in Korean) into the SFT model (/opt/mm/ref/sft.pth) in the chat format (mmkit.chat_text(…, add_generation_prompt=True)), set torch.manual_seed(0) at each temperature of 0.3, 1.0, and 1.5, generate with do_sample=True, top_k=0, top_p=1.0, max_new_tokens=24, num_return_sequences=30, and write the number of distinct answers, cut before <|im_end|>, to /root/mm/infer/temp.json as {"0.3": {"distinct": n}, "1.0": …, "1.5": …}.

Starting each temperature from the same random numbers lets you see only the effect of the temperature. At temperature 1.5, almost all thirty are different and some answers are garbled — once a rare token is picked, a small model cannot get back on track.

The candidates top-p leaves

From the logits at the last position after feeding [bos] + '민수는' (the Korean word means "Minsu, as the topic") into the reference pretrained model, write the number of candidates that survive at p=0.5, 0.9, and 0.99 under the top-p rule of MiniMind's generate (cut from the position where the cumulative sum in descending probability exceeds p, but shift by one slot to keep the token that first exceeds p, and always keep the top one) to /root/mm/infer/topp.json as {"0.5": n, "0.9": n, "0.99": n}.

You can copy the three top_p lines from generate in /opt/minimind/model/model_minimind.py as they are. After the word for "Minsu, as the topic," there are many branches such as verbs and places, so many candidates remain. Compare with a sharp position (for example, after the phrase meaning "the village's specialty is").

A record of inference cost and picking

In /root/mm/infer/report.md, write the three sections ## KV 캐시, ## 온도, and ## top-p (the Korean headings mean "KV cache", "Temperature", and "top-p"), and include the speedup from step 3 and the measured_bytes from step 4 as numbers.

Write side by side how many times the cache helped in token counts and in time, and add one line on the reason for the difference.