Модуль 10. KV-кэши и эффективный инференс

from __future__ import annotations

from dataclasses import dataclass
import hashlib
import json
import math
from typing import Hashable, Sequence

REFERENCE_JSON = '{"capacity_upper_bound":{"device_binary_gib":80,"max_requests":{"1000000":0,"131072":3,"32768":14,"8192":57},"runtime_reserve_binary_gib":8},"compression":{"base_bits_per_coordinate":16,"base_kib_per_token":128.0,"ceiling_128k_tokens_per_second":174.286494109339,"compressed_bits_per_coordinate":3,"compressed_kib_per_token":24.0},"context_memory_70b_gqa":{"1000000":{"binary_gib":305.17578125,"bytes":327680000000,"decimal_gb":327.68},"10000000":{"binary_gib":3051.7578125,"bytes":3276800000000,"decimal_gb":3276.8},"131072":{"binary_gib":40.0,"bytes":42949672960,"decimal_gb":42.94967296},"32768":{"binary_gib":10.0,"bytes":10737418240,"decimal_gb":10.73741824},"8192":{"binary_gib":2.5,"bytes":2684354560,"decimal_gb":2.68435456}},"eviction":{"attention_mass":[0.625,0.897,0.776,0.225,0.3,0.874,0.005,0.821,0.797,0.468,0.303,0.278,0.255,0.445,0.505,0.553],"heavy_hitters":[0,1,2,5,7,8,14,15],"lost_heavy_hitters":[5,7,8],"streaming":[0,1,2,3,12,13,14,15]},"hardware":{"bandwidth_decimal_tb_per_second":3.35,"device_binary_gib":80.0,"peak_decimal_tflops":990.0,"runtime_reserve_binary_gib":8.0,"weight_binary_gib":14.901161193848,"weight_bytes":16000000000,"weight_decimal_gb":16.0},"kv_ladder":{"70B MHA-equivalent (H=64)":{"bytes_per_token":2621440,"kib_per_token":2560.0},"DeepSeek-V3 MLA (latent+RoPE)":{"bytes_per_token":70272,"kib_per_token":68.625},"LLaMA-3.3-70B GQA":{"bytes_per_token":327680,"kib_per_token":320.0},"Llama-3-8B GQA":{"bytes_per_token":131072,"kib_per_token":128.0}},"paged":{"exact_2048_waste_fraction":0.0,"expected_paged_waste_fraction":0.003648747263,"naive_waste_fraction":0.9375,"preallocation_capacity_factor":16.0,"trace":{"after_a":{"allocated_blocks":2,"free_blocks":6,"internal_waste_tokens":2,"logical_tokens":6,"physical_capacity_tokens":8,"physical_payload_tokens":6,"sequences":{"A":{"blocks":[0,1],"length":6}},"sharing_savings_tokens":0},"after_a_release":{"allocated_blocks":3,"free_blocks":5,"internal_waste_tokens":3,"logical_tokens":9,"physical_capacity_tokens":12,"physical_payload_tokens":9,"sequences":{"B":{"blocks":[0,2,3],"length":9}},"sharing_savings_tokens":0},"after_b_append":{"allocated_blocks":4,"free_blocks":4,"internal_waste_tokens":5,"logical_tokens":15,"physical_capacity_tokens":16,"physical_payload_tokens":11,"sequences":{"A":{"blocks":[0,1],"length":6},"B":{"blocks":[0,2,3],"length":9}},"sharing_savings_tokens":4},"after_share":{"allocated_blocks":2,"free_blocks":6,"internal_waste_tokens":2,"logical_tokens":12,"physical_capacity_tokens":8,"physical_payload_tokens":6,"sequences":{"A":{"blocks":[0,1],"length":6},"B":{"blocks":[0,1],"length":6}},"sharing_savings_tokens":6}}},"prefix_and_interference":{"arithmetic_speedup":16.0,"cached_causal_suffix_ms":10.426705047014,"cached_suffix_prefill_ms":8.274747474747,"causal_arithmetic_ratio":14.402250086446,"chunked_itl_ms":13.050866877733,"decode_weight_floor_ms":4.776119402985,"full_causal_prefill_ms":150.168013664711,"full_prefill_ms":132.39595959596,"illustrative_input_cost":{"amortized_cost_factor":3.786982248521,"baseline":81920,"cached":21632.0,"warm_request_cost_ratio":0.15625},"monolithic_spike_ms":137.172078998945},"roofline":{"batch_32_8k":{"aggregate_tokens_per_second":2128.684609452179,"ideal_without_kv_tokens_per_second":6700.0,"kv_tax_factor":3.147483648},"batch_knees":{"131072":0.931322574615,"32768":3.725290298462,"8192":14.901161193848},"ceilings":{"0":{"aggregate_tokens_per_second":209.375,"kv_bytes":0},"1000000":{"aggregate_tokens_per_second":22.777959094865,"kv_bytes":131072000000},"131072":{"aggregate_tokens_per_second":100.96483447305,"kv_bytes":17179869184},"8192":{"aggregate_tokens_per_second":196.207722626508,"kv_bytes":1073741824}},"intensity_flops_per_byte_128k":0.482220104946,"intensity_including_attention_128k":2.553339685162,"ridge_flops_per_byte":295.522388059701,"ridge_to_full_intensity":115.739550744859,"ridge_to_intensity":612.837136047761},"schema_version":3,"unit_policy":{"binary":"KiB/MiB/GiB/TiB use powers of two","decimal":"GB/TB and bandwidth use powers of ten"}}'
REFERENCE_SHA256 = '1a283dbcb4061ab38b42bdf2d48f2ad1b1c844ffe88faab9137703985ee9ca5b'
actual_sha256 = hashlib.sha256(REFERENCE_JSON.encode("utf-8")).hexdigest()
assert actual_sha256 == REFERENCE_SHA256
REF = json.loads(REFERENCE_JSON)
assert REF["schema_version"] == 3

KIB, MIB, GIB, TIB = 2**10, 2**20, 2**30, 2**40
DECIMAL_GB, DECIMAL_TB = 10**9, 10**12
print('Встроенный контракт М10 проверен', actual_sha256[:12])

def positive_int(name: str, value: int) -> int:
    """Validate a strictly positive integer without accepting bools."""

    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return value


def non_negative_int(name: str, value: int) -> int:
    """Validate a non-negative integer without accepting bools."""

    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        raise ValueError(f"{name} must be a non-negative integer")
    return value


def positive_finite(name: str, value: float) -> float:
    """Validate a positive finite numeric scalar."""

    if isinstance(value, bool) or not isinstance(value, (int, float)):
        raise TypeError(f"{name} must be numeric")
    number = float(value)
    if not math.isfinite(number) or number <= 0.0:
        raise ValueError(f"{name} must be positive and finite")
    return number


@dataclass(frozen=True)
class AttentionKVGeometry:
    """KV geometry for layers that materialize ordinary K and V tensors."""

    name: str
    attention_layers: int
    kv_heads: int
    head_dim: int
    bytes_per_element: int = 2

    def __post_init__(self) -> None:
        """Validate the declared ordinary-attention geometry."""

        if not isinstance(self.name, str) or not self.name.strip():
            raise ValueError("name must be a non-empty string")
        positive_int("attention_layers", self.attention_layers)
        positive_int("kv_heads", self.kv_heads)
        positive_int("head_dim", self.head_dim)
        positive_int("bytes_per_element", self.bytes_per_element)

    @property
    def bytes_per_token(self) -> int:
        """Return K+V bytes for one token across all attention layers."""

        return (2 * self.attention_layers * self.kv_heads
                * self.head_dim * self.bytes_per_element)


@dataclass(frozen=True)
class MLAKVGeometry:
    """KV geometry for a stored latent and separate positional component."""

    name: str
    attention_layers: int
    latent_dim: int
    rope_dim: int
    bytes_per_element: int = 2

    def __post_init__(self) -> None:
        """Validate the declared latent-attention geometry."""

        if not isinstance(self.name, str) or not self.name.strip():
            raise ValueError("name must be a non-empty string")
        positive_int("attention_layers", self.attention_layers)
        positive_int("latent_dim", self.latent_dim)
        non_negative_int("rope_dim", self.rope_dim)
        positive_int("bytes_per_element", self.bytes_per_element)

    @property
    def bytes_per_token(self) -> int:
        """Return latent-plus-position bytes for one token across all layers."""

        return (self.attention_layers * (self.latent_dim + self.rope_dim)
                * self.bytes_per_element)


def kv_payload_bytes(geometry, context_tokens: int, batch_size: int = 1,
                     kv_scale: float = 1.0) -> int:
    """Return an exact integer KV payload for equal-length requests."""

    tokens = non_negative_int("context_tokens", context_tokens)
    batch = positive_int("batch_size", batch_size)
    scale = positive_finite("kv_scale", kv_scale)
    raw = geometry.bytes_per_token * tokens * batch * scale
    rounded = round(raw)
    if not math.isclose(raw, rounded, rel_tol=0.0, abs_tol=1e-9):
        raise ValueError("scaled KV payload is not an integer byte count")
    return int(rounded)


GEOMETRIES = {
    "70B MHA-equivalent (H=64)": AttentionKVGeometry(
        "70B MHA-equivalent (H=64)", 80, 64, 128, 2),
    "LLaMA-3.3-70B GQA": AttentionKVGeometry(
        "LLaMA-3.3-70B GQA", 80, 8, 128, 2),
    "Llama-3-8B GQA": AttentionKVGeometry(
        "Llama-3-8B GQA", 32, 8, 128, 2),
    "DeepSeek-V3 MLA (latent+RoPE)": MLAKVGeometry(
        "DeepSeek-V3 MLA (latent+RoPE)", 61, 512, 64, 2),
}
print('геометрия KV и функции преобразования единиц определены')

for name, geometry in GEOMETRIES.items():
    expected = REF["kv_ladder"][name]
    assert geometry.bytes_per_token == expected["bytes_per_token"]
    assert math.isclose(geometry.bytes_per_token / KIB,
                        expected["kib_per_token"], abs_tol=1e-12)
    print(f"{name:36s} {geometry.bytes_per_token / KIB:8.3f} KiB/token")

g70 = GEOMETRIES["LLaMA-3.3-70B GQA"]
for token_text, expected in REF["context_memory_70b_gqa"].items():
    tokens = int(token_text)
    payload = kv_payload_bytes(g70, tokens)
    assert payload == expected["bytes"]
    assert math.isclose(payload / DECIMAL_GB, expected["decimal_gb"], abs_tol=1e-12)
    assert math.isclose(payload / GIB, expected["binary_gib"], abs_tol=1e-12)
    print(f"N={tokens:>10,}: {payload / DECIMAL_GB:9.3f} GB"
          f" = {payload / GIB:9.3f} GiB")
print('Лестница и объёмы контекста воспроизведены')

"""Small, dependency-free FLOP and billing models; no hardware timing claims."""
import math


def attention_flops(layers, query_heads, head_dim, queries, cached_prefix=0):
    """QK and AV for causal suffix queries, two FLOPs per multiply-add.

    Query j sees cached_prefix + j + 1 keys. Softmax, norms, projections,
    cache writes, and kernel padding are outside this count.
    """
    for name, value in (("layers", layers), ("query_heads", query_heads),
                        ("head_dim", head_dim), ("queries", queries),
                        ("cached_prefix", cached_prefix)):
        if isinstance(value, bool) or not isinstance(value, int) or value < 0:
            raise ValueError(f"{name} must be a non-negative integer")
    if min(layers, query_heads, head_dim) == 0:
        raise ValueError("attention geometry must be positive")
    pairs = queries * cached_prefix + queries * (queries + 1) // 2
    return 4 * layers * query_heads * head_dim * pairs


def decode_flops(parameters, layers, query_heads, head_dim, keys, batch=1):
    """Projection estimate plus QK/AV for one query attending to `keys`."""
    for value in (parameters, batch, keys):
        if isinstance(value, bool) or not isinstance(value, int) or value < 0:
            raise ValueError("counts must be non-negative integers")
    if parameters == 0 or batch == 0:
        raise ValueError("parameters and batch must be positive")
    # Zero keys is a synthetic weights-only baseline, not a real decode query.
    attn = attention_flops(layers, query_heads, head_dim, 1, max(keys - 1, 0))
    return batch * (2 * parameters + (attn if keys else 0))


def prefill_flops(parameters, layers, query_heads, head_dim, suffix, prefix=0):
    """Only suffix projections are new; suffix attention still reads the prefix."""
    if isinstance(parameters, bool) or not isinstance(parameters, int) or parameters <= 0:
        raise ValueError("parameters must be a positive integer")
    return (attention_flops(layers, query_heads, head_dim, suffix, prefix)
            + 2 * parameters * suffix)


def cached_input_cost(prefix, suffix, requests, read_rate=0.1, write_rate=1.25):
    """Input-token equivalents: one prefix write, then complete prefix hits.

    Each suffix is uncached. No output cost, expiry, or extra cache writes.
    Multipliers are illustrative inputs, not an API tariff lookup.
    """
    for value in (prefix, suffix, requests):
        if isinstance(value, bool) or not isinstance(value, int) or value < 0:
            raise ValueError("counts must be non-negative integers")
    if requests == 0 or prefix + suffix == 0:
        raise ValueError("need a request with non-empty input")
    for rate in (read_rate, write_rate):
        if isinstance(rate, bool) or not isinstance(rate, (int, float)) or not math.isfinite(rate) or rate <= 0:
            raise ValueError("rates must be positive and finite")
    baseline = requests * (prefix + suffix)
    cached = prefix * (write_rate + (requests - 1) * read_rate) + requests * suffix
    warm = prefix * read_rate + suffix
    return {"baseline": baseline, "cached": cached,
            "amortized_cost_factor": baseline / cached,
            "warm_request_cost_ratio": warm / (prefix + suffix)}

@dataclass(frozen=True)
class ServingHardware:
    """Inputs for arithmetic serving ceilings and capacity upper bounds."""

    weight_bytes: int
    bandwidth_bytes_per_second: float
    peak_flops_per_second: float
    parameter_count: int
    device_capacity_bytes: int
    runtime_reserve_bytes: int = 0

    def __post_init__(self) -> None:
        """Validate arithmetic hardware inputs and the reserve budget."""

        positive_int("weight_bytes", self.weight_bytes)
        positive_finite("bandwidth_bytes_per_second", self.bandwidth_bytes_per_second)
        positive_finite("peak_flops_per_second", self.peak_flops_per_second)
        positive_int("parameter_count", self.parameter_count)
        positive_int("device_capacity_bytes", self.device_capacity_bytes)
        non_negative_int("runtime_reserve_bytes", self.runtime_reserve_bytes)
        if self.weight_bytes + self.runtime_reserve_bytes > self.device_capacity_bytes:
            raise ValueError("weights plus reserve exceed device capacity")


@dataclass(frozen=True)
class DecodeCeiling:
    """Traffic and token-rate fields for one declared decode scenario."""

    batch_size: int
    context_tokens: int
    kv_bytes: int
    total_bytes_per_step: int
    aggregate_tokens_per_second: float
    per_request_tokens_per_second: float


def decode_ceiling(hardware: ServingHardware, geometry, context_tokens: int,
                   batch_size: int = 1, kv_scale: float = 1.0) -> DecodeCeiling:
    """Compute a weight-plus-KV bandwidth ceiling for one decode step."""

    batch = positive_int("batch_size", batch_size)
    tokens = non_negative_int("context_tokens", context_tokens)
    kv_bytes = kv_payload_bytes(geometry, tokens, batch, kv_scale)
    total = hardware.weight_bytes + kv_bytes
    aggregate = batch * hardware.bandwidth_bytes_per_second / total
    return DecodeCeiling(batch, tokens, kv_bytes, total, aggregate,
                         aggregate / batch)


def arithmetic_intensity(hardware: ServingHardware, geometry,
                         context_tokens: int, batch_size: int = 1,
                         kv_scale: float = 1.0) -> float:
    """Projection-only intensity: context attention is added in the next cell."""

    scenario = decode_ceiling(hardware, geometry, context_tokens,
                              batch_size, kv_scale)
    return (2 * hardware.parameter_count * batch_size
            / scenario.total_bytes_per_step)


def batch_knee(hardware: ServingHardware, geometry,
               context_tokens: int, kv_scale: float = 1.0) -> float:
    """Return B* where one request's KV traffic equals weight traffic."""

    tokens = positive_int("context_tokens", context_tokens)
    return hardware.weight_bytes / kv_payload_bytes(
        geometry, tokens, 1, kv_scale)


HARDWARE = ServingHardware(
    weight_bytes=16_000_000_000,
    bandwidth_bytes_per_second=3.35e12,
    peak_flops_per_second=990e12,
    parameter_count=8_000_000_000,
    device_capacity_bytes=80 * GIB,
    runtime_reserve_bytes=8 * GIB,
)
g8 = GEOMETRIES["Llama-3-8B GQA"]
for token_text, expected in REF["roofline"]["ceilings"].items():
    scenario = decode_ceiling(HARDWARE, g8, int(token_text))
    assert scenario.kv_bytes == expected["kv_bytes"]
    assert math.isclose(scenario.aggregate_tokens_per_second,
                        expected["aggregate_tokens_per_second"], abs_tol=1e-9)
    print(f"N={int(token_text):>9,}:"
          f" {scenario.aggregate_tokens_per_second:8.3f} tok/s")

intensity = arithmetic_intensity(HARDWARE, g8, 131_072)
ridge = HARDWARE.peak_flops_per_second / HARDWARE.bandwidth_bytes_per_second
assert math.isclose(intensity,
                    REF["roofline"]["intensity_flops_per_byte_128k"], abs_tol=1e-12)
assert math.isclose(ridge, REF["roofline"]["ridge_flops_per_byte"], abs_tol=1e-12)
for token_text, expected in REF["roofline"]["batch_knees"].items():
    assert math.isclose(batch_knee(HARDWARE, g8, int(token_text)),
                        expected, abs_tol=1e-12)
b32 = decode_ceiling(HARDWARE, g8, 8_192, batch_size=32)
assert math.isclose(b32.aggregate_tokens_per_second,
                    REF["roofline"]["batch_32_8k"]["aggregate_tokens_per_second"],
                    abs_tol=1e-9)
print(f"128K intensity={intensity:.6f} FLOP/B; ridge={ridge:.3f} FLOP/B")
full_intensity = decode_flops(8_000_000_000, 32, 32, 128, 131072) / decode_ceiling(HARDWARE, g8, 131072).total_bytes_per_step
assert math.isclose(full_intensity, REF["roofline"]["intensity_including_attention_128k"], abs_tol=1e-12)
print(f"Including QK/AV: {full_intensity:.6f} FLOP/B")
print(f"Batch 32: aggregate={b32.aggregate_tokens_per_second:.3f} tok/s; per request={b32.per_request_tokens_per_second:.3f} tok/s")
print('Модель roofline, интенсивность и точки насыщения пакета проверены')

def max_resident_requests(hardware: ServingHardware, geometry,
                          context_tokens: int, kv_scale: float = 1.0) -> int:
    """Return a KV-only capacity upper bound after weights and reserve."""

    tokens = positive_int("context_tokens", context_tokens)
    free = (hardware.device_capacity_bytes - hardware.weight_bytes
            - hardware.runtime_reserve_bytes)
    return free // kv_payload_bytes(geometry, tokens, 1, kv_scale)


for token_text, expected in REF["capacity_upper_bound"]["max_requests"].items():
    actual = max_resident_requests(HARDWARE, g8, int(token_text))
    assert actual == expected
    print(f"N={int(token_text):>9,}: at most {actual} resident requests"
          " (KV-only upper bound)")
assert math.isclose(HARDWARE.weight_bytes / GIB,
                    REF["hardware"]["weight_binary_gib"], abs_tol=1e-12)
print('Верхние границы вместимости только по KV проверены')

_positive_int = positive_int
_non_negative_int = non_negative_int
@dataclass
class _PhysicalBlock:
    """Internal physical block state for the teaching page-pool simulator."""

    filled_slots: int = 0
    ref_count: int = 0


@dataclass
class _SequenceState:
    """Internal logical sequence state."""

    length: int
    blocks: list[int]


class PagedKVPool:
    """A validated page-table simulator with prefix sharing and copy-on-write.

    It models block ownership and logical lengths; it does not store real tensors.  A
    shared partial final block is copied before either sequence appends into it.  All
    allocations are preflighted so an out-of-memory error leaves the pool unchanged.
    """

    def __init__(self, total_blocks: int, block_size: int = 16) -> None:
        self.total_blocks = _positive_int("total_blocks", total_blocks)
        self.block_size = _positive_int("block_size", block_size)
        self._blocks = [_PhysicalBlock() for _ in range(self.total_blocks)]
        self._free = list(range(self.total_blocks - 1, -1, -1))
        self._sequences: dict[Hashable, _SequenceState] = {}

    def _require_sequence(self, sequence_id: Hashable) -> _SequenceState:
        try:
            return self._sequences[sequence_id]
        except KeyError as exc:
            raise KeyError(f"unknown sequence: {sequence_id!r}") from exc

    def create(self, sequence_id: Hashable) -> None:
        """Create an empty sequence with a unique hashable identifier."""

        try:
            hash(sequence_id)
        except TypeError as exc:
            raise TypeError("sequence_id must be hashable") from exc
        if sequence_id in self._sequences:
            raise ValueError(f"sequence already exists: {sequence_id!r}")
        self._sequences[sequence_id] = _SequenceState(length=0, blocks=[])

    def _blocks_needed_for_append(self, state: _SequenceState, token_count: int) -> int:
        if token_count == 0:
            return 0
        needed = 0
        remainder = state.length % self.block_size
        available_in_last = 0
        if state.blocks and remainder:
            last = self._blocks[state.blocks[-1]]
            if last.ref_count > 1:
                needed += 1  # copy-on-write replacement for the shared partial block
            available_in_last = self.block_size - remainder
        remaining = max(0, token_count - available_in_last)
        needed += math.ceil(remaining / self.block_size)
        return needed

    def _allocate_block(self, filled_slots: int = 0) -> int:
        block_id = self._free.pop()
        block = self._blocks[block_id]
        block.filled_slots = filled_slots
        block.ref_count = 1
        return block_id

    def append(self, sequence_id: Hashable, token_count: int) -> None:
        """Append logical tokens, allocating arbitrary physical blocks as needed."""

        count = _non_negative_int("token_count", token_count)
        state = self._require_sequence(sequence_id)
        required = self._blocks_needed_for_append(state, count)
        if required > len(self._free):
            raise MemoryError(
                f"need {required} free blocks, but only {len(self._free)} remain"
            )
        if count == 0:
            return

        remainder = state.length % self.block_size
        if state.blocks and remainder:
            last_id = state.blocks[-1]
            last = self._blocks[last_id]
            if last.ref_count > 1:
                last.ref_count -= 1
                replacement = self._allocate_block(filled_slots=remainder)
                state.blocks[-1] = replacement
                last_id = replacement
                last = self._blocks[last_id]
            take = min(count, self.block_size - remainder)
            last.filled_slots = remainder + take
            state.length += take
            count -= take

        while count:
            take = min(count, self.block_size)
            block_id = self._allocate_block(filled_slots=take)
            state.blocks.append(block_id)
            state.length += take
            count -= take
        self._refresh_fills()
        self.validate()

    def share_prefix(
        self,
        source_id: Hashable,
        target_id: Hashable,
        prefix_tokens: int,
    ) -> None:
        """Create ``target_id`` by sharing source blocks that cover a logical prefix."""

        source = self._require_sequence(source_id)
        prefix = _non_negative_int("prefix_tokens", prefix_tokens)
        if prefix > source.length:
            raise ValueError("prefix_tokens exceeds source length")
        if target_id in self._sequences:
            raise ValueError(f"target sequence already exists: {target_id!r}")
        try:
            hash(target_id)
        except TypeError as exc:
            raise TypeError("target_id must be hashable") from exc
        block_count = math.ceil(prefix / self.block_size) if prefix else 0
        table = source.blocks[:block_count]
        for block_id in table:
            self._blocks[block_id].ref_count += 1
        self._sequences[target_id] = _SequenceState(prefix, list(table))
        self.validate()

    def release(self, sequence_id: Hashable) -> None:
        """Release a sequence and return unreferenced physical blocks to the free list."""

        state = self._require_sequence(sequence_id)
        for block_id in state.blocks:
            block = self._blocks[block_id]
            block.ref_count -= 1
            if block.ref_count == 0:
                block.filled_slots = 0
                self._free.append(block_id)
        del self._sequences[sequence_id]
        self._refresh_fills()
        self.validate()

    def _refresh_fills(self) -> None:
        """Count live slots, including unequal prefixes sharing a physical block."""
        fills = [0] * self.total_blocks
        for live in self._sequences.values():
            for j, block_id in enumerate(live.blocks):
                fills[block_id] = max(fills[block_id], min(self.block_size, live.length - j * self.block_size))
        for block_id, block in enumerate(self._blocks):
            block.filled_slots = fills[block_id]

    def sequence_length(self, sequence_id: Hashable) -> int:
        """Return logical length in tokens."""

        return self._require_sequence(sequence_id).length

    def page_table(self, sequence_id: Hashable) -> tuple[int, ...]:
        """Return an immutable page table for inspection."""

        return tuple(self._require_sequence(sequence_id).blocks)

    @property
    def allocated_blocks(self) -> int:
        """Return the number of physical blocks with at least one reference."""

        return self.total_blocks - len(self._free)

    @property
    def free_blocks(self) -> int:
        """Return the number of currently free physical blocks."""

        return len(self._free)

    @property
    def physical_capacity_tokens(self) -> int:
        """Return token slots reserved by allocated physical blocks."""

        return self.allocated_blocks * self.block_size

    @property
    def physical_payload_tokens(self) -> int:
        """Return filled token slots stored once per physical block."""

        return sum(block.filled_slots for block in self._blocks if block.ref_count)

    @property
    def logical_tokens(self) -> int:
        """Return total logical tokens across all sequences, counting shared prefixes."""

        return sum(state.length for state in self._sequences.values())

    @property
    def internal_waste_tokens(self) -> int:
        """Return unused slots inside allocated physical blocks."""

        return self.physical_capacity_tokens - self.physical_payload_tokens

    @property
    def sharing_savings_tokens(self) -> int:
        """Return logical tokens not duplicated physically because of prefix sharing."""

        return self.logical_tokens - self.physical_payload_tokens

    def snapshot(self) -> dict[str, object]:
        """Return a deterministic JSON-friendly state summary."""

        sequence_items = sorted(self._sequences.items(), key=lambda item: repr(item[0]))
        return {
            "sequences": {
                str(sequence_id): {
                    "length": state.length,
                    "blocks": list(state.blocks),
                }
                for sequence_id, state in sequence_items
            },
            "allocated_blocks": self.allocated_blocks,
            "free_blocks": self.free_blocks,
            "physical_capacity_tokens": self.physical_capacity_tokens,
            "physical_payload_tokens": self.physical_payload_tokens,
            "logical_tokens": self.logical_tokens,
            "internal_waste_tokens": self.internal_waste_tokens,
            "sharing_savings_tokens": self.sharing_savings_tokens,
        }

    def validate(self) -> None:
        """Check page-table, reference-count, fill, and capacity invariants."""

        if len(set(self._free)) != len(self._free):
            raise AssertionError("free list contains duplicates")
        if any(block_id < 0 or block_id >= self.total_blocks for block_id in self._free):
            raise AssertionError("free list contains an invalid block id")
        expected_refs = [0] * self.total_blocks
        for state in self._sequences.values():
            if state.length < 0:
                raise AssertionError("sequence length is negative")
            expected_count = math.ceil(state.length / self.block_size) if state.length else 0
            if len(state.blocks) != expected_count:
                raise AssertionError("page-table length does not match sequence length")
            if len(set(state.blocks)) != len(state.blocks):
                raise AssertionError("one sequence references a physical block twice")
            for block_id in state.blocks:
                if block_id < 0 or block_id >= self.total_blocks:
                    raise AssertionError("page table contains an invalid block id")
                expected_refs[block_id] += 1
        free_set = set(self._free)
        for block_id, block in enumerate(self._blocks):
            if block.ref_count != expected_refs[block_id]:
                raise AssertionError("physical ref_count does not match page tables")
            if block.ref_count == 0:
                if block_id not in free_set or block.filled_slots != 0:
                    raise AssertionError("unreferenced block is not clean and free")
            else:
                if block_id in free_set:
                    raise AssertionError("referenced block appears in free list")
                if not 1 <= block.filled_slots <= self.block_size:
                    raise AssertionError("referenced block has invalid fill")
        if self.logical_tokens < self.physical_payload_tokens:
            raise AssertionError("physical payload exceeds logical payload")

def naive_preallocation_waste(max_tokens: int, used_tokens: int) -> float:
    """Return unused fraction of a fixed maximum-length allocation."""

    maximum = positive_int("max_tokens", max_tokens)
    used = non_negative_int("used_tokens", used_tokens)
    if used > maximum:
        raise ValueError("used_tokens cannot exceed max_tokens")
    return (maximum - used) / maximum


def expected_paged_waste(avg_tokens: int, block_size: int) -> float:
    """Return expected final-block waste under a uniform remainder."""

    average = positive_int("avg_tokens", avg_tokens)
    block = positive_int("block_size", block_size)
    tail = (block - 1) / 2
    return tail / (average + tail)


pool = PagedKVPool(total_blocks=8, block_size=4)
pool.create("A")
pool.append("A", 6)
assert pool.snapshot() == REF["paged"]["trace"]["after_a"]
pool.share_prefix("A", "B", 6)
assert pool.snapshot() == REF["paged"]["trace"]["after_share"]
pool.append("B", 3)
assert pool.snapshot() == REF["paged"]["trace"]["after_b_append"]
assert pool.page_table("A")[-1] != pool.page_table("B")[1]
pool.release("A")
assert pool.snapshot() == REF["paged"]["trace"]["after_a_release"]

naive = naive_preallocation_waste(32_768, 2_048)
paged = expected_paged_waste(2_048, 16)
assert math.isclose(naive, REF["paged"]["naive_waste_fraction"], abs_tol=1e-12)
assert math.isclose(paged, REF["paged"]["expected_paged_waste_fraction"], abs_tol=1e-12)
print(f"naive waste={naive:.2%}; expected paged waste={paged:.3%}")
print('Совместный префикс, копирование при записи, освобождение и потери проверены')

def prefill_floor_seconds(parameter_count: int, prompt_tokens: int,
                          peak_flops_per_second: float) -> float:
    """Return the dense 2*P*N/peak_FLOPs arithmetic floor."""

    parameters = positive_int("parameter_count", parameter_count)
    tokens = non_negative_int("prompt_tokens", prompt_tokens)
    peak = positive_finite("peak_flops_per_second", peak_flops_per_second)
    return 2 * parameters * tokens / peak


full = prefill_floor_seconds(8_000_000_000, 8_192, 990e12)
suffix = prefill_floor_seconds(8_000_000_000, 512, 990e12)
decode_weight = 16_000_000_000 / 3.35e12
spike = full + decode_weight
chunked = suffix + decode_weight
expected = REF["prefix_and_interference"]
assert math.isclose(full * 1e3, expected["full_prefill_ms"], abs_tol=1e-9)
assert math.isclose(suffix * 1e3, expected["cached_suffix_prefill_ms"], abs_tol=1e-9)
assert math.isclose(full / suffix, expected["arithmetic_speedup"], abs_tol=1e-12)
assert math.isclose(spike * 1e3, expected["monolithic_spike_ms"], abs_tol=1e-9)
assert math.isclose(chunked * 1e3, expected["chunked_itl_ms"], abs_tol=1e-9)
print(f"prefill={full*1e3:.3f} ms → suffix={suffix*1e3:.3f} ms")
print(f"monolithic spike={spike*1e3:.3f} ms; chunked floor={chunked*1e3:.3f} ms")
full_causal = prefill_flops(8_000_000_000, 32, 32, 128, 8192) / 990e12
suffix_causal = prefill_flops(8_000_000_000, 32, 32, 128, 512, 7680) / 990e12
assert math.isclose(full_causal*1e3, expected["full_causal_prefill_ms"], abs_tol=1e-9)
assert math.isclose(suffix_causal*1e3, expected["cached_causal_suffix_ms"], abs_tol=1e-9)
# All chunks together contain exactly the same causal query-key pairs.
chunks = [prefill_flops(8_000_000_000, 32, 32, 128, 512, p) for p in range(0,8192,512)]
assert sum(chunks) == prefill_flops(8_000_000_000, 32, 32, 128, 8192)
assert chunks[-1] > chunks[0]
cost = cached_input_cost(7680, 512, 10)
for key, value in cost.items():
    assert math.isclose(value, expected["illustrative_input_cost"][key], abs_tol=1e-9)
print(f"Including attention: {full_causal*1e3:.3f} / {suffix_causal*1e3:.3f} ms; ratio={full_causal/suffix_causal:.3f}")
print("Input-token equivalents:", cost)
print('Арифметика кэширования префикса и взаимного влияния фаз воспроизведена')

def heavy_hitter_indices(attention_mass: Sequence[float], budget: int) -> tuple[int, ...]:
    """Return top-mass indices with deterministic lower-index tie-breaking."""

    values = tuple(attention_mass)
    if not values:
        raise ValueError("attention_mass must be non-empty")
    keep = positive_int("budget", budget)
    if keep > len(values):
        raise ValueError("budget exceeds token count")
    normalized = []
    for index, value in enumerate(values):
        if isinstance(value, bool) or not isinstance(value, (int, float)):
            raise TypeError(f"attention_mass[{index}] must be numeric")
        number = float(value)
        if not math.isfinite(number) or number < 0:
            raise ValueError("attention masses must be finite and non-negative")
        normalized.append(number)
    ranked = sorted(range(len(values)), key=lambda i: (-normalized[i], i))
    return tuple(sorted(ranked[:keep]))


def streaming_keep_indices(length: int, budget: int,
                           sink_tokens: int) -> tuple[int, ...]:
    """Return sink tokens plus the most recent window under a fixed budget."""

    n = non_negative_int("length", length)
    keep = positive_int("budget", budget)
    sinks = non_negative_int("sink_tokens", sink_tokens)
    if sinks > keep:
        raise ValueError("sink_tokens cannot exceed budget")
    if n <= keep:
        return tuple(range(n))
    sink_count = min(sinks, n)
    recent_count = keep - sink_count
    recent_start = max(sink_count, n - recent_count)
    result = tuple(range(sink_count)) + tuple(range(recent_start, n))
    assert len(result) == keep and len(set(result)) == keep
    return result


compression = REF["compression"]
scale = compression["compressed_bits_per_coordinate"] / compression["base_bits_per_coordinate"]
compressed = decode_ceiling(HARDWARE, g8, 131_072, kv_scale=scale)
assert math.isclose(g8.bytes_per_token * scale / KIB,
                    compression["compressed_kib_per_token"], abs_tol=1e-12)
assert math.isclose(compressed.aggregate_tokens_per_second,
                    compression["ceiling_128k_tokens_per_second"], abs_tol=1e-9)

mass = REF["eviction"]["attention_mass"]
heavy = heavy_hitter_indices(mass, 8)
streaming = streaming_keep_indices(len(mass), 8, 4)
lost = tuple(index for index in heavy if index not in streaming)
assert list(heavy) == REF["eviction"]["heavy_hitters"]
assert list(streaming) == REF["eviction"]["streaming"]
assert list(lost) == REF["eviction"]["lost_heavy_hitters"]
print(f"KV: {g8.bytes_per_token/KIB:g} → {g8.bytes_per_token*scale/KIB:g} KiB/token")
print(f"heavy={heavy}; streaming={streaming}; lost={lost}")
print('Трассировка сжатия и вытеснения проверена')

def expect_error(error_type, fn, *args, **kwargs) -> None:
    """Require one declared exception from a negative contract test."""

    try:
        fn(*args, **kwargs)
    except error_type:
        return
    raise AssertionError(f"expected {error_type.__name__} from {fn.__name__}")


expect_error(ValueError, AttentionKVGeometry, "bad", 0, 8, 128, 2)
expect_error(ValueError, kv_payload_bytes, g8, -1)
expect_error(ValueError, kv_payload_bytes, g8, 1, 1, 1/3)
expect_error(ValueError, ServingHardware, 80*GIB, 1.0, 1.0, 1, 80*GIB, 1)
expect_error(ValueError, heavy_hitter_indices, [0.1, 0.2], 3)
expect_error(ValueError, heavy_hitter_indices, [0.1, float("nan")], 1)
expect_error(ValueError, streaming_keep_indices, 16, 4, 5)

test_pool = PagedKVPool(total_blocks=2, block_size=4)
expect_error(KeyError, test_pool.append, "missing", 1)
test_pool.create("x")
expect_error(ValueError, test_pool.create, "x")
test_pool.append("x", 6)
before_oom = test_pool.snapshot()
expect_error(MemoryError, test_pool.append, "x", 3)
assert test_pool.snapshot() == before_oom  # preflight prevents partial mutation
expect_error(ValueError, test_pool.share_prefix, "x", "y", 7)
test_pool.validate()
# A short child survives release or COW of its longer parent.
for action in ("release", "append"):
    unequal = PagedKVPool(4, 4)
    unequal.create("parent"); unequal.append("parent", 3)
    unequal.share_prefix("parent", "child", 1)
    if action == "release":
        unequal.release("parent")
        assert unequal.physical_payload_tokens == 1
    else:
        unequal.append("parent", 1)
        assert unequal.physical_payload_tokens == 5
    unequal.append("child", 2)
    unequal.validate()
fixed = PagedKVPool(128,16); fixed.create("fixed"); fixed.append("fixed",2048)
assert fixed.internal_waste_tokens == 0
expect_error(ValueError, attention_flops, 32, 32, 128, -1)
expect_error(ValueError, cached_input_cost, 7680, 512, 0)
print('Негативные тесты пройдены')

"""Lecture-level geometry, bandwidth bound, and private page allocation."""
from dataclasses import dataclass, field
import math


def count(value, minimum=0):
    if isinstance(value, bool) or not isinstance(value, int) or value < minimum:
        raise ValueError(f"expected an integer >= {minimum}")
    return value


@dataclass(frozen=True)
class KVGeometry:
    attention_layers: int
    kv_heads: int
    head_dim: int
    bytes_per_element: int = 2

    def __post_init__(self):
        for value in (self.attention_layers, self.kv_heads, self.head_dim, self.bytes_per_element):
            count(value, 1)

    @property
    def bytes_per_token(self):
        return 2 * self.attention_layers * self.kv_heads * self.head_dim * self.bytes_per_element


def decode_bandwidth_ceiling(context_tokens, weight_bytes, geometry, bandwidth_Bps,
                             batch_size=1, kv_scale=1.0):
    """Aggregate tok/s, assuming one weight read and one KV read per step."""
    count(context_tokens); count(weight_bytes, 1); count(batch_size, 1)
    for value in (bandwidth_Bps, kv_scale):
        if isinstance(value, bool) or not math.isfinite(value) or value <= 0:
            raise ValueError("bandwidth and scale must be positive and finite")
    kv = geometry.bytes_per_token * context_tokens * batch_size * kv_scale
    if not math.isclose(kv, round(kv), rel_tol=0, abs_tol=1e-9):
        raise ValueError("the scaled payload must contain an integer number of bytes")
    return batch_size * bandwidth_Bps / (weight_bytes + round(kv))


@dataclass
class PagedKV:
    """Private logical pages only; the notebook adds sharing and copy-on-write."""
    n_blocks: int
    block_size: int = 16
    free: list = field(init=False)
    tables: dict = field(default_factory=dict, init=False)
    lengths: dict = field(default_factory=dict, init=False)

    def __post_init__(self):
        count(self.n_blocks, 1); count(self.block_size, 1)
        self.free = list(range(self.n_blocks - 1, -1, -1))

    def append_to(self, seq_id, total_tokens):
        count(total_tokens)
        hash(seq_id)
        if total_tokens < self.lengths.get(seq_id, 0):
            raise ValueError("append_to cannot shorten a sequence")
        table = self.tables.get(seq_id, [])
        needed = (total_tokens + self.block_size - 1) // self.block_size - len(table)
        if needed > len(self.free):
            raise MemoryError("not enough free blocks")
        self.tables.setdefault(seq_id, table)
        table.extend(self.free.pop() for _ in range(needed))
        self.lengths[seq_id] = total_tokens

    def release(self, seq_id):
        self.free.extend(self.tables.pop(seq_id))
        del self.lengths[seq_id]

geometry = KVGeometry(32,8,128)
assert geometry.bytes_per_token == 131072
assert math.isclose(decode_bandwidth_ceiling(8192,16_000_000_000,geometry,3.35e12),196.207722626508)
pages = PagedKV(128,16); pages.append_to("a",2048)
assert len(pages.tables["a"]) == 128
before = (list(pages.free), dict(pages.lengths), {k:list(v) for k,v in pages.tables.items()})
expect_error(MemoryError, pages.append_to, "a", 2049)
assert before == (pages.free,pages.lengths,pages.tables)
pages.release("a"); assert len(pages.free) == 128
expect_error(ValueError, pages.append_to, "x", -1)
print("Lecture core: PASS")

import importlib.util
import subprocess
import sys

missing = []
if importlib.util.find_spec("numpy") is None:
    missing.append("numpy>=1.26")
if importlib.util.find_spec("torch") is None:
    missing.append("torch>=2.5")
if missing:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", *missing])

import numpy as np
import torch
import torch.nn.functional as F
import math

class NumpyKVCacheExplicit:
    """Owned [B,H,T,D] floating state; truncation cannot recover evicted tokens."""

    def __init__(self, *, window: int | None = None) -> None:
        if window is not None and (isinstance(window, bool) or not isinstance(window, int) or window <= 0):
            raise ValueError("window must be a positive integer or None")
        self.window = window
        self.k = self.v = None
        self.start_pos = self.next_pos = 0

    def append(self, k_new, v_new):
        k_new, v_new = np.asarray(k_new), np.asarray(v_new)
        if k_new.ndim != 4 or k_new.shape != v_new.shape or min(k_new.shape) <= 0:
            raise ValueError("K and V need the same non-empty [B,H,T,D] shape")
        if k_new.dtype != v_new.dtype or not np.issubdtype(k_new.dtype, np.floating):
            raise ValueError("K and V need the same floating dtype")
        if self.k is not None:
            if self.k.shape[:2] + self.k.shape[3:] != k_new.shape[:2] + k_new.shape[3:]:
                raise ValueError("cache geometry changed")
            if self.k.dtype != k_new.dtype:
                raise ValueError("cache dtype changed")
        # Validate first; keep old state intact if an input is rejected.
        k = k_new.copy() if self.k is None else np.concatenate([self.k, k_new], axis=-2)
        v = v_new.copy() if self.v is None else np.concatenate([self.v, v_new], axis=-2)
        drop = max(0, k.shape[-2] - self.window) if self.window is not None else 0
        # A slice alone retains its backing allocation. Make owned bounded copies.
        self.k = k[..., drop:, :].copy() if drop else k
        self.v = v[..., drop:, :].copy() if drop else v
        self.start_pos += drop
        self.next_pos += k_new.shape[-2]
        return self.k, self.v

    def truncate(self, absolute_end: int) -> None:
        if isinstance(absolute_end, bool) or not isinstance(absolute_end, int):
            raise ValueError("absolute_end must be an integer")
        if not self.start_pos <= absolute_end <= self.next_pos:
            raise ValueError("rollback lies outside the retained interval")
        keep = absolute_end - self.start_pos
        if self.k is not None:
            self.k = self.k[..., :keep, :].copy()
            self.v = self.v[..., :keep, :].copy()
        self.next_pos = absolute_end

    @property
    def positions(self):
        return np.arange(self.start_pos, self.next_pos, dtype=np.int64)

class TorchKVCacheExplicit:
    """Owned [B,H,T,D] floating state; truncation cannot recover evicted tokens."""

    def __init__(self, *, window: int | None = None) -> None:
        if window is not None and (isinstance(window, bool) or not isinstance(window, int) or window <= 0):
            raise ValueError("window must be a positive integer or None")
        self.window = window
        self.k = self.v = None
        self.start_pos = self.next_pos = 0

    def append(self, k_new, v_new):
        if not isinstance(k_new, torch.Tensor) or not isinstance(v_new, torch.Tensor):
            raise TypeError("K and V must be tensors")
        if k_new.ndim != 4 or k_new.shape != v_new.shape or min(k_new.shape) <= 0:
            raise ValueError("K and V need the same non-empty [B,H,T,D] shape")
        if k_new.dtype != v_new.dtype or k_new.device != v_new.device or not k_new.is_floating_point():
            raise ValueError("K and V need the same floating dtype and device")
        if k_new.requires_grad or v_new.requires_grad:
            raise ValueError("this cache is for inference; detach inputs first")
        if self.k is not None:
            if self.k.shape[:2] + self.k.shape[3:] != k_new.shape[:2] + k_new.shape[3:]:
                raise ValueError("cache geometry changed")
            if self.k.dtype != k_new.dtype or self.k.device != k_new.device:
                raise ValueError("cache dtype or device changed")
        # Validate first; keep old state intact if an input is rejected.
        k = k_new.clone() if self.k is None else torch.cat([self.k, k_new], dim=-2)
        v = v_new.clone() if self.v is None else torch.cat([self.v, v_new], dim=-2)
        drop = max(0, k.shape[-2] - self.window) if self.window is not None else 0
        # A slice alone retains its backing allocation. Make owned bounded copies.
        self.k = k[..., drop:, :].clone() if drop else k
        self.v = v[..., drop:, :].clone() if drop else v
        self.start_pos += drop
        self.next_pos += k_new.shape[-2]
        return self.k, self.v

    def truncate(self, absolute_end: int) -> None:
        if isinstance(absolute_end, bool) or not isinstance(absolute_end, int):
            raise ValueError("absolute_end must be an integer")
        if not self.start_pos <= absolute_end <= self.next_pos:
            raise ValueError("rollback lies outside the retained interval")
        keep = absolute_end - self.start_pos
        if self.k is not None:
            self.k = self.k[..., :keep, :].clone()
            self.v = self.v[..., :keep, :].clone()
        self.next_pos = absolute_end

    @property
    def positions(self):
        device = self.k.device if self.k is not None else None
        return torch.arange(self.start_pos, self.next_pos, dtype=torch.long, device=device)

np_cache = NumpyKVCacheExplicit(window=3); t_cache = TorchKVCacheExplicit(window=3)
for value in range(5):
    k_np = np.full((1,1,1,2), value, dtype=np.float64); v_np = k_np + 100
    np_cache.append(k_np, v_np); t_cache.append(torch.tensor(k_np), torch.tensor(v_np))
np.testing.assert_allclose(np_cache.k, t_cache.k.numpy())
np.testing.assert_array_equal(np_cache.positions, t_cache.positions.numpy())
assert (np_cache.start_pos, np_cache.next_pos) == (2,5)
np_cache.truncate(4); t_cache.truncate(4)
np.testing.assert_allclose(np_cache.k, t_cache.k.numpy())
rng = np.random.default_rng(17)
keys = rng.standard_normal((1,1,6,4)); values = rng.standard_normal((1,1,6,4)); queries = rng.standard_normal((1,1,6,4))
cache = TorchKVCacheExplicit(window=3)
for step in range(6):
    cache.append(torch.tensor(keys[..., step:step+1, :]), torch.tensor(values[..., step:step+1, :]))
    cached = F.scaled_dot_product_attention(torch.tensor(queries[..., step:step+1, :]), cache.k, cache.v)
    lo = max(0, step - 2)
    full = F.scaled_dot_product_attention(torch.tensor(queries[..., step:step+1, :]), torch.tensor(keys[..., lo:step+1, :]), torch.tensor(values[..., lo:step+1, :]))
    torch.testing.assert_close(cached, full)

# Rejected input must not mutate state; rollback cannot resurrect evicted tokens.
for cache_type, convert in ((NumpyKVCacheExplicit, np.asarray), (TorchKVCacheExplicit, torch.tensor)):
    c = cache_type(window=3)
    c.truncate(0)
    x = convert(np.ones((1,1,5,2), dtype=np.float64))
    c.append(x,x)
    assert (c.start_pos,c.next_pos) == (2,5)
    for bad in (-1,1,6,2.5,True):
        try: c.truncate(bad)
        except ValueError: pass
        else: raise AssertionError("invalid rollback accepted")
    assert (c.start_pos,c.next_pos) == (2,5)
    try: c.append(convert(np.zeros((1,2,1,2))),convert(np.zeros((1,2,1,2))))
    except ValueError: pass
    else: raise AssertionError("geometry change accepted")
    c.truncate(2)
    assert c.k.shape[-2] == 0
    c.append(convert(np.ones((1,1,1,2))),convert(np.ones((1,1,1,2))))
    assert (c.start_pos,c.next_pos) == (2,3)
    for window in (0,-1,2.5,True):
        try: cache_type(window=window)
        except ValueError: pass
        else: raise AssertionError("invalid window accepted")
a = NumpyKVCacheExplicit(window=3)
a.append(np.ones((1,1,10,2)),np.ones((1,1,10,2)))
assert a.k.flags.owndata and a.k.nbytes == 3*2*8
t = TorchKVCacheExplicit(window=3)
t.append(torch.ones(1,1,10,2),torch.ones(1,1,10,2))
assert t.k.untyped_storage().nbytes() == t.k.numel()*t.k.element_size()

print("B17 window, ownership, rollback, and invalid inputs: PASS")

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"воспроизводимый тестовый запуск; seed={SEED}")
DEMO_LANGUAGE = 'ru'

plt.rcParams.update({"font.family":"DejaVu Sans", "font.size":12, "axes.titlesize":15, "axes.labelsize":13, "xtick.labelsize":11, "ytick.labelsize":11, "figure.dpi":150})

import hashlib
rng = np.random.default_rng(SEED)
lengths = rng.integers(64, 4097, size=500)
page = 16
allocated = ((lengths + page - 1) // page) * page
reserved = np.full_like(lengths, 4096)
assert np.all(allocated - lengths < page)
print("waste / allocated:", (allocated-lengths).sum()/allocated.sum(),
      "fixed-reservation waste:", (reserved-lengths).sum()/reserved.sum())
assert (2048 + page - 1) // page * page == 2048

# Complete block boundaries. The key includes ALL preceding token IDs.
# No real KV tensors, tokenizer, network, TTL, eviction, or serving engine.
def prefix_key(tokens, end, scope, revision):
    return (scope, revision, tuple(tokens[:end]))

def run_prefix_trace(isolated):
    cache = set(); hit_requests = cached_tokens = input_tokens = 0
    for user in range(100):
        for turn in range(5):
            tokens = [100 + user % 10] + list(range(31)) + [1000+user, 2000+turn] + list(range(14))
            scope = user if isolated else "shared-public-template"
            hit = 0
            for end in range(page, len(tokens)+1, page):
                if prefix_key(tokens,end,scope,"model-r1") in cache:
                    hit = end
                else:
                    break
            for end in range(page,len(tokens)+1,page):
                cache.add(prefix_key(tokens,end,scope,"model-r1"))
            hit_requests += hit > 0
            cached_tokens += hit
            input_tokens += len(tokens)
    return hit_requests/500, cached_tokens/input_tokens

shared = run_prefix_trace(False); isolated = run_prefix_trace(True)
assert np.isclose(shared[0], .98) and np.isclose(isolated[0], .8)
assert np.isclose(shared[1], 490*32/(500*48))
print("request-hit fraction / reused-input fraction:", shared, isolated)
a = list(range(32)); b = a.copy(); b[0] = 999
assert a[16:32] == b[16:32]  # same block content, different preceding context
assert prefix_key(a,32,"scope","r1") != prefix_key(b,32,"scope","r1")
assert prefix_key(a,32,"scope","r1") != prefix_key(a,32,"scope","r2")
print("Context and model-revision identity: PASS")

rng = np.random.default_rng(SEED + 10)
n, d = 128, 32
keys = rng.normal(size=(n,d)); values = rng.normal(size=(n,d))
queries = rng.normal(size=(12,d))
def attention(q,k,v):
    logits = q @ k.T / np.sqrt(k.shape[1])
    logits -= logits.max(axis=-1,keepdims=True)
    prob = np.exp(logits); prob /= prob.sum(axis=-1,keepdims=True)
    return prob @ v, prob

def fake_quantize_rows(x,bits):
    limit = 2**(bits-1)-1
    scale = np.max(np.abs(x),axis=-1,keepdims=True)/limit
    scale = np.where(scale == 0, 1.0, scale)
    codes = np.clip(np.rint(x/scale),-limit,limit).astype(np.int8)
    # Float64 dequantization here is a numerical experiment, not a packed kernel.
    return codes.astype(float)*scale

reference,_ = attention(queries,keys,values)
for bits in (8,3):
    kq,vq = fake_quantize_rows(keys,bits),fake_quantize_rows(values,bits)
    out,_ = attention(queries,kq,vq)
    error = np.linalg.norm(out-reference)/np.linalg.norm(reference)
    ideal_codes = 2*n*d*bits/8
    fp32_scales = 2*n*4
    print(f"{bits} bits: relative output error={error:.5f}; ideal packed codes+scales={ideal_codes+fp32_scales:.0f} B")
    if bits == 8: assert error < .03
np.testing.assert_array_equal(fake_quantize_rows(np.zeros((2,4)),3),0)

# A controlled query shift: past attention favors token 0; the next query needs token 90.
k = np.zeros((n,d)); k[0,0] = 12; k[90,1] = 12
v = np.zeros((n,2)); v[0,0] = 1; v[90,1] = 1
past = np.zeros((1,d)); past[0,0] = 4
future = np.zeros((1,d)); future[0,1] = 4
_, mass = attention(past,k,v)
budget = 16
ranked = np.argsort(-mass[0],kind="stable")[:budget]
streaming = np.r_[np.arange(4),np.arange(n-(budget-4),n)]
exact,_ = attention(future,k,v)
for name,idx in (("past top-mass",ranked),("sink + recent",streaming)):
    out,_ = attention(future,k[idx],v[idx])
    error = np.linalg.norm(out-exact)
    print(name,"retains future target:",90 in idx,"output error:",error)
    assert 90 not in idx and error > .9
# Keeping the target has a low error on this example, but requires future knowledge.
oracle = np.r_[90,np.arange(budget-1)]
recovered,_ = attention(future,k[oracle],v[oracle])
assert np.linalg.norm(recovered-exact) < .03
print("Query-shift failure and oracle control: PASS")