QKV：
Transformer
计算图谱

kv_lab.py · 与上方保存按钮的内容一致

#!/usr/bin/env python3
"""Small, offline KV-cache laboratory. This is NOT a pretrained language model.

Run: python kv_lab.py --output results.json
Requires Python >=3.10 and NumPy. No downloads, credentials, or API requests.
The model is a randomly initialized, 3-layer causal decoder with GQA and RoPE.
Its purpose is to test state reuse, not language ability or production speed.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import math
import platform
import sys
from collections import OrderedDict
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

Array = np.ndarray
State = list[tuple[Array, Array]]


def softmax(x: Array) -> Array:
    e = np.exp(x - x.max(axis=-1, keepdims=True))
    return e / e.sum(axis=-1, keepdims=True)


def rmsnorm(x: Array) -> Array:
    return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + 1e-6)


def rotate(x: Array, positions: Array) -> Array:
    """x: [heads, tokens, head_dim]; interleaved even/odd RoPE pairs."""
    d = x.shape[-1]
    freq = 10000.0 ** (-np.arange(0, d, 2, dtype=np.float64) / d)
    angles = positions[:, None] * freq[None, :]
    c, s = np.cos(angles).astype(x.dtype), np.sin(angles).astype(x.dtype)
    y = np.empty_like(x)
    y[..., 0::2] = x[..., 0::2] * c - x[..., 1::2] * s
    y[..., 1::2] = x[..., 0::2] * s + x[..., 1::2] * c
    return y


@dataclass(frozen=True)
class Config:
    vocab: int = 41
    width: int = 32
    layers: int = 3
    q_heads: int = 4
    kv_heads: int = 2
    head_dim: int = 8
    seed: int = 7
    dtype: str = "float64"


class TinyDecoder:
    def __init__(self, cfg: Config = Config()):
        if (cfg.width != cfg.q_heads * cfg.head_dim or
            cfg.q_heads % cfg.kv_heads or cfg.head_dim % 2):
            raise ValueError("Invalid head dimensions or GQA grouping")
        self.cfg = cfg
        dtype = np.dtype(cfg.dtype)
        if dtype not in (np.dtype('float32'), np.dtype('float64')):
            raise ValueError("This lab supports float32 and float64 only")
        rng = np.random.default_rng(cfg.seed)
        self.embed = (rng.normal(size=(cfg.vocab, cfg.width)) * 0.3).astype(dtype)
        shapes = [(cfg.width, cfg.q_heads * cfg.head_dim),
                  (cfg.width, cfg.kv_heads * cfg.head_dim),
                  (cfg.width, cfg.kv_heads * cfg.head_dim),
                  (cfg.width, cfg.width), (cfg.width, cfg.width * 2),
                  (cfg.width * 2, cfg.width)]
        self.weights = [[(rng.normal(size=sh) / math.sqrt(sh[0])).astype(dtype)
                         for sh in shapes] for _ in range(cfg.layers)]
        self.lm_head = (rng.normal(size=(cfg.width, cfg.vocab)) /
                        math.sqrt(cfg.width)).astype(dtype)
        h = hashlib.sha256(json.dumps(cfg.__dict__, sort_keys=True).encode())
        for a in [self.embed, self.lm_head] + [w for row in self.weights for w in row]:
            h.update(a.tobytes())
        # Includes actual weights, not just a model name. RoPE rule is fixed by code.
        self.identity = 'tiny-gqa-rope-v1:' + h.hexdigest()
        self.processed_positions = 0

    def empty_state(self) -> State:
        c = self.cfg
        return [(np.empty((c.kv_heads, 0, c.head_dim), dtype=c.dtype),
                 np.empty((c.kv_heads, 0, c.head_dim), dtype=c.dtype))
                for _ in range(c.layers)]

    def forward(self, tokens: list[int], past: State | None = None) -> tuple[Array, State]:
        """Processes ONLY tokens supplied here; past has already been computed.

        Token indices in the causal mask are ABSOLUTE, including the cached offset.
        This function returns new arrays; it never mutates a supplied cache.
        """
        c = self.cfg
        if not tokens or any(type(t) is not int or not 0 <= t < c.vocab for t in tokens):
            raise ValueError("Provide a nonempty list of integer IDs within the vocabulary")
        past = self.empty_state() if past is None else past
        if len(past) != c.layers:
            raise ValueError("Cache layer count mismatch")
        start = past[0][0].shape[1]
        for k, v in past:
            expected = (c.kv_heads, start, c.head_dim)
            if k.shape != expected or v.shape != expected:
                raise ValueError("Inconsistent cache shapes")
            if k.dtype != np.dtype(c.dtype) or v.dtype != np.dtype(c.dtype):
                raise ValueError("Cache dtype mismatch")
        pos = np.arange(start, start + len(tokens))
        x = self.embed[tokens].copy()
        out_state = []
        self.processed_positions += len(tokens)
        for layer, (wq, wk, wv, wo, w1, w2) in enumerate(self.weights):
            h = rmsnorm(x)
            q = (h @ wq).reshape(-1, c.q_heads, c.head_dim).transpose(1, 0, 2)
            k = (h @ wk).reshape(-1, c.kv_heads, c.head_dim).transpose(1, 0, 2)
            v = (h @ wv).reshape(-1, c.kv_heads, c.head_dim).transpose(1, 0, 2)
            q, k = rotate(q, pos), rotate(k, pos)
            k = np.concatenate([past[layer][0], k], axis=1)
            v = np.concatenate([past[layer][1], v], axis=1)
            out_state.append((k, v))  # Persist KV heads, not repeated Q heads.
            kr = np.repeat(k, c.q_heads // c.kv_heads, axis=0)
            vr = np.repeat(v, c.q_heads // c.kv_heads, axis=0)
            scores = (q @ kr.transpose(0, 2, 1)) / math.sqrt(c.head_dim)
            allowed = np.arange(k.shape[1])[None, :] <= pos[:, None]
            scores = np.where(allowed[None, :, :], scores, -np.inf)
            a = (softmax(scores) @ vr).transpose(1, 0, 2).reshape(-1, c.width)
            x = x + a @ wo
            z = rmsnorm(x) @ w1
            x = x + (z / (1 + np.exp(-z))) @ w2
        return rmsnorm(x) @ self.lm_head, out_state


@dataclass
class Entry:
    tokens: tuple[int, ...]  # Collision guard; each entry stores a full block only.
    state: State            # This block's K/V across all layers, never a whole prefix.


class PrefixCache:
    """Serial educational prefix cache. Hits are copied before model execution.

    Source blocks cannot be evicted midway because this implementation is serial.
    Production async transfers need ownership protection until the copy completes.
    This class stores neither final logits nor hidden states, so it leaves >=1 input.
    """
    def __init__(self, model: TinyDecoder, block_size: int = 16,
                 capacity_blocks: int = 128, namespace: str = "local-demo"):
        if block_size < 1 or capacity_blocks < 1:
            raise ValueError("Block size and capacity must be positive")
        self.model = model
        self.block_size = block_size
        self.capacity = capacity_blocks
        self.base = hashlib.sha256((model.identity + ':' + namespace).encode()).digest()
        self.blocks: OrderedDict[bytes, Entry] = OrderedDict()

    def keys(self, tokens: list[int]):
        parent = self.base
        b = self.block_size
        for end in range(b, len(tokens) + 1, b):
            block = tuple(tokens[end-b:end])
            raw = np.asarray(block, dtype='<i8').tobytes()
            key = hashlib.sha256(parent + raw).digest()
            yield end, key, block
            parent = key

    def lookup(self, tokens: list[int]) -> tuple[int, State | None]:
        rows: list[State] = []
        hit = 0
        for end, key, block in self.keys(tokens):
            if end >= len(tokens):
                break  # Need a final hidden state, not only K/V.
            entry = self.blocks.get(key)
            if entry is None or entry.tokens != block:
                break
            self.blocks.move_to_end(key)
            rows.append(entry.state)
            hit = end
        if not rows:
            return 0, None
        state = [(np.concatenate([r[i][0] for r in rows], axis=1).copy(),
                  np.concatenate([r[i][1] for r in rows], axis=1).copy())
                 for i in range(self.model.cfg.layers)]
        return hit, state

    def store(self, tokens: list[int], state: State) -> None:
        # Reverse insertion makes the beginning newer than the tail in this toy LRU.
        # It preserves a useful prefix when one old block must be evicted.
        for end, key, block in reversed(list(self.keys(tokens))):
            if key not in self.blocks:
                if len(self.blocks) >= self.capacity:
                    self.blocks.popitem(last=False)
                sliced = [(k[:, end-self.block_size:end].copy(),
                           v[:, end-self.block_size:end].copy()) for k, v in state]
                self.blocks[key] = Entry(block, sliced)
            self.blocks.move_to_end(key)

    def request(self, tokens: list[int]) -> tuple[Array, State, dict[str, int]]:
        hit, state = self.lookup(tokens)
        logits, state = self.model.forward(tokens[hit:], state)
        self.store(tokens, state)
        return logits[-1], state, {"input_tokens": len(tokens), "hit_tokens": hit,
                                   "computed_tokens": len(tokens)-hit}


def online_attention(q: Array, k: Array, v: Array, block: int = 7) -> Array:
    """One query, all keys visible. Online global softmax; not a GPU kernel."""
    m = -np.inf
    l = 0.0
    acc = np.zeros(v.shape[-1], dtype=np.float64)
    for i in range(0, len(k), block):
        s = k[i:i+block] @ q / math.sqrt(q.size)
        new_m = max(m, float(s.max()))
        alpha = math.exp(m-new_m)
        p = np.exp(s-new_m)
        acc = alpha * acc + p @ v[i:i+block]
        l = alpha * l + float(p.sum())
        m = new_m
    return acc / l


def pack_int4(x: Array, group: int = 16) -> tuple[Array, Array, tuple[int, ...]]:
    """Uniform signed INT4, NOT NVFP4. One FP32 scale per group, no zero point.
    Codes use [-7,7]; +8 offset enables two codes per uint8. Pad zeros at the tail.
    """
    if group < 2 or group % 2:
        raise ValueError("Group must be positive and even")
    f = np.asarray(x, dtype=np.float32).reshape(-1)
    if not np.isfinite(f).all():
        raise ValueError("Quantization requires finite input")
    n = ((len(f) + group - 1) // group) * group
    a = np.pad(f, (0, n-len(f))).reshape(-1, group)
    scales = np.max(np.abs(a), axis=1).astype(np.float32) / 7
    scales[scales == 0] = 1
    codes = (np.rint(a / scales[:, None]).clip(-7, 7).astype(np.int8) + 8).astype(np.uint8)
    codes = codes.reshape(-1)
    packed = codes[0::2] | (codes[1::2] << 4)
    return packed, scales, x.shape


def unpack_int4(packed: Array, scales: Array, shape: tuple[int, ...], group: int = 16) -> Array:
    codes = np.empty(packed.size * 2, dtype=np.int16)
    codes[0::2] = packed & 15
    codes[1::2] = packed >> 4
    a = ((codes - 8).reshape(-1, group) * scales[:, None]).reshape(-1)
    return a[:math.prod(shape)].reshape(shape)


class PagePool:
    """Small single-threaded page-ownership simulation (not model execution)."""
    def __init__(self, capacity: int = 4):
        self.capacity = capacity
        self.pages: dict[int, Array] = {}
        self.refs: dict[int, int] = {}

    def allocate(self, data: Array) -> int:
        for p in range(self.capacity):
            if self.refs.get(p, 0) == 0:
                self.pages[p] = data.copy()
                self.refs[p] = 1
                return p
        raise MemoryError("All pages are held by readers")

    def retain(self, p: int) -> None:
        if self.refs.get(p, 0) <= 0:
            raise RuntimeError("Cannot retain a free page")
        self.refs[p] += 1

    def release(self, p: int) -> None:
        if self.refs.get(p, 0) <= 0:
            raise RuntimeError("Double release or unknown page")
        self.refs[p] -= 1

    def private_for_write(self, p: int) -> int:
        if self.refs.get(p, 0) <= 0:
            raise RuntimeError("Cannot write a free page")
        if self.refs[p] == 1:
            return p
        # Allocate before dropping the old reference, so failure is safe.
        new_p = self.allocate(self.pages[p])
        self.release(p)
        return new_p


def max_error(a: Array, b: Array) -> float:
    return float(np.max(np.abs(a-b)))


def experiments() -> dict[str, Any]:
    rng = np.random.default_rng(20261009)
    prefix = rng.integers(0, 41, size=42).tolist()
    request_a = prefix + [4, 5, 6]
    request_b = prefix + [7, 8, 9, 10, 11]
    rows = []
    for dtype in ('float64', 'float32'):
        model = TinyDecoder(Config(dtype=dtype))
        cache = PrefixCache(model)
        cache.request(request_a)
        cached, _, usage = cache.request(request_b)
        full, _ = model.forward(request_b)
        _, pre = model.forward(request_b[:usage['hit_tokens']])
        split, _ = model.forward(request_b[usage['hit_tokens']:], pre)
        tolerance = 1e-10 if dtype == 'float64' else 2e-5
        assert max_error(cached, full[-1]) < tolerance
        assert max_error(cached, split[-1]) < tolerance
        rows.append({'dtype': dtype, **usage,
                     'cached_vs_full_max_abs': max_error(cached, full[-1]),
                     'split_vs_full_max_abs': max_error(split[-1], full[-1]),
                     'cached_vs_split_max_abs': max_error(cached, split[-1])})

    model = TinyDecoder()
    full, state = model.forward(request_b)
    _, old = model.forward([1, 2, 3, 4] + prefix)
    _, new = model.forward([5, 6, 7, 8] + prefix)
    first = max_error(old[0][0][:,4:], new[0][0][:,4:])
    higher = max_error(old[1][0][:,4:], new[1][0][:,4:])
    assert first < 1e-10 and higher > 1e-6
    # Same-length prefixes: isolate content dependency without a RoPE-offset change.
    _, lead = model.forward([5, 6, 7, 8])
    wrong = [(np.concatenate([lead[i][0], old[i][0][:,4:]], axis=1),
              np.concatenate([lead[i][1], old[i][1][:,4:]], axis=1))
             for i in range(model.cfg.layers)]
    wrong_logits, _ = model.forward([9], wrong)
    right_logits, _ = model.forward([9], new)
    wrong_error = max_error(wrong_logits, right_logits)
    assert wrong_error > 1e-6
    # Replay document under its new prefix: exact repair baseline, not CacheBlend.
    repaired_logits, repaired = model.forward(prefix, lead)
    repaired_next, _ = model.forward([9], repaired)
    assert max_error(repaired_next, right_logits) < 1e-10

    q = rng.normal(size=16)
    k = rng.normal(size=(37,16))
    v = rng.normal(size=(37,12))
    dense = softmax(k @ q / 4) @ v
    online = online_attention(q,k,v)
    online_error = max_error(dense, online)
    assert online_error < 1e-12
    data = rng.normal(size=(4,32,16)).astype(np.float32)
    packed, scales, shape = pack_int4(data)
    recovered = unpack_int4(packed, scales, shape)
    # Quantized model state errors are measured, never interpreted as task quality.
    quant_state = []
    for kk,vv in state:
        quant_state.append(tuple(unpack_int4(*pack_int4(a)).astype('float64') for a in (kk,vv)))
    ref_next,_ = model.forward([12],state)
    quant_next,_ = model.forward([12],quant_state)
    return {'lab_version':'1.0', 'run_date':'2026-10-09',
            'environment': {'python':platform.python_version(), 'numpy':np.__version__,
                            'platform':platform.platform(), 'device':'CPU'},
            'scope':'Random tiny model; numerical and state-management tests only. No API or pretrained-model benchmark.',
            'config':model.cfg.__dict__, 'model_identity':model.identity,
            'prefix_tests':rows,
            'context_dependency': {'first_layer_document_k_max_abs': first,
                                    'second_layer_document_k_max_abs': higher,
                                    'wrong_splice_next_logits_max_abs': wrong_error,
                                    'full_replay_repair_max_abs': max_error(repaired_next,right_logits)},
            'online_softmax_max_abs': online_error,
            'quantization': {'format':'uniform signed INT4; FP32 scale per 16 values',
                             'original_fp32_bytes':int(data.nbytes),
                             'hypothetical_fp16_bytes':int(data.size * 2),
                             'packed_data_bytes':int(packed.nbytes),
                             'scale_bytes':int(scales.nbytes),
                             'total_payload_bytes':int(packed.nbytes + scales.nbytes),
                             'reconstruction_max_abs':max_error(data,recovered),
                             'next_logits_max_abs':max_error(ref_next,quant_next)},
            'timing_note':'No latency benchmark: toy CPU timing is not evidence of GPU speedup.'}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--output', type=Path, default=Path('results.json'))
    args = parser.parse_args()
    result = experiments()
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2)+'\n',encoding='utf-8')
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == '__main__':
    main()

test_kv_lab.py · 与上方保存按钮的内容一致

"""Run with python -m unittest -v. Tests exercise deliberate failure controls."""
import unittest
import numpy as np
from kv_lab import (TinyDecoder, Config, PrefixCache, PagePool, softmax,
                    online_attention, pack_int4, unpack_int4, max_error)

class KVLabTests(unittest.TestCase):
    def setUp(self):
        self.model = TinyDecoder()
        self.tokens = (list(range(31)) * 2)[:53]

    def test_incremental_matches_full(self):
        full, _ = self.model.forward(self.tokens)
        _, pre = self.model.forward(self.tokens[:32])
        tail, _ = self.model.forward(self.tokens[32:], pre)
        self.assertLess(max_error(full[32:], tail), 1e-10)

    def test_prefix_hit_processes_only_suffix(self):
        cache = PrefixCache(self.model)
        cache.request(self.tokens)
        b = self.tokens[:42] + [39, 40, 38]
        before = self.model.processed_positions
        logits, _, usage = cache.request(b)
        self.assertEqual(usage['hit_tokens'], 32)
        self.assertEqual(self.model.processed_positions-before, 13)
        ref,_ = self.model.forward(b)
        self.assertLess(max_error(logits,ref[-1]),1e-10)

    def test_full_block_hit_leaves_one_block(self):
        cache = PrefixCache(self.model)
        cache.request(self.tokens[:48])
        _,_,u = cache.request(self.tokens[:48])
        self.assertEqual(u['hit_tokens'], 32)
        self.assertEqual(u['computed_tokens'], 16)

    def test_early_change_breaks_hash_chain(self):
        cache = PrefixCache(self.model)
        cache.request(self.tokens)
        changed = self.tokens.copy();changed[0]=40
        self.assertEqual(cache.lookup(changed)[0], 0)

    def test_identity_includes_weights_and_namespace(self):
        c1=PrefixCache(self.model,namespace='alpha')
        c2=PrefixCache(TinyDecoder(Config(seed=8)),namespace='alpha')
        c3=PrefixCache(self.model,namespace='beta')
        self.assertNotEqual(c1.base,c2.base)
        self.assertNotEqual(c1.base,c3.base)

    def test_past_not_mutated_and_no_private_alias(self):
        _, past = self.model.forward(self.tokens[:32])
        saved=past[0][0].copy()
        self.model.forward(self.tokens[32:],past)
        np.testing.assert_array_equal(saved,past[0][0])
        cache=PrefixCache(self.model);cache.request(self.tokens)
        _,hit=cache.lookup(self.tokens)
        hit[0][0][:]=0
        _,hit2=cache.lookup(self.tokens)
        self.assertGreater(np.max(np.abs(hit2[0][0])),0)

    def test_cache_shape_counts_kv_heads_not_q_heads(self):
        _,state=self.model.forward(self.tokens)
        c=self.model.cfg
        size=sum(k.nbytes+v.nbytes for k,v in state)
        self.assertEqual(size,2*c.layers*len(self.tokens)*c.kv_heads*c.head_dim*8)

    def test_online_global_softmax(self):
        rng=np.random.default_rng(2);q=rng.normal(size=8)
        k=rng.normal(size=(29,8));v=rng.normal(size=(29,5))
        ref=softmax(k@q/np.sqrt(8))@v
        for block in [1,3,7,29,50]:
            self.assertLess(max_error(ref,online_attention(q,k,v,block)),1e-12)

    def test_packed_int4_bytes_and_bound(self):
        x=np.linspace(-3,3,64,dtype=np.float32).reshape(4,16)
        packed,scales,shape=pack_int4(x)
        y=unpack_int4(packed,scales,shape)
        self.assertEqual(packed.nbytes,32)
        self.assertEqual(scales.nbytes,16)
        self.assertLessEqual(max_error(x,y),float(scales.max())/2+1e-6)
        z=np.zeros(17,dtype=np.float32)
        np.testing.assert_array_equal(unpack_int4(*pack_int4(z)),z)

    def test_reference_count_protects_second_reader(self):
        pool=PagePool(1);p=pool.allocate(np.array([1.0,2.0]))
        pool.retain(p);pool.release(p)
        with self.assertRaises(MemoryError):pool.allocate(np.array([9.0]))
        np.testing.assert_array_equal(pool.pages[p],[1.0,2.0])
        pool.release(p)
        self.assertEqual(pool.allocate(np.array([9.0])),p)

    def test_copy_on_write_separates_branches(self):
        pool=PagePool(2);p=pool.allocate(np.array([1.0,2.0,0.0]))
        pool.retain(p);branch=pool.private_for_write(p)
        pool.pages[branch][2]=9.0
        self.assertEqual(pool.pages[p][2],0.0)
        self.assertEqual(pool.refs[p],1);self.assertEqual(pool.refs[branch],1)

    def test_double_release_rejected(self):
        pool=PagePool();p=pool.allocate(np.zeros(3));pool.release(p)
        with self.assertRaises(RuntimeError):pool.release(p)

    def test_wrong_splice_is_detectable(self):
        doc=[8,9,10,11,12]
        _,old=self.model.forward([1,2]+doc)
        _,newlead=self.model.forward([3,4])
        wrong=[(np.concatenate([newlead[i][0],k[:,2:]],axis=1),
                np.concatenate([newlead[i][1],v[:,2:]],axis=1))
               for i,(k,v) in enumerate(old)]
        bad,_=self.model.forward([13],wrong)
        good,_=self.model.forward([3,4]+doc+[13])
        self.assertGreater(max_error(bad[-1],good[-1]),1e-6)

    def test_eviction_recovers_correctly(self):
        cache=PrefixCache(self.model,capacity_blocks=2)
        cache.request(self.tokens)
        self.assertLessEqual(len(cache.blocks),2)
        out,_,_=cache.request(self.tokens)
        full,_=self.model.forward(self.tokens)
        self.assertLess(max_error(out,full[-1]),1e-10)

    def test_invalid_inputs_rejected(self):
        for t in [[],[-1],[41],[1.2]]:
            with self.assertRaises(ValueError):self.model.forward(t)

if __name__=='__main__':unittest.main()

results.json · 与上方保存按钮的内容一致

{
  "lab_version": "1.0",
  "run_date": "2026-10-09",
  "environment": {
    "python": "3.13.5",
    "numpy": "2.3.5",
    "platform": "Linux-6.18.44-x86_64-with-glibc2.41",
    "device": "CPU"
  },
  "scope": "Random tiny model; numerical and state-management tests only. No API or pretrained-model benchmark.",
  "config": {
    "vocab": 41,
    "width": 32,
    "layers": 3,
    "q_heads": 4,
    "kv_heads": 2,
    "head_dim": 8,
    "seed": 7,
    "dtype": "float64"
  },
  "model_identity": "tiny-gqa-rope-v1:d71a873dd3c45a619daae621b2fb66f86897f5f6b16f068a738f212adb3e60f2",
  "prefix_tests": [
    {
      "dtype": "float64",
      "input_tokens": 47,
      "hit_tokens": 32,
      "computed_tokens": 15,
      "cached_vs_full_max_abs": 0.0,
      "split_vs_full_max_abs": 0.0,
      "cached_vs_split_max_abs": 0.0
    },
    {
      "dtype": "float32",
      "input_tokens": 47,
      "hit_tokens": 32,
      "computed_tokens": 15,
      "cached_vs_full_max_abs": 0.0,
      "split_vs_full_max_abs": 0.0,
      "cached_vs_split_max_abs": 0.0
    }
  ],
  "context_dependency": {
    "first_layer_document_k_max_abs": 0.0,
    "second_layer_document_k_max_abs": 2.112007194820824,
    "wrong_splice_next_logits_max_abs": 1.3119541099416636,
    "full_replay_repair_max_abs": 0.0
  },
  "online_softmax_max_abs": 5.551115123125783e-17,
  "quantization": {
    "format": "uniform signed INT4; FP32 scale per 16 values",
    "original_fp32_bytes": 8192,
    "hypothetical_fp16_bytes": 4096,
    "packed_data_bytes": 1024,
    "scale_bytes": 512,
    "total_payload_bytes": 1536,
    "reconstruction_max_abs": 0.21534329652786255,
    "next_logits_max_abs": 0.09491987950877667
  },
  "timing_note": "No latency benchmark: toy CPU timing is not evidence of GPU speedup."
}

tests.txt · 与上方保存按钮的内容一致

test_cache_shape_counts_kv_heads_not_q_heads (test_kv_lab.KVLabTests.test_cache_shape_counts_kv_heads_not_q_heads) ... ok
test_copy_on_write_separates_branches (test_kv_lab.KVLabTests.test_copy_on_write_separates_branches) ... ok
test_double_release_rejected (test_kv_lab.KVLabTests.test_double_release_rejected) ... ok
test_early_change_breaks_hash_chain (test_kv_lab.KVLabTests.test_early_change_breaks_hash_chain) ... ok
test_eviction_recovers_correctly (test_kv_lab.KVLabTests.test_eviction_recovers_correctly) ... ok
test_full_block_hit_leaves_one_block (test_kv_lab.KVLabTests.test_full_block_hit_leaves_one_block) ... ok
test_identity_includes_weights_and_namespace (test_kv_lab.KVLabTests.test_identity_includes_weights_and_namespace) ... ok
test_incremental_matches_full (test_kv_lab.KVLabTests.test_incremental_matches_full) ... ok
test_invalid_inputs_rejected (test_kv_lab.KVLabTests.test_invalid_inputs_rejected) ... ok
test_online_global_softmax (test_kv_lab.KVLabTests.test_online_global_softmax) ... ok
test_packed_int4_bytes_and_bound (test_kv_lab.KVLabTests.test_packed_int4_bytes_and_bound) ... ok
test_past_not_mutated_and_no_private_alias (test_kv_lab.KVLabTests.test_past_not_mutated_and_no_private_alias) ... ok
test_prefix_hit_processes_only_suffix (test_kv_lab.KVLabTests.test_prefix_hit_processes_only_suffix) ... ok
test_reference_count_protects_second_reader (test_kv_lab.KVLabTests.test_reference_count_protects_second_reader) ... ok
test_wrong_splice_is_detectable (test_kv_lab.KVLabTests.test_wrong_splice_is_detectable) ... ok

----------------------------------------------------------------------
Ran 15 tests in 0.086s

OK

README.md · 与上方保存按钮的内容一致

# KV 缓存最小实验：先把状态复用做出来

这是补充教材的可运行伴随实验，版本 1.0，整理于 2026-10-09。
**它不是原附录 A 中 Qwen3-0.6B 实验的复现，也不是训练好的语言模型。**

## 跑什么

`kv_lab.py` 从固定随机种子构建三层因果 Transformer：4 个 Q 头、2 个 KV 头、每头 8 维、RoPE、RMSNorm、前馈层和词表输出。词表为 41 个整数 ID，无分词器、外部数据集、模型下载或付费 API。

它包含完整前缀缓存（16 token 一块、链式 SHA-256、权重和命名空间身份、私有复制）、在线 softmax、实际打包的 INT4，以及独立的共享页／写时复制模拟器。

## 运行

需要 Python 3.10 及以上。本次验证环境是 Python 3.13.5、NumPy 2.3.5、CPU。
已有 NumPy 时可直接运行；需要安装时由你执行：

```bash
python -m pip install numpy==2.3.5
python -m unittest -v
python kv_lab.py --output my_results.json
```

在 macOS / Linux 上可用 `OPENBLAS_NUM_THREADS=1 python kv_lab.py --output my_results.json` 固定 BLAS 线程数。本次记录使用该设置。测试不要求每一种硬件都与所附 JSON 逐位相同，而是检查 FP64 / FP32 对应的数值容差。NumPy 2.3.5 的 Python 支持范围以安装环境为准；Python 3.10 用户可使用支持该版本的 NumPy，实际结果应记录版本。

## 先预测，再看结果

1. 两个输入共享前 42 个 token，缓存以 16 token 为单位，只保存整块。第二个输入共 47 token。命中 32，重算 15；不是共享的 42 都能复用。
2. “不查表，只把输入分两段计算”是必要对照：它能分离分段执行带来的数值差异与缓存管理错误。
3. 把文档的旧 K/V 接到长度相同、内容不同的前文后，位置不变仍然可能错。重放文档是精确修复对照；本实验**没有实现 CacheBlend**。
4. `PagePool` 中先让两个使用者共享一页，再结束其中一个；另一人没结束之前，不能覆盖。写共享尾页之前先复制。
5. INT4 两个码打包为 1 字节。每 16 个数一个 FP32 缩放系数，因此 2,048 个数的载荷是 1,024 + 512 = 1,536 字节。**这不是 NVFP4 格式**，也未计 Python 对象、对齐及工作区。

## 文件

- `kv_lab.py`：完整模型与缓存实验；会写入你指定的结果 JSON。
- `test_kv_lab.py`：15 项自动测试。
- `results.json`：本次实际运行的测量值与环境，不是示例数字。
- `tests.txt`：本次测试原始输出。
- `requirements.txt`：本次 NumPy 版本。
- `SOURCE_NOTES.md`：机制依据与实验范围。

## 有意没有实现的东西

没有真实 GPU 内核、异步网络传输、跨进程并发、语义正确率评测、生产安全隔离、缓存持久化或训练。`PrefixCache` 的复制与回收完全串行；`PagePool` 是另一个独立生命周期模拟，不是 `TinyDecoder` 的分页注意力后端。不能据此声称已经复现 vLLM、FlashAttention、NVFP4、CacheBlend 或云端服务。

缓存权重身份包含实际数组的 SHA-256，而不是只用模型名字；实验中的命名空间是本地隔离示例，不是线上鉴权或秘密 salt 的替代。JSON 中的 `run_date` 是本版记录的整理日；重跑时应另外保存自己的日期／环境。