diff --git a/README.md b/README.md index c6c73f1..8e4a440 100644 --- a/README.md +++ b/README.md @@ -170,7 +170,7 @@ Defined in [`cacheseek/reuse/exact_prefix/config.py`](./cacheseek/reuse/exact_pr | `window_chunks` | `W` — the local attention window, measured in chunks | | `sink_chunks` | Pinned window head (the chunks at the head that are never evicted) | | `break_even_k` | Smallest reused-prefix length that pays off; below it, skip the cache and just generate — harmless | -| `quant` | `none` (bf16, lossless default) / `int8` / `int4` — enable quantization only after a quality A/B passes | +| `quant` | `none` (bf16, lossless default) / `kivi-int8` / `kivi-int4` — enable quantization only after a quality A/B passes | | `commit_tier` | Where committed KV lives (e.g. Fluxon DRAM) | ### Approximate reuse — `CacheConfig` @@ -248,6 +248,7 @@ cacheseek/ ├── cacheseek/ │ ├── service/ ← CacheService orchestrator, CacheConfig, Protocol interfaces │ ├── reuse/ ← reuse strategies: approximate (video DiT) / exact_prefix (trie) +│ ├── quant/ ← quantization methods (KIVI, etc.) │ ├── stores/ ← KV / tensor byte stores (memory / local_file / fluxon) │ ├── backends/ ← vector (faiss/qdrant), encoder (Qwen3-VL), metadata, audit │ ├── adapters/ ← framework adapters (telefuser) diff --git a/benchmarks/kvcache_quantization/README.md b/benchmarks/kvcache_quantization/README.md new file mode 100644 index 0000000..ceb8285 --- /dev/null +++ b/benchmarks/kvcache_quantization/README.md @@ -0,0 +1,100 @@ +# LingBot Exact-Reuse KV Quantization Benchmark + +`lingbot_exact_reuse.py` measures whether KIVI-compressed exact-prefix KV reuse keeps LingBot World output quality close to an unquantized reuse baseline while reducing forked-request latency versus a cold fork. + +The script runs five requests: + +- `main_cache_seed_quant`: cold main trajectory that writes the selected KIVI KV cache. +- `fork_cache_reuse_quant`: forked trajectory that reuses the quantized prefix. +- `main_cache_seed_none`: cold main trajectory for the unquantized baseline. +- `fork_cache_reuse_none`: forked trajectory that reuses unquantized prefix KV. +- `fork_cold_reference`: forked trajectory with an empty cache. + +Quality metrics are aligned by absolute chunk index, so the fast-forwarded suffix is compared against the matching cold-reference chunks instead of frame zero. + +## Requirements + +- CUDA-capable environment with LingBot/Telefuser dependencies installed. +- `LINGBOT_WORLD_CHECKPOINT_DIR` set to the LingBot World checkpoint directory. +- If `telefuser` is not importable, set `TELEFUSER=/path/to/telefuser-internal`. +- If editable installs resolve to another checkout, set `WORLDKV_REPO_ROOTS=/path/to/telefuser:/path/to/CacheSeek`. + +## Basic Usage + +```bash +export LINGBOT_WORLD_CHECKPOINT_DIR=/path/to/lingbot-world-base-cam +python benchmarks/kvcache_quantization/lingbot_exact_reuse.py \ + --quant kivi_int4 \ + --out-dir /tmp/worldkv_quant_benchmark_kivi_int4 +``` + +`--quant` is required and accepts `kivi_int4` or `kivi_int8`. The unquantized `none` baseline is always run automatically. + +The command prints a JSON summary and writes: + +- `manifest.json`: run configuration, timings, speedups, quality metrics, cache stats, validity checks, and profiling metadata. +- Fork videos for quantized reuse, unquantized reuse, and the cold reference. +- Optional per-frame images when `--save-frames` is enabled. + +The process exits with status `0` only when all prefix-hit, alignment, and timing-validity checks pass. + +## Useful Options + +- `--frame-num N`: total requested frames, default `37`. +- `--prefix-chunks N`: reusable prefix chunks before the fork, default `2`. +- `--prompt TEXT`, `--seed N`: generation inputs. +- `--store inmem|localdisk|fluxon`: cache backing store, default `inmem`. +- `--group-size N`: KIVI quantization group size, default `64`. +- `--video-format mp4|gif`: output video format, default `mp4`. +- `--aux-device cuda:1`: place text encoder and VAE on another CUDA device. + +## Profiling + +Benchmark timings are collected without profiler overhead. To diagnose a stage, run a separate profiling pass after the benchmark: + +```bash +python benchmarks/kvcache_quantization/lingbot_exact_reuse.py \ + --quant kivi_int4 \ + --profile cprofile \ + --profile-target fork_cache_reuse_quant \ + --profile-scope create_runtime +``` + +Use `--torch-profile fork_cache_reuse_quant` for a separate `torch.profiler` trace. cProfile and torch profiling are intentionally mutually exclusive. + +To choose where diagnostic artifacts are saved, pass `--profile-dir` for cProfile outputs or `--torch-profile-dir` for torch traces: + +```bash +python benchmarks/kvcache_quantization/lingbot_exact_reuse.py \ + --quant kivi_int4 \ + --torch-profile fork_cache_reuse_quant \ + --profile-scope create_runtime \ + --torch-profile-dir /tmp/worldkv_quant_traces/kivi_int4 +``` + +## Repeated Trials + +This script currently runs one benchmark suite per invocation and records `"repeated_trials": false` in `manifest.json`. It does not yet have an in-script trial loop for repeated measurements, aggregation, or variance reporting. + +For timing numbers intended for comparison, keep profiler flags disabled. Use a separate invocation only when collecting cProfile or torch traces, because profiler overhead is intentionally excluded from benchmark timings. + +## A Thorough Example + +```bash +export LINGBOT_WORLD_CHECKPOINT_DIR=/path/to/lingbot-world-base-cam +export ASSETS=/path/to/lingbot-world-assets +source /path/to/telefuser-env/bin/activate +python benchmarks/kvcache_quant/lingbot_exact_prefix_quant.py \ + --image-path "$ASSETS/image.jpg" \ + --action-path "$ASSETS" \ + --quant kivi_int8 \ + --group-size 64 \ + --frame-num 45 \ + --prefix-chunks 3 \ + --store inmem \ + --out-dir /tmp/benchmark_int8 \ + --torch-profile-chrome-trace \ + --torch-profile fork_cache_reuse_quant \ + --profile-scope create_runtime \ + --no-torch-profile-detail \ +``` diff --git a/benchmarks/kvcache_quantization/lingbot_exact_reuse.py b/benchmarks/kvcache_quantization/lingbot_exact_reuse.py new file mode 100644 index 0000000..4501f39 --- /dev/null +++ b/benchmarks/kvcache_quantization/lingbot_exact_reuse.py @@ -0,0 +1,1616 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +# ruff: noqa: E402 (imports intentionally follow the sys.path/editable-finder shim) +"""Quality/performance benchmark for LingBot exact-prefix KV quantization. + +The benchmark deliberately separates *measurement* from *profiling*: + +* The benchmark suite always runs without cProfile or ``torch.profiler`` so + request latency is not polluted by profiler overhead. +* Optional profiler passes are executed afterwards on fresh cache managers. + Those passes exist only for bottleneck diagnosis and are never used in the + reported speed comparison. + +The unprofiled benchmark contains two cache-reuse branches plus one cold +reference: + + main_cache_seed_quant + Main trajectory, empty cache, full denoise, writes the selected KIVI + representation. + + fork_cache_reuse_quant + Forked trajectory, reuses the selected quantized prefix KV, then denoises + only the suffix. + + main_cache_seed_none + Internal seed run for the automatically included unquantized baseline. + + fork_cache_reuse_none + Forked trajectory, reuses an FP/unquantized prefix KV. This baseline is + always measured, so ``--quant none`` is intentionally not a valid CLI + choice. + + fork_cold_reference + Same forked trajectory with an empty cache and full denoising. + +Generated images are indexed by the runtime's *absolute chunk index*. Quality +metrics compare only common absolute chunks, preventing a fast-forwarded suffix +from being accidentally compared with frame zero of the cold reference. The +manifest records per-phase latency, aligned quality metrics, cache occupancy, +profiling metadata, and the checks required for a valid prefix-fork run. +""" +from __future__ import annotations + +import argparse +import cProfile +import hashlib +import inspect +import json +import math +import os +import pstats +import sys +import time +from pathlib import Path +from typing import Any, Callable, TypeVar + +_WORLDKV_ROOTS = [r for r in os.environ.get("WORLDKV_REPO_ROOTS", "").split(":") if r] +_WORLDKV_REAL_ROOTS = [os.path.realpath(r) for r in _WORLDKV_ROOTS] +if _WORLDKV_ROOTS: + sys.meta_path = [ + f for f in sys.meta_path + if not any( + tag in getattr(f, "__module__", "") + for tag in ("__editable___telefuser", "__editable___cacheseek") + ) + ] + for _r in reversed(_WORLDKV_ROOTS): + if _r not in sys.path: + sys.path.insert(0, _r) + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +import numpy as np +import torch +from PIL import Image + +try: + import telefuser # noqa: F401 +except ImportError: + _tf = os.environ.get("TELEFUSER", "") + assert _tf, "cannot import telefuser: set TELEFUSER=/path/to/telefuser-internal" + sys.path.insert(0, _tf) + +from telefuser.pipelines.lingbot_world_fast import ( + LingBotWorldFastPipeline, + LingBotWorldFastPipelineConfig, + LingBotWorldFastSessionConfig, +) + +from cacheseek.reuse.exact_prefix import NamespaceForest, WorldKVManager +from cacheseek.reuse.exact_prefix.telefuser_lingbot import ( + LingBotWorldKVBinding, + make_world_kv_config, +) +from cacheseek.stores import InMemoryTierStore + +if _WORLDKV_ROOTS: + import cacheseek as _cs + import telefuser as _tf + + def _is_under_worldkv_root(path: str) -> bool: + real_path = os.path.realpath(path) + return any( + real_path == root or real_path.startswith(root + os.sep) + for root in _WORLDKV_REAL_ROOTS + ) + + for _mod in (_tf, _cs): + _src = inspect.getsourcefile(_mod) or "" + assert _is_under_worldkv_root(_src), ( + f"{_mod.__name__} resolved to {_src} (editable finder not stripped?)" + ) + print(f"[worldkv-shim] {_mod.__name__} -> {_src}", flush=True) + +DEFAULT_PROMPT = "a quiet stone courtyard, ancient walls, soft afternoon light" +DEFAULT_SEED = 42 +ORIG_H, ORIG_W = 720, 1280 + +# The LingBot control path maps one reusable KV chunk to three latent frames, +# with a temporal expansion factor of four and two initial conditioning frames. +# Keep these values named because the fork position must remain consistent with +# the binding's chunk geometry. +KV_CHUNK_SIZE = 3 +TEMPORAL_DOWNSAMPLE = 4 +INITIAL_FRAME_OFFSET = 2 + +RUN_MAIN_CACHE_SEED_QUANT = "main_cache_seed_quant" +RUN_FORK_CACHE_REUSE_QUANT = "fork_cache_reuse_quant" +RUN_MAIN_CACHE_SEED_NONE = "main_cache_seed_none" +RUN_FORK_CACHE_REUSE_NONE = "fork_cache_reuse_none" +RUN_FORK_COLD_REFERENCE = "fork_cold_reference" +RUN_DESCRIPTIONS = { + RUN_MAIN_CACHE_SEED_QUANT: ( + "main trajectory cold run that populates the selected quantized KV cache" + ), + RUN_FORK_CACHE_REUSE_QUANT: ( + "forked trajectory reusing the selected quantized exact-prefix KV" + ), + RUN_MAIN_CACHE_SEED_NONE: ( + "main trajectory cold run that populates the unquantized baseline cache" + ), + RUN_FORK_CACHE_REUSE_NONE: ( + "forked trajectory reusing an unquantized exact-prefix KV baseline" + ), + RUN_FORK_COLD_REFERENCE: "forked trajectory cold reference with an empty cache", +} + +PROFILE_RUN_TARGETS = ( + RUN_MAIN_CACHE_SEED_QUANT, + RUN_FORK_CACHE_REUSE_QUANT, + RUN_MAIN_CACHE_SEED_NONE, + RUN_FORK_CACHE_REUSE_NONE, + RUN_FORK_COLD_REFERENCE, +) +PROFILE_SCOPES = ("full_request", "create_runtime", "generation", "chunk") + +T = TypeVar("T") + +PROFILE_TARGETS: tuple[tuple[str, str, str], ...] = ( + ( + "manager_encode_payload", + "cacheseek/reuse/exact_prefix/manager.py", + "_encode_kv_payload", + ), + ( + "manager_decode_payload", + "cacheseek/reuse/exact_prefix/manager.py", + "_decode_layer_payload", + ), + ("kivi_encode_layer", "cacheseek/quant/kivi.py", "encode_layer"), + ("kivi_decode_layer", "cacheseek/quant/kivi.py", "decode_layer"), + ("kivi_encode_tensor", "cacheseek/quant/kivi.py", "_encode_tensor"), + ("kivi_decode_tensor", "cacheseek/quant/kivi.py", "_decode_tensor"), + ("kivi_quantize_grouped", "cacheseek/quant/kivi.py", "_quantize_grouped"), + ("kernel_quantize_grouped", "cacheseek/quant/kernel.py", "quantize_grouped"), + ("kernel_dequantize_grouped", "cacheseek/quant/kernel.py", "dequantize_grouped"), + ("kernel_pack_int4", "cacheseek/quant/kernel.py", "pack_int4_to_int32"), + ("kernel_unpack_int4", "cacheseek/quant/kernel.py", "unpack_int4_from_int32"), +) + + +def _profile_path(path: str) -> str: + return path.replace(os.sep, "/") + + +def _profile_target_label(filename: str, funcname: str) -> str | None: + normalized = _profile_path(filename) + for label, suffix, target_func in PROFILE_TARGETS: + if normalized.endswith(suffix) and funcname == target_func: + return label + return None + + +def write_cprofile_outputs( + profiler: cProfile.Profile, + out_dir: Path, + tag: str, + *, + sort_by: str, + top_n: int, +) -> dict[str, Any]: + out_dir.mkdir(parents=True, exist_ok=True) + stats_path = out_dir / f"{tag}.prof" + summary_path = out_dir / f"{tag}.txt" + profiler.dump_stats(str(stats_path)) + + stats = pstats.Stats(profiler) + targets: list[dict[str, Any]] = [] + for (filename, line, funcname), raw in stats.stats.items(): + primitive_calls, total_calls, inline_s, cumulative_s, _callers = raw + label = _profile_target_label(filename, funcname) + if label is None: + continue + targets.append( + { + "label": label, + "file": _profile_path(filename), + "line": int(line), + "function": funcname, + "primitive_calls": int(primitive_calls), + "total_calls": int(total_calls), + "inline_s": float(inline_s), + "cumulative_s": float(cumulative_s), + } + ) + targets.sort(key=lambda row: row["cumulative_s"], reverse=True) + + with summary_path.open("w", encoding="utf-8") as fh: + fh.write(f"# cProfile summary for request {tag}\n") + fh.write(f"# sort={sort_by} top_n={top_n}\n\n") + pstats.Stats(profiler, stream=fh).strip_dirs().sort_stats(sort_by).print_stats(top_n) + + return { + "stats_path": str(stats_path), + "summary_path": str(summary_path), + "targets": targets, + } + + +class RequestProfileController: + """Profile exactly one semantic request stage. + + A controller is attached only to the separate diagnostic pass. The normal + benchmark passes use no controller, which keeps their timing free from + profiler instrumentation. ``run()`` can be nested: when a whole request or + the generation loop is under ``torch.profiler``, nested stages are emitted as + readable ``world_kv:*`` annotations instead of starting another profiler. + """ + + def __init__( + self, + *, + kind: str, + tag: str, + scope: str, + chunk_index: int, + cprofile_dir: Path | None, + torch_profile_dir: Path | None, + sort_by: str, + top_n: int, + export_chrome_trace: bool, + torch_detail: bool, + ) -> None: + self.kind = kind + self.tag = tag + self.scope = scope + self.chunk_index = int(chunk_index) + self.cprofile_dir = cprofile_dir + self.torch_profile_dir = torch_profile_dir + self.sort_by = sort_by + self.top_n = max(int(top_n), 1) + self.export_chrome_trace = bool(export_chrome_trace) + self.torch_detail = bool(torch_detail) + + self._captured = False + self._torch_active = False + self._cprofiler = cProfile.Profile() if kind == "cprofile" else None + self._torch_profiler: Any | None = None + self._torch_activities: list[Any] = [] + + @property + def enabled(self) -> bool: + return self.kind != "none" + + def _matches(self, stage: str, chunk_index: int | None) -> bool: + """Return whether this stage is the single requested capture window.""" + + if self._captured: + return False + if self.scope == "chunk": + return stage == "chunk" and chunk_index == self.chunk_index + return stage == self.scope + + @staticmethod + def _stage_label(stage: str, chunk_index: int | None) -> str: + if stage == "chunk" and chunk_index is not None: + return f"world_kv:chunk:{chunk_index}" + return f"world_kv:{stage}" + + def run( + self, + stage: str, + fn: Callable[[], T], + *, + chunk_index: int | None = None, + ) -> T: + """Run ``fn`` and profile it only when it matches the configured scope.""" + + label = self._stage_label(stage, chunk_index) + + # A parent torch-profiler scope is already collecting. Emit nested + # semantic ranges so Perfetto can collapse/search the KV-specific stages. + if self._torch_active: + with torch.profiler.record_function(label): + return fn() + + if not self._matches(stage, chunk_index): + return fn() + + self._captured = True + if self.kind == "cprofile": + assert self._cprofiler is not None + self._cprofiler.enable() + try: + return fn() + finally: + self._cprofiler.disable() + + if self.kind == "torch": + self._torch_activities = [torch.profiler.ProfilerActivity.CPU] + if torch.cuda.is_available(): + self._torch_activities.append(torch.profiler.ProfilerActivity.CUDA) + + trace_handler = None + run_trace_dir = self._torch_run_dir() + if not self.export_chrome_trace: + trace_handler = torch.profiler.tensorboard_trace_handler(str(run_trace_dir)) + + # Shape, memory and Python-stack collection create a much denser + # trace. They are opt-in because the default use case is locating + # KV lookup/decode/copy gaps rather than inspecting every operator. + with torch.profiler.profile( + activities=self._torch_activities, + record_shapes=self.torch_detail, + profile_memory=self.torch_detail, + with_stack=self.torch_detail, + on_trace_ready=trace_handler, + ) as profiler: + self._torch_profiler = profiler + self._torch_active = True + try: + with torch.profiler.record_function(label): + return fn() + finally: + self._torch_active = False + + return fn() + + def _torch_run_dir(self) -> Path: + assert self.torch_profile_dir is not None + run_dir = self.torch_profile_dir / self.tag + run_dir.mkdir(parents=True, exist_ok=True) + return run_dir + + def finalize(self) -> dict[str, Any]: + """Write profiler artifacts after the diagnostic request has completed.""" + + base = { + "kind": self.kind, + "target": self.tag, + "scope": self.scope, + "chunk_index": self.chunk_index if self.scope == "chunk" else None, + "captured": self._captured, + "timing_valid_for_speedup": False, + } + if not self._captured: + base["warning"] = ( + "The requested profile scope was not reached; for a chunk scope, " + "the chunk may have been skipped by exact-prefix fast-forward." + ) + return base + + artifact_tag = f"{self.tag}.{self.scope}" + if self.scope == "chunk": + artifact_tag += f"_{self.chunk_index}" + + if self.kind == "cprofile": + assert self._cprofiler is not None + assert self.cprofile_dir is not None + outputs = write_cprofile_outputs( + self._cprofiler, + self.cprofile_dir, + artifact_tag, + sort_by=self.sort_by, + top_n=self.top_n, + ) + return {**base, **outputs} + + if self.kind == "torch": + assert self._torch_profiler is not None + run_dir = self._torch_run_dir() + chrome_trace_path = None + if self.export_chrome_trace: + chrome_trace_path = run_dir / f"{artifact_tag}.chrome_trace.json" + self._torch_profiler.export_chrome_trace(str(chrome_trace_path)) + return { + **base, + "trace_dir": str(run_dir), + "chrome_trace_path": ( + str(chrome_trace_path) if chrome_trace_path is not None else None + ), + "activities": [activity.name for activity in self._torch_activities], + "detail": self.torch_detail, + } + + return base + + +def run_profiled_request( + controller: RequestProfileController, + fn: Callable[[RequestProfileController], dict[str, Any]], +) -> tuple[dict[str, Any], dict[str, Any]]: + """Execute one diagnostic request and return its result plus artifacts.""" + + result = controller.run("full_request", lambda: fn(controller)) + return result, controller.finalize() + + +def parse_group_axis(axis: str) -> str | int: + axis = axis.strip() + if axis.lstrip("-").isdigit(): + return int(axis) + return axis + + +def load_image(image_path: str) -> Image.Image: + if image_path: + return Image.open(image_path).convert("RGB") + y, x = np.mgrid[0:ORIG_H, 0:ORIG_W] + arr = np.stack( + [ + (x / ORIG_W * 255).astype(np.uint8), + (y / ORIG_H * 255).astype(np.uint8), + ((x + y) % 256).astype(np.uint8), + ], + axis=-1, + ) + arr[180:300, 320:520] = (200, 60, 60) + return Image.fromarray(arr) + + +def make_trajectory(frame_num: int, action_path: str) -> tuple[np.ndarray, np.ndarray]: + if action_path: + root = Path(action_path) + poses = np.load(root / "poses.npy").astype(np.float32)[:frame_num] + intrinsics = np.load(root / "intrinsics.npy").astype(np.float32) + if intrinsics.ndim > 1: + intrinsics = intrinsics[:frame_num] + assert len(poses) >= frame_num, ( + f"poses.npy has only {len(poses)} frames < frame_num={frame_num}" + ) + return poses, intrinsics + + poses = np.tile(np.eye(4, dtype=np.float32), (frame_num, 1, 1)) + for i in range(frame_num): + poses[i, 2, 3] = 0.05 * i + intrinsics = np.tile( + np.array([400.0, 400.0, ORIG_W / 2, ORIG_H / 2], dtype=np.float32), + (frame_num, 1), + ) + return poses, intrinsics + + +def _apply_fork(poses: np.ndarray, fork_from: int) -> np.ndarray: + """Apply a pure yaw rotation after fork_from frames. + + The prefix remains byte-identical at the control-key level. Avoid changing + translation because LingBot normalizes camera translation over the whole + trajectory, which would contaminate the prefix and correctly miss the cache. + """ + + forked = poses.copy() + for i in range(fork_from, len(forked)): + j = i - fork_from + yaw = 0.03 * j + c, s = np.cos(yaw), np.sin(yaw) + rot = np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]], dtype=forked.dtype) + forked[i, :3, :3] = rot @ forked[i, :3, :3] + return forked + + +def make_trajectories( + frame_num: int, + prefix_chunks: int, + action_path: str, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Build byte-identical prefix controls followed by a yaw-only fork.""" + + poses, intrinsics = make_trajectory(frame_num, action_path) + fork_from = ( + prefix_chunks * KV_CHUNK_SIZE * TEMPORAL_DOWNSAMPLE + + INITIAL_FRAME_OFFSET + ) + if fork_from >= frame_num: + raise ValueError( + f"fork_from={fork_from} must be smaller than frame_num={frame_num}; " + "increase --frame-num or reduce --prefix-chunks" + ) + return poses, _apply_fork(poses, fork_from), intrinsics + + +def cache_stats(forest: NamespaceForest) -> dict[str, int]: + """Summarize cache blobs reachable from the current namespace forest.""" + + ready_nodes = 0 + ready_bytes = 0 + total_nodes = 0 + for namespace in forest.snapshot().get("namespaces", []): + for node in namespace.get("nodes", []): + total_nodes += 1 + blob = node.get("blob") + if blob is not None: + ready_nodes += 1 + ready_bytes += int(blob.get("nbytes", 0)) + return { + "ready_blob_nodes": ready_nodes, + "ready_blob_bytes": ready_bytes, + "total_nodes": total_nodes, + } + + +def _cuda_synchronize() -> None: + """Synchronize only at coarse phase boundaries used for wall-clock timing.""" + + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _run_profile_stage( + controller: RequestProfileController | None, + stage: str, + fn: Callable[[], T], + *, + chunk_index: int | None = None, +) -> T: + """Dispatch a stage through the optional diagnostic profiler controller.""" + + if controller is None: + return fn() + return controller.run(stage, fn, chunk_index=chunk_index) + + +def run_request( + pipeline: Any, + binding: LingBotWorldKVBinding, + forest: NamespaceForest, + *, + prompt: str, + seed: int, + frame_num: int, + poses: np.ndarray, + intrinsics: np.ndarray, + image_path: str, + tag: str, + profile_controller: RequestProfileController | None = None, +) -> dict[str, Any]: + """Execute one request and retain generated frames by absolute chunk index. + + ``create_runtime_s`` includes exact-prefix lookup/materialization and any + runtime initialization. ``generation_s`` covers only generated suffix + chunks. ``flush_s`` is kept separate because asynchronous stores may defer + persistence beyond the online generation path. CUDA is synchronized only + at these coarse boundaries; per-chunk timing remains host-observed so the + benchmark does not introduce a synchronization after every chunk. + """ + + cache_before = cache_stats(forest) + session = LingBotWorldFastSessionConfig( + prompt=prompt, + image=load_image(image_path), + control_mode="cam", + frame_num=frame_num, + seed=seed, + sample_shift=10, # Set 10 for better quality + poses=poses, + intrinsics=intrinsics, + world_kv_binding=binding, + ) + + _cuda_synchronize() + request_start = time.perf_counter() + + create_start = time.perf_counter() + runtime = _run_profile_stage( + profile_controller, + "create_runtime", + lambda: pipeline.create_runtime(session), + ) + _cuda_synchronize() + create_runtime_s = time.perf_counter() - create_start + + # A reuse request may begin at ``prefix_chunks`` rather than zero. Store + # this initial absolute index before generation mutates the runtime cursor. + initial_chunk_index = int(runtime.current_chunk_index) + frames: list[Image.Image] = [] + frames_by_chunk: dict[int, list[Image.Image]] = {} + chunk_host_times_s: dict[int, float] = {} + + def _generate_all_chunks() -> None: + while runtime.active: + # Capture the cursor before ``generate_next_chunk`` advances it. + # This is the absolute chunk id used to align reuse and cold output. + chunk_index = int(runtime.current_chunk_index) + chunk_start = time.perf_counter() + chunk_frames = list( + _run_profile_stage( + profile_controller, + "chunk", + lambda: pipeline.generate_next_chunk(runtime), + chunk_index=chunk_index, + ) + ) + chunk_host_times_s[chunk_index] = time.perf_counter() - chunk_start + frames_by_chunk[chunk_index] = chunk_frames + frames.extend(chunk_frames) + + generation_start = time.perf_counter() + _run_profile_stage( + profile_controller, + "generation", + _generate_all_chunks, + ) + _cuda_synchronize() + generation_s = time.perf_counter() - generation_start + + store = binding.mgr.store + + def _flush_store() -> None: + if hasattr(store, "flush"): + store.flush() + + flush_start = time.perf_counter() + _run_profile_stage(profile_controller, "store_flush", _flush_store) + flush_s = time.perf_counter() - flush_start + total_with_flush_s = time.perf_counter() - request_start + cache_after = cache_stats(forest) + + return { + "tag": tag, + "frames": frames, + "frames_by_chunk": frames_by_chunk, + "fast_forward_k": binding.last_fast_forward, + "initial_chunk_index": initial_chunk_index, + "final_chunk_index": int(runtime.current_chunk_index), + "generated_chunk_indices": sorted(frames_by_chunk), + "generated_chunk_frame_counts": { + str(chunk_index): len(chunk_frames) + for chunk_index, chunk_frames in sorted(frames_by_chunk.items()) + }, + "n_generated_chunks": len(frames_by_chunk), + "n_frames": len(frames), + "create_runtime_s": create_runtime_s, + "generation_s": generation_s, + "flush_s": flush_s, + "total_without_flush_s": create_runtime_s + generation_s, + "total_with_flush_s": total_with_flush_s, + "timing_valid_for_speedup": profile_controller is None, + "chunk_host_times_s": { + str(chunk_index): duration + for chunk_index, duration in sorted(chunk_host_times_s.items()) + }, + "cache": { + "before": cache_before, + "after": cache_after, + "delta_ready_blob_bytes": ( + cache_after["ready_blob_bytes"] - cache_before["ready_blob_bytes"] + ), + }, + "frame_sha256_by_chunk": { + str(chunk_index): [ + hashlib.sha256(np.asarray(frame).tobytes()).hexdigest() + for frame in chunk_frames + ] + for chunk_index, chunk_frames in sorted(frames_by_chunk.items()) + }, + } + + +def save_gif(frames: list[Image.Image], path: Path, fps: int) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + duration_ms = max(int(1000 / max(fps, 1)), 1) + frames[0].save( + path, + save_all=True, + append_images=frames[1:], + duration=duration_ms, + loop=0, + ) + return path + + +def save_video(frames: list[Image.Image], path: Path, fps: int) -> dict[str, Any]: + path.parent.mkdir(parents=True, exist_ok=True) + if path.suffix.lower() == ".gif": + return {"path": str(save_gif(frames, path, fps)), "format": "gif"} + arr = np.stack([np.asarray(f.convert("RGB")) for f in frames], axis=0) + try: + try: + import imageio.v3 as iio + + iio.imwrite(path, arr, fps=fps, codec="libx264", quality=8) + except TypeError: + import imageio.v3 as iio + + iio.imwrite(path, arr, fps=fps) + return {"path": str(path), "format": path.suffix.lstrip(".").lower()} + except Exception as exc: # pragma: no cover - optional codec dependent + fallback = path.with_suffix(".gif") + save_gif(frames, fallback, fps) + return { + "path": str(fallback), + "format": "gif", + "requested_path": str(path), + "warning": f"MP4 save failed; wrote GIF fallback: {type(exc).__name__}: {exc}", + } + + +def save_frames_by_chunk( + frames_by_chunk: dict[int, list[Image.Image]], + root: Path, +) -> None: + """Save frames with absolute chunk ids encoded in their file names.""" + + root.mkdir(parents=True, exist_ok=True) + for chunk_index, chunk_frames in sorted(frames_by_chunk.items()): + for frame_index, frame in enumerate(chunk_frames): + frame.save( + root / f"chunk_{chunk_index:04d}_frame_{frame_index:04d}.jpg", + quality=92, + ) + + +def frame_metrics(a: list[Image.Image], b: list[Image.Image]) -> dict[str, Any]: + """Compare two frame lists without silently hiding a length mismatch. + + Metrics are computed over the common prefix only, but both original lengths + and ``frame_count_match`` are returned. Callers therefore cannot mistake a + truncated comparison for a complete one. ``global_psnr`` is derived from + the aggregate MSE, while ``finite_psnr_*`` excludes exactly equal frames + whose mathematical PSNR is infinite. + """ + + n = min(len(a), len(b)) + base = { + "frames_a": len(a), + "frames_b": len(b), + "frame_count_match": len(a) == len(b), + "n_compared_frames": n, + } + if n == 0: + return { + **base, + "byte_equal_frames": 0, + "first_different_frame": None, + "mae_mean": None, + "mae_max": None, + "mse_mean": None, + "mse_max": None, + "global_psnr": None, + "finite_psnr_mean": None, + "finite_psnr_min": None, + } + + maes: list[float] = [] + mses: list[float] = [] + finite_psnrs: list[float] = [] + equal = 0 + first_diff: int | None = None + squared_error_sum = 0.0 + pixel_count = 0 + + for i in range(n): + aa = np.asarray(a[i].convert("RGB"), dtype=np.float32) + bb = np.asarray(b[i].convert("RGB"), dtype=np.float32) + if aa.shape != bb.shape: + raise ValueError( + f"frame shape mismatch at index {i}: {aa.shape} != {bb.shape}" + ) + + diff = aa - bb + absolute = np.abs(diff) + squared = diff * diff + mae = float(np.mean(absolute)) + mse = float(np.mean(squared)) + maes.append(mae) + mses.append(mse) + squared_error_sum += float(np.sum(squared, dtype=np.float64)) + pixel_count += int(squared.size) + + if mse == 0: + equal += 1 + else: + finite_psnrs.append(float(20 * math.log10(255.0 / math.sqrt(mse)))) + if first_diff is None: + first_diff = i + + global_mse = squared_error_sum / pixel_count + global_psnr = ( + None + if global_mse == 0 + else float(20 * math.log10(255.0 / math.sqrt(global_mse))) + ) + return { + **base, + "byte_equal_frames": equal, + "first_different_frame": first_diff, + "mae_mean": float(np.mean(maes)), + "mae_max": float(np.max(maes)), + "mse_mean": float(np.mean(mses)), + "mse_max": float(np.max(mses)), + "global_psnr": global_psnr, + "finite_psnr_mean": ( + None if not finite_psnrs else float(np.mean(finite_psnrs)) + ), + "finite_psnr_min": ( + None if not finite_psnrs else float(np.min(finite_psnrs)) + ), + } + + +def aligned_chunk_metrics( + a_by_chunk: dict[int, list[Image.Image]], + b_by_chunk: dict[int, list[Image.Image]], +) -> dict[str, Any]: + """Compare only matching absolute chunks from two generation requests. + + Exact-prefix reuse skips already materialized chunks. Consequently, list + position zero in a reuse result is not necessarily absolute chunk zero. By + intersecting the runtime chunk ids first, this function compares like with + like and also exposes missing chunks and per-chunk frame-count mismatches. + """ + + chunks_a = sorted(a_by_chunk) + chunks_b = sorted(b_by_chunk) + common_chunks = sorted(set(chunks_a) & set(chunks_b)) + only_a = sorted(set(chunks_a) - set(chunks_b)) + only_b = sorted(set(chunks_b) - set(chunks_a)) + + per_chunk = { + str(chunk_index): frame_metrics( + a_by_chunk[chunk_index], + b_by_chunk[chunk_index], + ) + for chunk_index in common_chunks + } + frame_count_match_per_chunk = all( + metrics["frame_count_match"] for metrics in per_chunk.values() + ) + + # Flatten only common chunks, preserving absolute chunk order, to provide a + # single overall quality number in addition to the diagnostic per-chunk view. + aligned_a: list[Image.Image] = [] + aligned_b: list[Image.Image] = [] + for chunk_index in common_chunks: + # Truncate independently inside each absolute chunk. Flattening the + # complete unequal lists would shift every later chunk by one frame and + # recreate the same temporal-alignment bug this function is meant to + # prevent. + n = min( + len(a_by_chunk[chunk_index]), + len(b_by_chunk[chunk_index]), + ) + aligned_a.extend(a_by_chunk[chunk_index][:n]) + aligned_b.extend(b_by_chunk[chunk_index][:n]) + return { + "chunks_a": chunks_a, + "chunks_b": chunks_b, + "common_chunks": common_chunks, + "chunks_only_in_a": only_a, + "chunks_only_in_b": only_b, + "has_common_chunks": bool(common_chunks), + "frame_count_match_per_chunk": frame_count_match_per_chunk, + "overall": frame_metrics(aligned_a, aligned_b), + "per_chunk": per_chunk, + } + + +def strip_frames(result: dict[str, Any]) -> dict[str, Any]: + """Remove PIL objects while retaining all JSON-serializable request data.""" + + return { + key: value + for key, value in result.items() + if key not in {"frames", "frames_by_chunk"} + } + + + +def make_store(args: argparse.Namespace, store_namespace: str) -> Any: + """Create an isolated store for one benchmark branch. + + In-memory managers are naturally isolated. Local-disk managers receive a + unique child directory so an earlier script invocation cannot accidentally + make a nominally cold reference observe stale physical objects. Fluxon is + externally configured; lookup isolation still comes from a fresh forest and + the WorldKV namespace/config fingerprint. + """ + + if args.store == "fluxon": + from cacheseek.stores import TensorStoreTierStore + from cacheseek.stores.fluxon import FluxonKVStore + + assert args.fluxon_config, "--store fluxon requires --fluxon-config" + return TensorStoreTierStore(FluxonKVStore(config_path=args.fluxon_config), async_put=True) + if args.store == "localdisk": + from cacheseek.stores import TensorStoreTierStore + from cacheseek.stores.tier import LocalDiskTensorStore + + disk_root = Path(args.disk_root) / store_namespace + return TensorStoreTierStore(LocalDiskTensorStore(str(disk_root)), async_put=True) + return InMemoryTierStore() + + +def build_manager( + args: argparse.Namespace, + cfg: LingBotWorldFastPipelineConfig, + key_group_axis: str | int, + value_group_axis: str | int, + *, + quant: str, + store_namespace: str, +) -> tuple[NamespaceForest, WorldKVManager]: + """Build one exact-prefix manager for a fixed quantization branch.""" + + forest = NamespaceForest() + mgr = WorldKVManager( + forest, + make_store(args, store_namespace), + make_world_kv_config( + local_attn_size=cfg.local_attn_size, + sink_size=cfg.sink_size, + chunk_size=KV_CHUNK_SIZE, + quant=quant, + group_size=args.group_size, + kv_layout=args.kv_layout, + key_group_axis=key_group_axis, + value_group_axis=value_group_axis, + scale_dtype=args.scale_dtype, + offset_dtype=args.offset_dtype, + ), + ) + return forest, mgr + + +def timing_comparison( + candidate: dict[str, Any], + reference: dict[str, Any], +) -> dict[str, Any]: + """Compare one faster-expected request against a reference request. + + ``speedup`` is reported as ``reference / candidate``; values above one mean + the candidate is faster. The profiler never feeds this function because all + entries supplied by the benchmark suite are collected without instrumentation. + """ + + phases = ( + "create_runtime_s", + "generation_s", + "flush_s", + "total_without_flush_s", + "total_with_flush_s", + ) + comparisons: dict[str, Any] = {} + for phase in phases: + candidate_s = float(candidate[phase]) + reference_s = float(reference[phase]) + comparisons[phase] = { + "candidate_s": candidate_s, + "reference_s": reference_s, + "saved_s": reference_s - candidate_s, + "speedup": None if candidate_s <= 0 else reference_s / candidate_s, + } + return { + "candidate": candidate["tag"], + "reference": reference["tag"], + "timing_valid_for_speedup": bool( + candidate["timing_valid_for_speedup"] + and reference["timing_valid_for_speedup"] + ), + "phases": comparisons, + } + + +def main() -> int: + """Run the unprofiled benchmark, then optional isolated diagnostic passes.""" + + ap = argparse.ArgumentParser( + description=( + "Benchmark selected KIVI exact-prefix reuse, an automatic unquantized " + "reuse baseline, and a cold fork reference." + ) + ) + ap.add_argument( + "--quant", + choices=["kivi_int4", "kivi_int8"], + required=True, + help=( + "Selected quantized mode. The unquantized 'none' branch is always " + "run automatically as the reuse baseline." + ), + ) + ap.add_argument("--frame-num", type=int, default=37) + ap.add_argument("--prefix-chunks", type=int, default=2) + ap.add_argument("--prompt", type=str, default=DEFAULT_PROMPT) + ap.add_argument("--seed", type=int, default=DEFAULT_SEED) + ap.add_argument("--out-dir", type=str, default="") + ap.add_argument("--fps", type=int, default=8) + ap.add_argument("--video-format", choices=["mp4", "gif"], default="mp4") + ap.add_argument("--save-frames", action=argparse.BooleanOptionalAction, default=False) + ap.add_argument("--image-path", type=str, default="") + ap.add_argument("--action-path", type=str, default="") + ap.add_argument( + "--store", + type=str, + default="inmem", + choices=["inmem", "localdisk", "fluxon"], + ) + ap.add_argument("--fluxon-config", type=str, default="") + ap.add_argument("--disk-root", type=str, default="") + ap.add_argument("--group-size", type=int, default=64) + ap.add_argument("--kv-layout", type=str, default="B,T,H,D") + ap.add_argument("--key-group-axis", type=str, default="T") + ap.add_argument("--value-group-axis", type=str, default="D") + ap.add_argument("--scale-dtype", type=str, default="float32") + ap.add_argument("--offset-dtype", type=str, default="float32") + ap.add_argument("--aux-device", type=str, default="") + + # cProfile and torch.profiler are diagnostic-only. Both run after the + # benchmark on fresh managers, so the benchmark's latency remains usable. + ap.add_argument( + "--profile", + choices=["none", "cprofile"], + default="none", + help="Run a separate cProfile diagnostic pass after the benchmark.", + ) + ap.add_argument( + "--profile-target", + choices=[*PROFILE_RUN_TARGETS, "all"], + default=RUN_FORK_CACHE_REUSE_QUANT, + help="Semantic request to rerun under cProfile.", + ) + ap.add_argument( + "--profile-dir", + type=str, + default="", + help="Directory for .prof/.txt artifacts; defaults to /profiles.", + ) + ap.add_argument( + "--profile-sort", + choices=["cumulative", "time", "calls"], + default="cumulative", + help="Sort key for generated cProfile text summaries.", + ) + ap.add_argument("--profile-top-n", type=int, default=80) + ap.add_argument( + "--profile-scope", + choices=PROFILE_SCOPES, + default="create_runtime", + help=( + "Capture only one stage. create_runtime is the recommended scope for " + "exact-prefix lookup/materialize/decode analysis." + ), + ) + ap.add_argument( + "--profile-chunk-index", + type=int, + default=-1, + help=( + "Absolute runtime chunk to capture when --profile-scope=chunk. " + "Default: first generated chunk for the selected target." + ), + ) + ap.add_argument( + "--torch-profile", + choices=["none", *PROFILE_RUN_TARGETS, "all"], + default="none", + help="Run a separate torch.profiler diagnostic pass after the benchmark.", + ) + ap.add_argument( + "--torch-profile-dir", + type=str, + default="", + help="Trace directory; defaults to /torch_traces.", + ) + ap.add_argument( + "--torch-profile-chrome-trace", + action=argparse.BooleanOptionalAction, + default=True, + help=( + "Export Chrome/Perfetto JSON. Disable to use the TensorBoard trace " + "handler instead." + ), + ) + ap.add_argument( + "--torch-profile-detail", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Collect shapes, memory events and Python stacks. Disabled by default " + "to keep Chrome/Perfetto traces readable." + ), + ) + args = ap.parse_args() + + if args.profile != "none" and args.torch_profile != "none": + ap.error( + "cProfile and torch.profiler must be run separately; enabling both " + "would distort the diagnostic trace and duplicate expensive passes" + ) + if args.frame_num <= 0: + ap.error("--frame-num must be positive") + if args.prefix_chunks < 0: + ap.error("--prefix-chunks must be non-negative") + if args.group_size <= 0: + ap.error("--group-size must be positive") + + if not args.out_dir: + args.out_dir = f"/tmp/worldkv_quant_benchmark_{args.quant}" + if not args.disk_root: + args.disk_root = f"/tmp/worldkv_quant_benchmark_diskstore_{args.quant}" + + ckpt = os.environ.get("LINGBOT_WORLD_CHECKPOINT_DIR", "") + assert ckpt, "Set LINGBOT_WORLD_CHECKPOINT_DIR" + out = Path(args.out_dir) + out.mkdir(parents=True, exist_ok=True) + profile_dir = Path(args.profile_dir) if args.profile_dir else out / "profiles" + torch_profile_dir = ( + Path(args.torch_profile_dir) + if args.torch_profile_dir + else out / "torch_traces" + ) + + # Initialize the heavyweight pipeline once. Cache managers remain separate + # per branch, which isolates cache state without reloading model weights. + pipeline = LingBotWorldFastPipeline(device="cuda") + cfg = LingBotWorldFastPipelineConfig(checkpoint_dir=ckpt) + if args.aux_device: + from telefuser.core.config import ModelRuntimeConfig + + _t, _i = (args.aux_device.split(":") + ["0"])[:2] + aux = dict(device_type=_t, device_id=int(_i)) + cfg.text_encoding_config = ModelRuntimeConfig(**aux) + cfg.vae_config = ModelRuntimeConfig(**aux) + print( + f"[2-gpu] text-encoder + VAE -> {args.aux_device}; DiT -> cuda:0", + flush=True, + ) + pipeline.init(None, cfg) + + poses_main, poses_fork, intrinsics = make_trajectories( + args.frame_num, + args.prefix_chunks, + args.action_path, + ) + fork_from_frame = ( + args.prefix_chunks * KV_CHUNK_SIZE * TEMPORAL_DOWNSAMPLE + + INITIAL_FRAME_OFFSET + ) + key_group_axis = parse_group_axis(args.key_group_axis) + value_group_axis = parse_group_axis(args.value_group_axis) + + # The run id creates fresh local-disk child directories without deleting any + # user-provided root. Managers that must share cache state reuse one manager. + run_id = f"pid{os.getpid()}_{time.time_ns()}" + + def _build_branch( + *, + quant: str, + branch: str, + ) -> tuple[NamespaceForest, WorldKVManager]: + return build_manager( + args, + cfg, + key_group_axis, + value_group_axis, + quant=quant, + store_namespace=f"{run_id}_{branch}", + ) + + def _execute( + *, + tag: str, + mgr: WorldKVManager, + forest: NamespaceForest, + poses: np.ndarray, + controller: RequestProfileController | None = None, + ) -> dict[str, Any]: + # Bindings are per request so ``last_fast_forward`` cannot leak from a + # previous trajectory while the manager/forest cache remains shared. + return run_request( + pipeline, + LingBotWorldKVBinding(mgr, forest), + forest, + prompt=args.prompt, + seed=args.seed, + frame_num=args.frame_num, + poses=poses, + intrinsics=intrinsics, + image_path=args.image_path, + tag=tag, + profile_controller=controller, + ) + + # ------------------------------------------------------------------ + # Unprofiled benchmark suite + # ------------------------------------------------------------------ + # Selected KIVI branch: seed the main path, then reuse its exact prefix on + # the forked path. + forest_quant, mgr_quant = _build_branch( + quant=args.quant, + branch="benchmark_quant", + ) + main_cache_seed_quant = _execute( + tag=RUN_MAIN_CACHE_SEED_QUANT, + mgr=mgr_quant, + forest=forest_quant, + poses=poses_main, + ) + fork_cache_reuse_quant = _execute( + tag=RUN_FORK_CACHE_REUSE_QUANT, + mgr=mgr_quant, + forest=forest_quant, + poses=poses_fork, + ) + + # Automatic unquantized baseline. It has its own cache so selected KIVI + # payloads can never be mistaken for the FP/no-quant representation. + forest_none, mgr_none = _build_branch( + quant="none", + branch="benchmark_none", + ) + main_cache_seed_none = _execute( + tag=RUN_MAIN_CACHE_SEED_NONE, + mgr=mgr_none, + forest=forest_none, + poses=poses_main, + ) + fork_cache_reuse_none = _execute( + tag=RUN_FORK_CACHE_REUSE_NONE, + mgr=mgr_none, + forest=forest_none, + poses=poses_fork, + ) + + # A fresh manager and empty forest guarantee no exact-prefix match. The + # manager uses quant=none because this request does not consume cached KV; + # any cache written during the cold run is irrelevant to the comparison. + forest_ref, mgr_ref = _build_branch( + quant="none", + branch="benchmark_cold_reference", + ) + fork_cold_reference = _execute( + tag=RUN_FORK_COLD_REFERENCE, + mgr=mgr_ref, + forest=forest_ref, + poses=poses_fork, + ) + + suffix = "." + args.video_format + videos = { + RUN_FORK_CACHE_REUSE_QUANT: save_video( + fork_cache_reuse_quant["frames"], + out / (RUN_FORK_CACHE_REUSE_QUANT + suffix), + args.fps, + ), + RUN_FORK_CACHE_REUSE_NONE: save_video( + fork_cache_reuse_none["frames"], + out / (RUN_FORK_CACHE_REUSE_NONE + suffix), + args.fps, + ), + RUN_FORK_COLD_REFERENCE: save_video( + fork_cold_reference["frames"], + out / (RUN_FORK_COLD_REFERENCE + suffix), + args.fps, + ), + } + if args.save_frames: + save_frames_by_chunk( + fork_cache_reuse_quant["frames_by_chunk"], + out / "frames" / RUN_FORK_CACHE_REUSE_QUANT, + ) + save_frames_by_chunk( + fork_cache_reuse_none["frames_by_chunk"], + out / "frames" / RUN_FORK_CACHE_REUSE_NONE, + ) + save_frames_by_chunk( + fork_cold_reference["frames_by_chunk"], + out / "frames" / RUN_FORK_COLD_REFERENCE, + ) + + # Quality comparisons answer three different questions: + # selected vs cold: end-to-end effect of quantized reuse; + # none vs cold: baseline effect of exact reuse without quantization; + # selected vs none: incremental effect attributable to KIVI quantization. + quant_vs_cold = aligned_chunk_metrics( + fork_cache_reuse_quant["frames_by_chunk"], + fork_cold_reference["frames_by_chunk"], + ) + none_vs_cold = aligned_chunk_metrics( + fork_cache_reuse_none["frames_by_chunk"], + fork_cold_reference["frames_by_chunk"], + ) + quant_vs_none = aligned_chunk_metrics( + fork_cache_reuse_quant["frames_by_chunk"], + fork_cache_reuse_none["frames_by_chunk"], + ) + metrics = { + "quant_reuse_vs_cold_reference": quant_vs_cold, + "unquantized_reuse_vs_cold_reference": none_vs_cold, + "quant_reuse_vs_unquantized_reuse": quant_vs_none, + } + + performance = { + "quant_reuse_vs_cold_reference": timing_comparison( + fork_cache_reuse_quant, + fork_cold_reference, + ), + "unquantized_reuse_vs_cold_reference": timing_comparison( + fork_cache_reuse_none, + fork_cold_reference, + ), + "quant_reuse_vs_unquantized_reuse": timing_comparison( + fork_cache_reuse_quant, + fork_cache_reuse_none, + ), + } + + checks = { + "trajectory_prefix_is_identical": bool( + np.array_equal( + poses_main[:fork_from_frame], + poses_fork[:fork_from_frame], + ) + ), + "trajectory_suffix_is_forked": bool( + not np.array_equal( + poses_main[fork_from_frame:], + poses_fork[fork_from_frame:], + ) + ), + "main_cache_seed_quant_is_cold": ( + main_cache_seed_quant["fast_forward_k"] == 0 + ), + "fork_cache_reuse_quant_prefix_hit": ( + fork_cache_reuse_quant["fast_forward_k"] == args.prefix_chunks + ), + "main_cache_seed_none_is_cold": ( + main_cache_seed_none["fast_forward_k"] == 0 + ), + "fork_cache_reuse_none_prefix_hit": ( + fork_cache_reuse_none["fast_forward_k"] == args.prefix_chunks + ), + "fork_cold_reference_is_cold": ( + fork_cold_reference["fast_forward_k"] == 0 + ), + "quant_vs_cold_has_common_chunks": quant_vs_cold["has_common_chunks"], + "quant_vs_cold_frame_counts_match": ( + quant_vs_cold["frame_count_match_per_chunk"] + ), + "none_vs_cold_has_common_chunks": none_vs_cold["has_common_chunks"], + "none_vs_cold_frame_counts_match": ( + none_vs_cold["frame_count_match_per_chunk"] + ), + "quant_vs_none_has_common_chunks": quant_vs_none["has_common_chunks"], + "quant_vs_none_frame_counts_match": ( + quant_vs_none["frame_count_match_per_chunk"] + ), + "benchmark_timings_are_unprofiled": all( + result["timing_valid_for_speedup"] + for result in ( + main_cache_seed_quant, + fork_cache_reuse_quant, + main_cache_seed_none, + fork_cache_reuse_none, + fork_cold_reference, + ) + ), + } + + # ------------------------------------------------------------------ + # Optional diagnostic passes + # ------------------------------------------------------------------ + # Every selected target gets a fresh manager. Reuse targets receive an + # unprofiled seed request first; only the requested target/scope is captured. + profile_kind = "none" + requested_profile_target = "none" + if args.profile == "cprofile": + profile_kind = "cprofile" + requested_profile_target = args.profile_target + elif args.torch_profile != "none": + profile_kind = "torch" + requested_profile_target = args.torch_profile + + profile_targets = ( + list(PROFILE_RUN_TARGETS) + if requested_profile_target == "all" + else ( + [] + if requested_profile_target == "none" + else [requested_profile_target] + ) + ) + profile_requests: dict[str, Any] = {} + + def _default_profile_chunk_index(target: str) -> int: + if args.profile_chunk_index >= 0: + return args.profile_chunk_index + if target in {RUN_FORK_CACHE_REUSE_QUANT, RUN_FORK_CACHE_REUSE_NONE}: + return args.prefix_chunks + return 0 + + for profile_target in profile_targets: + target_quant = ( + args.quant + if profile_target + in {RUN_MAIN_CACHE_SEED_QUANT, RUN_FORK_CACHE_REUSE_QUANT} + else "none" + ) + profile_forest, profile_mgr = _build_branch( + quant=target_quant, + branch=f"profile_{profile_kind}_{profile_target}", + ) + + # Seed only when profiling a reuse request. The seed remains outside + # every profiler scope, including full_request. + if profile_target == RUN_FORK_CACHE_REUSE_QUANT: + _execute( + tag=f"profile_setup_{RUN_MAIN_CACHE_SEED_QUANT}", + mgr=profile_mgr, + forest=profile_forest, + poses=poses_main, + ) + elif profile_target == RUN_FORK_CACHE_REUSE_NONE: + _execute( + tag=f"profile_setup_{RUN_MAIN_CACHE_SEED_NONE}", + mgr=profile_mgr, + forest=profile_forest, + poses=poses_main, + ) + + target_poses = ( + poses_main + if profile_target + in {RUN_MAIN_CACHE_SEED_QUANT, RUN_MAIN_CACHE_SEED_NONE} + else poses_fork + ) + controller = RequestProfileController( + kind=profile_kind, + tag=profile_target, + scope=args.profile_scope, + chunk_index=_default_profile_chunk_index(profile_target), + cprofile_dir=profile_dir if profile_kind == "cprofile" else None, + torch_profile_dir=( + torch_profile_dir if profile_kind == "torch" else None + ), + sort_by=args.profile_sort, + top_n=args.profile_top_n, + export_chrome_trace=args.torch_profile_chrome_trace, + torch_detail=args.torch_profile_detail, + ) + profiled_result, profile_artifacts = run_profiled_request( + controller, + lambda active_controller, target=profile_target, poses=target_poses: _execute( + tag=target, + mgr=profile_mgr, + forest=profile_forest, + poses=poses, + controller=active_controller, + ), + ) + profile_requests[profile_target] = { + "artifacts": profile_artifacts, + "request_summary": { + "fast_forward_k": profiled_result["fast_forward_k"], + "initial_chunk_index": profiled_result["initial_chunk_index"], + "generated_chunk_indices": profiled_result[ + "generated_chunk_indices" + ], + "timing_valid_for_speedup": profiled_result[ + "timing_valid_for_speedup" + ], + }, + } + + cprofile_manifest = { + "enabled": profile_kind == "cprofile", + "target": args.profile_target if args.profile == "cprofile" else "none", + "scope": args.profile_scope, + "output_dir": str(profile_dir) if profile_kind == "cprofile" else None, + "sort_by": args.profile_sort, + "top_n": max(int(args.profile_top_n), 1), + "requests": profile_requests if profile_kind == "cprofile" else {}, + "separate_from_benchmark": True, + } + torch_profile_manifest = { + "enabled": profile_kind == "torch", + "target": args.torch_profile, + "scope": args.profile_scope, + "output_dir": ( + str(torch_profile_dir) if profile_kind == "torch" else None + ), + "chrome_trace": bool(args.torch_profile_chrome_trace), + "detail": bool(args.torch_profile_detail), + "requests": profile_requests if profile_kind == "torch" else {}, + "separate_from_benchmark": True, + } + + results = { + RUN_MAIN_CACHE_SEED_QUANT: strip_frames(main_cache_seed_quant), + RUN_FORK_CACHE_REUSE_QUANT: strip_frames(fork_cache_reuse_quant), + RUN_MAIN_CACHE_SEED_NONE: strip_frames(main_cache_seed_none), + RUN_FORK_CACHE_REUSE_NONE: strip_frames(fork_cache_reuse_none), + RUN_FORK_COLD_REFERENCE: strip_frames(fork_cold_reference), + } + manifest = { + "selected_quant": args.quant, + "automatic_baseline_quant": "none", + "frame_num": args.frame_num, + "prefix_chunks": args.prefix_chunks, + "fork_from_frame": fork_from_frame, + "prompt": args.prompt, + "seed": args.seed, + "store": args.store, + "run_descriptions": RUN_DESCRIPTIONS, + "benchmark": { + "profiled": False, + "timing_valid_for_speedup": True, + "repeated_trials": False, + }, + "pipeline": { + "checkpoint_dir": ckpt, + "local_attn_size": cfg.local_attn_size, + "sink_size": cfg.sink_size, + "kv_chunk_size": KV_CHUNK_SIZE, + }, + "world_kv_quant": { + "selected_quant": args.quant, + "baseline_quant": "none", + "group_size": args.group_size, + "kv_layout": args.kv_layout, + "key_group_axis": key_group_axis, + "value_group_axis": value_group_axis, + "scale_dtype": args.scale_dtype, + "offset_dtype": args.offset_dtype, + }, + "environment": { + "python": sys.version, + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "device": ( + torch.cuda.get_device_name(0) + if torch.cuda.is_available() + else "cuda_unavailable" + ), + }, + "videos": videos, + "profile": cprofile_manifest, + "torch_profile": torch_profile_manifest, + "results": results, + "quality_metrics": metrics, + "performance": performance, + "checks": checks, + "all_checks_pass": all(checks.values()), + } + manifest_path = out / "manifest.json" + manifest_path.write_text( + json.dumps(manifest, indent=2, allow_nan=False), + encoding="utf-8", + ) + + print( + json.dumps( + { + "manifest": str(manifest_path), + "all_checks_pass": manifest["all_checks_pass"], + "selected_quant": args.quant, + "timing_s": { + tag: { + "create_runtime": result["create_runtime_s"], + "generation": result["generation_s"], + "flush": result["flush_s"], + "total_without_flush": result["total_without_flush_s"], + "total_with_flush": result["total_with_flush_s"], + } + for tag, result in results.items() + }, + "speedup": { + name: comparison["phases"]["total_without_flush_s"][ + "speedup" + ] + for name, comparison in performance.items() + }, + "profile": cprofile_manifest, + "torch_profile": torch_profile_manifest, + "checks": checks, + }, + indent=2, + allow_nan=False, + ), + flush=True, + ) + return 0 if manifest["all_checks_pass"] else 1 + + +if __name__ == "__main__": + sys.exit(main()) \ No newline at end of file diff --git a/cacheseek/quant/__init__.py b/cacheseek/quant/__init__.py new file mode 100644 index 0000000..eccc72f --- /dev/null +++ b/cacheseek/quant/__init__.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""KV-cache quantization codecs and payload schemas.""" + +from .factory import ( + build_kv_codec, + build_kv_codec_from_config, + quant_fingerprint_from_config, +) +from .kivi import KIVICodec +from .types import ( + KVLayerCodec, + KVQuantizedLayer, + QuantDType, + QuantScheme, + QuantTensor, + QuantTensorSpec, + TensorRole, + TensorSpec, +) + +__all__ = [ + "TensorRole", + "QuantScheme", + "QuantDType", + "TensorSpec", + "QuantTensorSpec", + "QuantTensor", + "KVQuantizedLayer", + "KVLayerCodec", + "KIVICodec", + "build_kv_codec", + "build_kv_codec_from_config", + "quant_fingerprint_from_config", +] diff --git a/cacheseek/quant/factory.py b/cacheseek/quant/factory.py new file mode 100644 index 0000000..2fdd72f --- /dev/null +++ b/cacheseek/quant/factory.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""Factory helpers for KV-cache quantization codecs.""" + +from __future__ import annotations + +from typing import Any + +from .types import KVLayerCodec, QuantScheme + + +def build_kv_codec( + *, + quant: str | QuantScheme = QuantScheme.NONE, + group_size: int = 64, + kv_layout: str = "H,T,D", + key_group_axis: str | int = "T", + value_group_axis: str | int = "D", + scale_dtype: str = "float32", + offset_dtype: str = "float32", +) -> KVLayerCodec | None: + """Build an optional per-layer KV codec from explicit parameters.""" + + if quant == QuantScheme.NONE.value: + return None + + if quant == QuantScheme.KIVI_INT4.value: + from .kivi import KIVICodec + + return KIVICodec( + bits=4, + group_size=group_size, + layout=kv_layout, + key_group_axis=key_group_axis, + value_group_axis=value_group_axis, + scale_dtype=scale_dtype, + offset_dtype=offset_dtype, + ) + + if quant == QuantScheme.KIVI_INT8.value: + from .kivi import KIVICodec + + return KIVICodec( + bits=8, + group_size=group_size, + layout=kv_layout, + key_group_axis=key_group_axis, + value_group_axis=value_group_axis, + scale_dtype=scale_dtype, + offset_dtype=offset_dtype, + ) + + raise ValueError(f"unsupported KV quantization scheme: {quant!r}") + + +def build_kv_codec_from_config(cfg: Any) -> KVLayerCodec | None: + """Build an optional per-layer KV codec from a config-like object. + + This helper expects attributes used by WorldKVConfig, but intentionally does + not import exact_prefix.config so quant/ can stay reusable. + """ + + return build_kv_codec( + quant=getattr(cfg, "quant", QuantScheme.NONE.value), + group_size=getattr(cfg, "group_size", 64), + kv_layout=getattr(cfg, "kv_layout", "H,T,D"), + key_group_axis=getattr(cfg, "key_group_axis", "T"), + value_group_axis=getattr(cfg, "value_group_axis", "D"), + scale_dtype=getattr(cfg, "scale_dtype", "float32"), + offset_dtype=getattr(cfg, "offset_dtype", "float32"), + ) + + +def quant_fingerprint_from_config(cfg: Any) -> dict[str, Any]: + """Return config fields that must isolate cache namespaces.""" + + return { + "quant": getattr(cfg, "quant", QuantScheme.NONE.value), + "group_size": getattr(cfg, "group_size", 64), + "kv_layout": getattr(cfg, "kv_layout", "H,T,D"), + "key_group_axis": getattr(cfg, "key_group_axis", "T"), + "value_group_axis": getattr(cfg, "value_group_axis", "D"), + "scale_dtype": getattr(cfg, "scale_dtype", "float32"), + "offset_dtype": getattr(cfg, "offset_dtype", "float32"), + } + + +__all__ = [ + "build_kv_codec", + "build_kv_codec_from_config", + "quant_fingerprint_from_config", +] diff --git a/cacheseek/quant/kernel.py b/cacheseek/quant/kernel.py new file mode 100644 index 0000000..a2351a9 --- /dev/null +++ b/cacheseek/quant/kernel.py @@ -0,0 +1,468 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""Optional kernels for KIVI-style KV-cache quantization. + +The public helpers preserve the storage contract established by +``cacheseek.quant.kivi``: + +* grouped affine quantization uses minimum offsets and unsigned q values; +* int4 packing is flat, low nibble to high nibble, eight values per int32; +* CPU tensors, CUDA tensors without Triton, and environments without Triton all + use PyTorch fallbacks. + +When Triton is available and tensors are CUDA resident, the grouped quant/dequant +and int4 pack/unpack paths use fused Triton kernels. Otherwise the same helpers +run with ordinary torch ops on the tensor's current device, which gives a CUDA +PyTorch fallback without copying full KV tensors back to CPU. +""" + +from __future__ import annotations + +from typing import Any + +try: # Triton is an optional acceleration dependency. + import triton + import triton.language as tl +except Exception: # pragma: no cover - depends on optional local install. + triton = None + tl = None + + +_TRITON_AVAILABLE = triton is not None and tl is not None + + +if _TRITON_AVAILABLE: + + @triton.jit + def _round_half_to_even(x): + floored = tl.floor(x) + frac = x - floored + floor_i = floored.to(tl.int32) + is_half = frac == 0.5 + floor_is_odd = (floor_i & 1) == 1 + round_up = (frac > 0.5) | (is_half & floor_is_odd) + return floored + round_up.to(tl.float32) + + @triton.jit + def _quantize_grouped_last_kernel( + x_ptr, + q_ptr, + scale_ptr, + offset_ptr, + GROUP_SIZE: tl.constexpr, + LEVELS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + group_id = tl.program_id(0) + offs = tl.arange(0, BLOCK_SIZE) + mask = offs < GROUP_SIZE + ptrs = x_ptr + group_id * GROUP_SIZE + offs + + x_min_load = tl.load(ptrs, mask=mask, other=float("inf")).to(tl.float32) + x_max_load = tl.load(ptrs, mask=mask, other=-float("inf")).to(tl.float32) + mn = tl.min(x_min_load, axis=0) + mx = tl.max(x_max_load, axis=0) + scale = (mx - mn) / LEVELS + scale_safe = tl.where(scale > 0.0, scale, 1.0) + + x = tl.load(ptrs, mask=mask, other=0.0).to(tl.float32) + q = _round_half_to_even((x - mn) / scale_safe) + q = tl.minimum(tl.maximum(q, 0.0), LEVELS).to(tl.uint8) + + tl.store(q_ptr + group_id * GROUP_SIZE + offs, q, mask=mask) + tl.store(scale_ptr + group_id, scale) + tl.store(offset_ptr + group_id, mn) + + @triton.jit + def _dequantize_grouped_last_kernel( + q_ptr, + scale_ptr, + offset_ptr, + out_ptr, + GROUP_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + group_id = tl.program_id(0) + offs = tl.arange(0, BLOCK_SIZE) + mask = offs < GROUP_SIZE + ptrs = group_id * GROUP_SIZE + offs + + q = tl.load(q_ptr + ptrs, mask=mask, other=0).to(tl.float32) + scale = tl.load(scale_ptr + group_id).to(tl.float32) + offset = tl.load(offset_ptr + group_id).to(tl.float32) + out = q * scale + offset + tl.store(out_ptr + ptrs, out, mask=mask) + + @triton.jit + def _pack_int4_flat_kernel( + q_ptr, + packed_ptr, + NUM_VALUES: tl.constexpr, + BLOCK_WORDS: tl.constexpr, + ): + word_offsets = tl.program_id(0) * BLOCK_WORDS + tl.arange(0, BLOCK_WORDS) + value_base = word_offsets * 8 + word_mask = value_base < NUM_VALUES + packed = tl.zeros((BLOCK_WORDS,), dtype=tl.uint32) + + for i in range(8): + value_offsets = value_base + i + q = tl.load(q_ptr + value_offsets, mask=value_offsets < NUM_VALUES, other=0).to( + tl.uint32 + ) + packed = packed | ((q & 0xF) << (4 * i)) + + tl.store(packed_ptr + word_offsets, packed.to(tl.int32), mask=word_mask) + + @triton.jit + def _unpack_int4_flat_kernel( + packed_ptr, + q_ptr, + NUM_VALUES: tl.constexpr, + BLOCK_VALUES: tl.constexpr, + ): + value_offsets = tl.program_id(0) * BLOCK_VALUES + tl.arange(0, BLOCK_VALUES) + mask = value_offsets < NUM_VALUES + word_offsets = value_offsets // 8 + shift = (value_offsets - word_offsets * 8) * 4 + packed = tl.load(packed_ptr + word_offsets, mask=mask, other=0).to(tl.uint32) + q = ((packed >> shift) & 0xF).to(tl.uint8) + tl.store(q_ptr + value_offsets, q, mask=mask) + + +def triton_available() -> bool: + """Return whether optional Triton kernels were imported successfully.""" + + return bool(_TRITON_AVAILABLE) + + +def can_use_triton(*tensors: Any) -> bool: + """Return whether all provided tensors can run through the Triton path.""" + + if not _TRITON_AVAILABLE or not tensors: + return False + devices = [] + for tensor in tensors: + if tensor is None or not getattr(tensor, "is_cuda", False): + return False + devices.append(getattr(tensor, "device", None)) + return all(device == devices[0] for device in devices) + + +def quantize_grouped( + tensor: Any, + *, + axis: int, + group_size: int, + bits: int, +) -> tuple[Any, Any, Any, tuple[int, ...]]: + """Quantize ``tensor`` in fixed-size groups along ``axis``. + + Returns ``(q, scale, offset, padded_shape)`` where ``q`` has the padded + original layout and dtype ``uint8``. ``scale`` and ``offset`` are stored in + grouped-stat layout: original axes with the grouped axis moved to the end and + replaced by ``num_groups``. + """ + + torch = _torch() + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"quantize_grouped expects torch.Tensor, got {type(tensor).__name__}") + if bits not in (4, 8): + raise ValueError(f"bits must be 4 or 8, got {bits}") + if group_size <= 0: + raise ValueError(f"group_size must be positive, got {group_size}") + + axis = _normalize_axis(axis, tensor.ndim) + x = _pad_to_group( + tensor.detach().to(torch.float32).contiguous(), + axis=axis, + group_size=group_size, + ) + padded_shape = tuple(int(s) for s in x.shape) + + if not can_use_triton(x): + q, scale, offset = _quantize_grouped_torch(x, axis=axis, group_size=group_size, bits=bits) + return q, scale, offset, padded_shape + + moved = x.movedim(axis, -1).contiguous() + groups = moved.shape[-1] // group_size + grouped_shape = (*moved.shape[:-1], groups, group_size) + grouped = moved.reshape(grouped_shape) + rows = grouped.numel() // group_size + + q_grouped = torch.empty_like(grouped, dtype=torch.uint8) + scale = torch.empty((*moved.shape[:-1], groups), dtype=torch.float32, device=x.device) + offset = torch.empty_like(scale) + + block_size = _next_power_of_2(group_size) + _quantize_grouped_last_kernel[(rows,)]( + grouped, + q_grouped, + scale, + offset, + GROUP_SIZE=group_size, + LEVELS=(1 << bits) - 1, + BLOCK_SIZE=block_size, + num_warps=_num_warps(block_size), + ) + + q = _restore_grouped(q_grouped, axis=axis) + return q.contiguous(), scale.contiguous(), offset.contiguous(), padded_shape + + +def dequantize_grouped( + qdata: Any, + scale: Any, + offset: Any, + *, + axis: int, + group_size: int, + original_shape: tuple[int, ...] | None = None, + dtype: Any | None = None, +) -> Any: + """Dequantize full ``uint8`` qdata from grouped KIVI metadata.""" + + torch = _torch() + if not isinstance(qdata, torch.Tensor): + raise TypeError(f"dequantize_grouped expects torch.Tensor, got {type(qdata).__name__}") + if offset is None: + raise ValueError("offset tensor is required") + if group_size <= 0: + raise ValueError(f"group_size must be positive, got {group_size}") + + axis = _normalize_axis(axis, qdata.ndim) + if dtype is None: + dtype = torch.float32 + + if not can_use_triton(qdata, scale, offset): + out = _dequantize_grouped_torch( + qdata, + scale, + offset, + axis=axis, + group_size=group_size, + ) + out = _slice_to_shape(out, original_shape) if original_shape is not None else out + return out.to(dtype=dtype).contiguous() + + q = qdata.contiguous() + moved = q.movedim(axis, -1).contiguous() + if moved.shape[-1] % group_size != 0: + raise ValueError( + f"grouped axis length {moved.shape[-1]} is not divisible by group_size={group_size}" + ) + + groups = moved.shape[-1] // group_size + expected_stats_shape = (*moved.shape[:-1], groups) + if tuple(scale.shape) != expected_stats_shape or tuple(offset.shape) != expected_stats_shape: + raise ValueError( + "scale/offset shape mismatch: " + f"expected {expected_stats_shape}, got scale={tuple(scale.shape)} " + f"offset={tuple(offset.shape)}" + ) + + grouped = moved.reshape(*moved.shape[:-1], groups, group_size) + out_grouped = torch.empty(grouped.shape, dtype=dtype, device=q.device) + block_size = _next_power_of_2(group_size) + rows = grouped.numel() // group_size + + _dequantize_grouped_last_kernel[(rows,)]( + grouped, + scale.contiguous(), + offset.contiguous(), + out_grouped, + GROUP_SIZE=group_size, + BLOCK_SIZE=block_size, + num_warps=_num_warps(block_size), + ) + + out = _restore_grouped(out_grouped, axis=axis) + out = _slice_to_shape(out, original_shape) if original_shape is not None else out + return out.contiguous() + + +def pack_int4_to_int32(qdata: Any) -> Any: + """Pack unsigned int4 values into int32 words, low nibble first.""" + + torch = _torch() + if not isinstance(qdata, torch.Tensor): + raise TypeError(f"pack_int4_to_int32 expects torch.Tensor, got {type(qdata).__name__}") + + flat = qdata.reshape(-1).contiguous() + num_values = int(flat.numel()) + num_words = (num_values + 7) // 8 + if num_words == 0: + return torch.empty(0, dtype=torch.int32, device=qdata.device) + + if can_use_triton(flat): + packed = torch.empty(num_words, dtype=torch.int32, device=flat.device) + block_words = 256 + _pack_int4_flat_kernel[(triton.cdiv(num_words, block_words),)]( + flat, + packed, + NUM_VALUES=num_values, + BLOCK_WORDS=block_words, + num_warps=4, + ) + return packed.contiguous() + + flat_i64 = flat.to(torch.int64) + pad = (-num_values) % 8 + if pad: + flat_i64 = torch.cat( + (flat_i64, torch.zeros(pad, dtype=flat_i64.dtype, device=flat_i64.device)) + ) + + nibbles = flat_i64.reshape(-1, 8) + packed = torch.zeros(nibbles.shape[0], dtype=torch.int64, device=flat_i64.device) + for i in range(8): + packed |= (nibbles[:, i] & 0xF) << (4 * i) + return packed.to(torch.int32).contiguous() + + +def unpack_int4_from_int32(packed: Any, num_values: int) -> Any: + """Unpack flat int32-packed int4 values into a ``uint8`` tensor.""" + + torch = _torch() + if not isinstance(packed, torch.Tensor): + raise TypeError(f"unpack_int4_from_int32 expects torch.Tensor, got {type(packed).__name__}") + num_values = int(num_values) + if num_values < 0: + raise ValueError(f"num_values must be non-negative, got {num_values}") + if num_values == 0: + return torch.empty(0, dtype=torch.uint8, device=packed.device) + + words = packed.reshape(-1).contiguous() + if can_use_triton(words): + out = torch.empty(num_values, dtype=torch.uint8, device=words.device) + block_values = 256 + _unpack_int4_flat_kernel[(triton.cdiv(num_values, block_values),)]( + words, + out, + NUM_VALUES=num_values, + BLOCK_VALUES=block_values, + num_warps=4, + ) + return out.contiguous() + + words_i64 = words.to(torch.int64) + parts = [] + for i in range(8): + parts.append(((words_i64 >> (4 * i)) & 0xF).to(torch.uint8)) + return torch.stack(parts, dim=1).reshape(-1)[:num_values].contiguous() + + +def _quantize_grouped_torch( + tensor: Any, + *, + axis: int, + group_size: int, + bits: int, +) -> tuple[Any, Any, Any]: + torch = _torch() + grouped = _reshape_grouped(tensor, axis=axis, group_size=group_size) + offset = grouped.amin(dim=-1) + maximum = grouped.amax(dim=-1) + levels = float((1 << bits) - 1) + scale = (maximum - offset) / levels + scale_safe = torch.where(scale > 0, scale, torch.ones_like(scale)) + q = torch.round((grouped - offset.unsqueeze(-1)) / scale_safe.unsqueeze(-1)) + q = q.clamp_(0, int(levels)).to(torch.uint8) + return _restore_grouped(q, axis=axis).contiguous(), scale.contiguous(), offset.contiguous() + + +def _dequantize_grouped_torch( + qdata: Any, + scale: Any, + offset: Any, + *, + axis: int, + group_size: int, +) -> Any: + torch = _torch() + q_grouped = _reshape_grouped(qdata.to(torch.float32), axis=axis, group_size=group_size) + decoded = q_grouped * scale.to(torch.float32).unsqueeze(-1) + offset.to( + torch.float32 + ).unsqueeze(-1) + return _restore_grouped(decoded, axis=axis) + + +def _pad_to_group(tensor: Any, *, axis: int, group_size: int) -> Any: + torch = _torch() + axis = _normalize_axis(axis, tensor.ndim) + size = int(tensor.shape[axis]) + if size <= 0: + raise ValueError("cannot quantize an empty tensor axis") + pad = (-size) % int(group_size) + if pad == 0: + return tensor + last = tensor.select(axis, size - 1).unsqueeze(axis) + shape = list(tensor.shape) + shape[axis] = pad + return torch.cat((tensor, last.expand(*shape)), dim=axis) + + +def _reshape_grouped(tensor: Any, *, axis: int, group_size: int) -> Any: + axis = _normalize_axis(axis, tensor.ndim) + moved = tensor.movedim(axis, -1).contiguous() + if moved.shape[-1] % group_size != 0: + raise ValueError( + f"grouped axis length {moved.shape[-1]} is not divisible by group_size={group_size}" + ) + groups = moved.shape[-1] // group_size + return moved.reshape(*moved.shape[:-1], groups, group_size) + + +def _restore_grouped(grouped: Any, *, axis: int) -> Any: + ndim = grouped.ndim - 1 + axis = _normalize_axis(axis, ndim) + flat = grouped.reshape(*grouped.shape[:-2], grouped.shape[-2] * grouped.shape[-1]) + return flat.movedim(-1, axis).contiguous() + + +def _slice_to_shape(tensor: Any, shape: tuple[int, ...] | None) -> Any: + if shape is None: + return tensor + return tensor[tuple(slice(0, int(size)) for size in shape)] + + +def _normalize_axis(axis: int, ndim: int) -> int: + if ndim <= 0: + raise ValueError(f"ndim must be positive, got {ndim}") + axis = int(axis) + if axis < 0: + axis += ndim + if axis < 0 or axis >= ndim: + raise ValueError(f"axis {axis} out of bounds for ndim={ndim}") + return axis + + +def _next_power_of_2(value: int) -> int: + value = int(value) + if value <= 1: + return 1 + return 1 << (value - 1).bit_length() + + +def _num_warps(block_size: int) -> int: + if block_size >= 2048: + return 8 + if block_size >= 512: + return 4 + return 2 + + +def _torch() -> Any: + import torch + + return torch + + +__all__ = [ + "can_use_triton", + "dequantize_grouped", + "pack_int4_to_int32", + "quantize_grouped", + "triton_available", + "unpack_int4_from_int32", +] diff --git a/cacheseek/quant/kivi.py b/cacheseek/quant/kivi.py new file mode 100644 index 0000000..673d464 --- /dev/null +++ b/cacheseek/quant/kivi.py @@ -0,0 +1,421 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""KIVI-style KV quantization with optional kernel acceleration. + +The CPU path remains the reference implementation and defines the payload +format. CUDA tensors use ``cacheseek.quant.kernel`` helpers, which choose +Triton when available and PyTorch CUDA fallback otherwise while preserving +the same flat storage layout. +""" + +from __future__ import annotations + +from typing import Any + +from . import kernel as _kernel +from .types import ( + KVQuantizedLayer, + QuantDType, + QuantScheme, + QuantTensor, + QuantTensorSpec, + TensorRole, + TensorSpec, +) + + +class KIVICodec: + """Naive KV codec following KIVI-style key/value grouping. + + For layout H,T,D: + - key_group_axis="T" implements grouped per-channel quantization. + - value_group_axis="D" implements grouped per-token quantization. + + This implementation is chunk/layer-level. It does not manage attention, + rolling windows, ring buffers, or runtime cache eviction. + """ + + def __init__( + self, + *, + bits: int = 4, + group_size: int = 64, + layout: str = "H,T,D", + key_group_axis: str | int = "T", + value_group_axis: str | int = "D", + scale_dtype: str = "float32", + offset_dtype: str = "float32", + ) -> None: + if bits not in (4, 8): + raise ValueError(f"KIVICodec supports bits=4 or bits=8, got {bits}") + if group_size <= 0: + raise ValueError(f"group_size must be positive, got {group_size}") + + self.bits = int(bits) + self.group_size = int(group_size) + self.layout = layout + self._layout_tokens = _parse_layout(layout) + self.key_group_axis = key_group_axis + self.value_group_axis = value_group_axis + self.scale_dtype = scale_dtype + self.offset_dtype = offset_dtype + + @property + def scheme(self) -> QuantScheme: + if self.bits == 4: + return QuantScheme.KIVI_INT4 + return QuantScheme.KIVI_INT8 + + def encode_layer(self, key: Any, value: Any) -> KVQuantizedLayer: + """Encode one KV layer into a CPU-resident storage payload. + + Quantization run on the input CUDA device, but the compressed + qdata/scale/offset tensors are moved to CPU before they are returned. + """ + return KVQuantizedLayer( + key=self._encode_tensor(key, TensorRole.KEY), + value=self._encode_tensor(value, TensorRole.VALUE), + ) + + def decode_layer( + self, + payload: KVQuantizedLayer, + *, + device: Any | None = None, + ) -> tuple[Any, Any]: + """Decode one KV layer onto ``device``. + + When ``device`` is omitted, decoding preserves the payload's current + device. Because :meth:`encode_layer` returns CPU-resident payloads, + the default remains backward-compatible CPU decoding. Runtime callers + should pass the destination KV-cache device explicitly, for example + ``device=window.key_cache.device``. + """ + if payload.key.quant.scheme != self.scheme or payload.value.quant.scheme != self.scheme: + raise ValueError( + "quantized layer scheme does not match codec: " + f"key={payload.key.quant.scheme.value} " + f"value={payload.value.quant.scheme.value} " + f"codec={self.scheme.value}" + ) + return ( + self._decode_tensor(payload.key, device=device), + self._decode_tensor(payload.value, device=device), + ) + + def _encode_tensor(self, tensor: Any, role: TensorRole) -> QuantTensor: + torch = _torch() + + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"KIVICodec expects torch.Tensor, got {type(tensor).__name__}") + + if tensor.ndim != len(self._layout_tokens): + raise ValueError( + f"tensor rank {tensor.ndim} does not match layout {self.layout!r} " + f"with {len(self._layout_tokens)} axes" + ) + + source = tensor.detach() + original = TensorSpec( + shape=tuple(int(s) for s in source.shape), + dtype=_dtype_name(source.dtype), + layout=self.layout, + ) + + group_axis = self._group_axis(role, source.ndim) + if getattr(source, "is_cuda", False): + q, scale, offset, padded_shape = _kernel.quantize_grouped( + source.contiguous(), + axis=group_axis, + group_size=self.group_size, + bits=self.bits, + ) + else: + x = source.to(device="cpu").contiguous() + x_padded = _pad_to_group( + x.to(torch.float32), + axis=group_axis, + group_size=self.group_size, + ) + q, scale, offset = self._quantize_grouped(x_padded, axis=group_axis) + padded_shape = tuple(int(s) for s in x_padded.shape) + + if self.bits == 4: + qdata = _pack_int4_to_int32(q) + storage_dtype = QuantDType.INT32_PACKED + else: + qdata = q.to(torch.uint8).contiguous() + storage_dtype = QuantDType.UINT8 + + quant = QuantTensorSpec( + role=role, + scheme=self.scheme, + bits=self.bits, + storage_dtype=storage_dtype, + group_size=self.group_size, + group_axis=group_axis, + padded_shape=tuple(int(s) for s in padded_shape), + pack_order="low_to_high", + scale_dtype=self.scale_dtype, + offset_dtype=self.offset_dtype, + offset_kind="minimum", + symmetric=False, + ) + + return QuantTensor( + tensor=original, + quant=quant, + qdata=qdata.to(device="cpu").contiguous(), + scale=scale.to(device="cpu", dtype=_torch_dtype(self.scale_dtype)).contiguous(), + offset=offset.to(device="cpu", dtype=_torch_dtype(self.offset_dtype)).contiguous(), + ) + + def _decode_tensor( + self, + payload: QuantTensor, + *, + device: Any | None = None + ) -> Any: + torch = _torch() + quant = payload.quant + + if quant.offset_kind != "minimum": + raise ValueError(f"KIVICodec only supports minimum offsets, got {quant.offset_kind!r}") + if payload.offset is None: + raise ValueError("quantized payload is missing offset tensor") + if quant.bits not in (4, 8): + raise ValueError(f"unsupported bit width: {quant.bits}") + + padded_shape = quant.padded_shape or payload.tensor.shape + if len(padded_shape) != len(payload.tensor.shape): + raise ValueError( + "padded_shape rank does not match original tensor rank: " + f"padded={padded_shape}, original={payload.tensor.shape}" + ) + + group_axis = _normalize_axis(quant.group_axis, len(padded_shape)) + + for axis, (original_dim, padded_dim) in enumerate(zip(payload.tensor.shape, padded_shape)): + original_dim = int(original_dim) + padded_dim = int(padded_dim) + if axis == group_axis: + if padded_dim < original_dim: + raise ValueError( + "padded grouped axis is shorter than original axis: " + f"axis={axis}, padded={padded_dim}, original={original_dim}" + ) + elif padded_dim != original_dim: + raise ValueError( + "only grouped axis may be padded: " + f"axis={axis}, padded={padded_dim}, original={original_dim}" + ) + + num_values = _numel(padded_shape) + target_device = payload.qdata.device if device is None else torch.device(device) + + + # Move the compact representation before dequantization. For GPU + # materialization this transfers int4/int8 data plus statistics rather + # than a full-precision KV tensor. + qdata = payload.qdata.to(device=target_device).contiguous() + scale_tensor = payload.scale.to(device=target_device).contiguous() + offset_tensor = payload.offset.to(device=target_device).contiguous() + + if quant.bits == 4: + if quant.storage_dtype is not QuantDType.INT32_PACKED: + raise ValueError( + f"int4 KIVI payload must use INT32_PACKED storage, got {quant.storage_dtype}" + ) + if quant.pack_order != "low_to_high": + raise ValueError(f"unsupported int4 pack_order={quant.pack_order!r}") + q = _unpack_int4_from_int32(qdata, num_values).reshape(padded_shape) + else: + if quant.storage_dtype is not QuantDType.UINT8: + raise ValueError( + f"int8 KIVI payload must use UINT8 storage, got {quant.storage_dtype}" + ) + q = qdata.reshape(padded_shape) + + out_dtype = _torch_dtype(payload.tensor.dtype) + if _kernel.can_use_triton(q, scale_tensor, offset_tensor): + return _kernel.dequantize_grouped( + q, + scale_tensor, + offset_tensor, + axis=group_axis, + group_size=quant.group_size, + original_shape=payload.tensor.shape, + dtype=out_dtype, + ) + + q_grouped = _reshape_grouped( + q.to(torch.float32), + axis=group_axis, + group_size=quant.group_size, + ) + + # scale/offset were stored in grouped layout, e.g.: + # original [H,T,D], group_axis=T -> stats [H,D,num_groups] + # original [H,T,D], group_axis=D -> stats [H,T,num_groups] + # Therefore they only need unsqueeze(-1), not movedim(group_axis, -1). + scale = _reshape_stats_for_group(scale_tensor) + offset = _reshape_stats_for_group(offset_tensor) + + decoded_grouped = q_grouped * scale.to(torch.float32) + offset.to(torch.float32) + decoded_padded = _restore_grouped(decoded_grouped, axis=group_axis) + decoded = _slice_to_shape(decoded_padded, payload.tensor.shape) + + return decoded.to(dtype=out_dtype).contiguous() + + def _quantize_grouped(self, tensor: Any, *, axis: int) -> tuple[Any, Any, Any]: + """Naive CPU implementation of KIVI quantization, called when encoding a CPU tensor.""" + torch = _torch() + + grouped = _reshape_grouped(tensor, axis=axis, group_size=self.group_size) + offset = grouped.amin(dim=-1) + maximum = grouped.amax(dim=-1) + + levels = float((1 << self.bits) - 1) + scale = (maximum - offset) / levels + + # Avoid division by zero for constant groups. q becomes all zeros and + # dequantizes back to offset, which equals the original constant value. + scale_safe = torch.where(scale > 0, scale, torch.ones_like(scale)) + + q = torch.round((grouped - offset.unsqueeze(-1)) / scale_safe.unsqueeze(-1)) + q = q.clamp_(0, int(levels)).to(torch.uint8) + q = _restore_grouped(q, axis=axis) + + return q.contiguous(), scale.contiguous(), offset.contiguous() + + def _group_axis(self, role: TensorRole, ndim: int) -> int: + axis = self.key_group_axis if role is TensorRole.KEY else self.value_group_axis + + if isinstance(axis, int): + return _normalize_axis(axis, ndim) + + token = axis.strip().upper() + try: + return self._layout_tokens.index(token) + except ValueError as exc: + raise ValueError(f"axis {axis!r} not present in layout {self.layout!r}") from exc + + +def _parse_layout(layout: str) -> tuple[str, ...]: + tokens = tuple(part.strip().upper() for part in layout.split(",") if part.strip()) + + if not tokens: + raise ValueError("layout must contain at least one axis name") + if len(set(tokens)) != len(tokens): + raise ValueError(f"layout axes must be unique, got {layout!r}") + + return tokens + + +def _normalize_axis(axis: int, ndim: int) -> int: + if ndim <= 0: + raise ValueError(f"ndim must be positive, got {ndim}") + + axis = int(axis) + if axis < 0: + axis += ndim + + if axis < 0 or axis >= ndim: + raise ValueError(f"axis {axis} out of bounds for ndim={ndim}") + + return axis + + +def _pad_to_group(tensor: Any, *, axis: int, group_size: int) -> Any: + torch = _torch() + + axis = _normalize_axis(axis, tensor.ndim) + size = int(tensor.shape[axis]) + + if size <= 0: + raise ValueError("cannot quantize an empty tensor axis") + + pad = (-size) % int(group_size) + if pad == 0: + return tensor + + last = tensor.select(axis, size - 1).unsqueeze(axis) + shape = list(tensor.shape) + shape[axis] = pad + padding = last.expand(*shape) + + return torch.cat((tensor, padding), dim=axis) + + +def _reshape_grouped(tensor: Any, *, axis: int, group_size: int) -> Any: + axis = _normalize_axis(axis, tensor.ndim) + + moved = tensor.movedim(axis, -1).contiguous() + + if moved.shape[-1] % group_size != 0: + raise ValueError( + f"grouped axis length {moved.shape[-1]} is not divisible by " + f"group_size={group_size}" + ) + + groups = moved.shape[-1] // group_size + return moved.reshape(*moved.shape[:-1], groups, group_size) + + +def _reshape_stats_for_group(stats: Any) -> Any: + return stats.unsqueeze(-1).contiguous() + + +def _restore_grouped(grouped: Any, *, axis: int) -> Any: + ndim = grouped.ndim - 1 + axis = _normalize_axis(axis, ndim) + + flat = grouped.reshape(*grouped.shape[:-2], grouped.shape[-2] * grouped.shape[-1]) + return flat.movedim(-1, axis).contiguous() + + +def _slice_to_shape(tensor: Any, shape: tuple[int, ...]) -> Any: + slices = tuple(slice(0, int(size)) for size in shape) + return tensor[slices] + + +def _pack_int4_to_int32(qdata: Any) -> Any: + return _kernel.pack_int4_to_int32(qdata) + + +def _unpack_int4_from_int32(packed: Any, num_values: int) -> Any: + return _kernel.unpack_int4_from_int32(packed, num_values) + + +def _numel(shape: tuple[int, ...]) -> int: + total = 1 + for dim in shape: + total *= int(dim) + return total + + +def _dtype_name(dtype: Any) -> str: + return str(dtype).replace("torch.", "") + + +def _torch_dtype(name: str) -> Any: + torch = _torch() + dtype_name = str(name).replace("torch.", "") + + try: + return getattr(torch, dtype_name) + except AttributeError as exc: + raise ValueError(f"unsupported torch dtype name {name!r}") from exc + + +def _torch() -> Any: + import torch + + return torch + + +__all__ = [ + "KIVICodec", + "_pack_int4_to_int32", + "_unpack_int4_from_int32", +] \ No newline at end of file diff --git a/cacheseek/quant/types.py b/cacheseek/quant/types.py new file mode 100644 index 0000000..b4accd2 --- /dev/null +++ b/cacheseek/quant/types.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the CacheSeek project +"""Basic data types for KV-cache quantization. + +This module is intentionally torch-free. Codec implementations may use torch, +but store/reuse layers should be able to inspect payload metadata without +importing or initializing tensor libraries. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Any, Literal, Protocol, runtime_checkable + + +class TensorRole(str, Enum): + """Which half of a per-layer attention KV pair a tensor represents.""" + + KEY = "k" + VALUE = "v" + + +class QuantScheme(str, Enum): + """Supported KV-cache quantization schemes.""" + + NONE = "none" + KIVI_INT4 = "kivi_int4" + KIVI_INT8 = "kivi_int8" + + +class QuantDType(str, Enum): + """Storage dtype for quantized payload tensors.""" + + UINT8 = "uint8" + INT8 = "int8" + INT32_PACKED = "int32_packed" + + +@dataclass(frozen=True, slots=True) +class TensorSpec: + """Original tensor view needed to reconstruct a decoded KV tensor. + + Attributes: + shape: Logical tensor shape before quantization. + dtype: Original tensor dtype, stored as a portable string such as + "bfloat16", "float16", or "float32". + layout: Logical layout string, for example "H,T,D". + """ + + shape: tuple[int, ...] + dtype: str + layout: str # e.g. "H,T,D" for (num_heads, num_tokens, head_dim) + + +@dataclass(frozen=True, slots=True) +class QuantTensorSpec: + """Quantization metadata for one tensor component. + + group_axis identifies the original tensor axis split into fixed-size + quantization groups. For layout "H,T,D", group_axis=1 means grouping + along token axis T. + """ + + role: TensorRole + scheme: QuantScheme + bits: int + storage_dtype: QuantDType + group_size: int + group_axis: int + + padded_shape: tuple[int, ...] | None = None + pack_order: str = "low_to_high" + + scale_dtype: str = "float32" + offset_dtype: str = "float32" + offset_kind: Literal["minimum", "zero_point"] = "minimum" + + symmetric: bool = False + + +@dataclass(frozen=True, slots=True) +class QuantTensor: + """Encoded representation of one key or value tensor. + + Tensor-bearing fields are typed as Any so this type can represent torch + tensors, numpy arrays, or backend-native tensor handles. + """ + + tensor: TensorSpec + quant: QuantTensorSpec + qdata: Any + scale: Any + offset: Any | None = None + + @property + def role(self) -> TensorRole: + """Return the tensor role carried by the quantization spec.""" + + return self.quant.role + + +@dataclass(frozen=True, slots=True) +class KVQuantizedLayer: + """Quantized payload for one attention layer's KV pair.""" + + key: QuantTensor + value: QuantTensor + + def as_pair(self) -> tuple[QuantTensor, QuantTensor]: + """Return (key, value) in the same order as runtime KV payloads.""" + + return self.key, self.value + + +@runtime_checkable +class KVLayerCodec(Protocol): + """Codec boundary used by reuse managers and KV stores.""" + + @property + def scheme(self) -> QuantScheme: + """Quantization scheme implemented by this codec.""" + + ... + + def encode_layer(self, key: Any, value: Any) -> KVQuantizedLayer: + """Encode one attention layer's key/value tensors.""" + + ... + + def decode_layer(self, payload: KVQuantizedLayer) -> tuple[Any, Any]: + """Decode one attention layer payload back to ordinary tensors.""" + + ... + + +__all__ = [ + "TensorRole", + "QuantScheme", + "QuantDType", + "TensorSpec", + "QuantTensorSpec", + "QuantTensor", + "KVQuantizedLayer", + "KVLayerCodec", +] \ No newline at end of file diff --git a/cacheseek/reuse/exact_prefix/config.py b/cacheseek/reuse/exact_prefix/config.py index a477826..4a96c99 100644 --- a/cacheseek/reuse/exact_prefix/config.py +++ b/cacheseek/reuse/exact_prefix/config.py @@ -12,6 +12,8 @@ from dataclasses import dataclass +from cacheseek.quant import quant_fingerprint_from_config + from .types import Tier # Quantization scheme. Default NONE (bf16 lossless): the impact of KV quantization on @@ -19,8 +21,8 @@ # is an opt-in to enable only after validation passes, not a fixed design choice. QUANT_BYTES_PER_ELEM = { "none": 2.0, # bf16/fp16, lossless original values (default) - "int8": 1.0, - "int4": 0.5, # enable only after a quality A/B passes + "kivi_int8": 1.0, + "kivi_int4": 0.5, } @@ -113,8 +115,18 @@ class WorldKVConfig: break_even_k: int # see above; below it, don't fast-forward (falls back to normal generation, harmless) quant: str = "none" # default bf16 lossless; enable quantization (int8/int4) only after a quality A/B passes group_size: int = 64 + kv_layout: str = "H,T,D" + key_group_axis: str | int = "T" + value_group_axis: str | int = "D" + scale_dtype: str = "float32" + offset_dtype: str = "float32" commit_tier: Tier = Tier.FLUXON_DRAM + def quant_fingerprint(self) -> dict[str, object]: + """Return quantization fields that must participate in namespace hashing.""" + + return quant_fingerprint_from_config(self) + @classmethod def from_geometry( cls, @@ -127,6 +139,11 @@ def from_geometry( fixed_overhead_s: float = 0.0, quant: str = "none", # default lossless; quantization not assumed group_size: int = 64, + kv_layout: str = "H,T,D", + key_group_axis: str | int = "T", + value_group_axis: str | int = "D", + scale_dtype: str = "float32", + offset_dtype: str = "float32", overlap_factor: float = 1.0, commit_tier: Tier = Tier.FLUXON_DRAM, ) -> WorldKVConfig: @@ -145,6 +162,11 @@ def from_geometry( fixed_overhead_s: Per-fast-forward fixed overhead (lookup/setup/first-block fetch). quant: KV quantization scheme; "none" is bf16 lossless (default). group_size: Quantization group size (ignored when quant="none"). + kv_layout: Logical key/value tensor layout used by the codec. + key_group_axis: Grouped axis for keys. + value_group_axis: Grouped axis for values. + scale_dtype: Storage dtype for quantization scale tensors. + offset_dtype: Storage dtype for quantization offset tensors. overlap_factor: Fraction of fetch not overlappable with compute. commit_tier: Storage tier new blobs are committed to. @@ -167,5 +189,10 @@ def from_geometry( break_even_k=k, quant=quant, group_size=group_size, + kv_layout=kv_layout, + key_group_axis=key_group_axis, + value_group_axis=value_group_axis, + scale_dtype=scale_dtype, + offset_dtype=offset_dtype, commit_tier=commit_tier, ) diff --git a/cacheseek/reuse/exact_prefix/manager.py b/cacheseek/reuse/exact_prefix/manager.py index f15ab50..8afab53 100644 --- a/cacheseek/reuse/exact_prefix/manager.py +++ b/cacheseek/reuse/exact_prefix/manager.py @@ -19,6 +19,8 @@ from dataclasses import dataclass from typing import Any, Protocol +from cacheseek.quant import KVLayerCodec, KVQuantizedLayer, build_kv_codec_from_config + from .config import WorldKVConfig from .trie import Namespace, NamespaceForest, PrefixMatch from .types import ActionKey, BlobHandle, NodeKey, RootHash, Skeleton, TrieNode @@ -66,6 +68,11 @@ class RollingWindow(Protocol): blobs are given as ``(chunk_depth, layer_payload)`` (oldest→newest); the adapter assembles them into the engine's own physical layout (full-length / rolling ring). """ + + @property + def device(self) -> Any | None: + """Device on which restored runtime KV tensors must reside.""" + ... def seed_layer(self, layer: int, blobs: Sequence[tuple[int, Any]], depth: int) -> None: """Seed one attention layer's KV pool with the sink + window historical KV. @@ -114,6 +121,7 @@ def __init__(self, forest: NamespaceForest, store: KVTierStore, cfg: WorldKVConf self.forest = forest self.store = store self.cfg = cfg + self._kv_codec: KVLayerCodec | None = build_kv_codec_from_config(cfg) self._now = clock or (lambda: 0.0) # injectable clock for testing/determinism # ---------------------------------------------------------------- Compute-saving entry @@ -148,18 +156,40 @@ def materialize(self, node: TrieNode, window: RollingWindow) -> bool: return False for n in path: n.ref_count += 1 # materialize in flight; eviction must not reclaim + try: - n_layers = path[-1].blob.n_layers # type: ignore[union-attr] + n_layers = path[-1].blob.n_layers # type: ignore[union-attr] + # Production adapters should expose ``window.device``. ``getattr`` + # keeps existing CPU-only test doubles backward compatible: when + # absent, the codec decodes on the stored payload's current device. + target_device = ( + getattr(window, "device", None) + if self._kv_codec is not None + else None + ) + for layer in range(n_layers): - blobs = [(n.depth, self.store.get_layer(n.blob, layer)) for n in path] + blobs: list[tuple[int, Any]] = [] + + for n in path: + stored_payload = self.store.get_layer(n.blob, layer) + runtime_payload = self._decode_layer_payload( + stored_payload, + device=target_device, + ) + blobs.append((n.depth, runtime_payload)) + window.seed_layer(layer, blobs, depth=node.depth) + window.set_resume_depth(node.depth) + now = self._now() for n in path: n.last_access = now finally: for n in path: n.ref_count -= 1 + return True # ---------------------------------------------------------------- Write path @@ -173,7 +203,6 @@ def ingest( kv_payload: Sequence[Any], latent: Any, *, - nbytes: int = 0, n_layers: int | None = None, ) -> TrieNode: """Write on chunk finalize (write-through; the two pools are decoupled: roll-off is ingest, not evict). @@ -190,11 +219,46 @@ def ingest( self.store.put_skeleton(loc + ":lat", latent) node.skeleton = Skeleton(latent_locator=loc + ":lat") # blob (heavy; store's write callback sets ready ⇒ only then published to readers) + + def _to_cpu(t: Any) -> Any: + """Copy an unquantized runtime tensor into CPU-owned storage.""" + if not hasattr(t, "detach"): + return t + return t.detach().to("cpu").contiguous().clone() + + def _prepare_for_encode(t: Any) -> Any: + """Detach a tensor while preserving its source device. + + CUDA inputs therefore use the accelerated codec path on their + current GPU. CPU inputs use the codec's CPU reference path. The + manager must not move data to the process-wide default CUDA device, + because that device may not match the runtime in multi-GPU setups. + """ + if not hasattr(t, "detach"): + return t + return t.detach().contiguous() + + if self._kv_codec is None: + kv_payload = [ + tuple(_to_cpu(t) for t in payload) + if isinstance(payload, (tuple, list)) + else _to_cpu(payload) + for payload in kv_payload + ] + else: + kv_payload = [ + tuple(_prepare_for_encode(t) for t in payload) + if isinstance(payload, (tuple, list)) + else _prepare_for_encode(payload) + for payload in kv_payload + ] + + stored_payload = self._encode_kv_payload(kv_payload) handle = BlobHandle( tier=self.cfg.commit_tier, locator=loc + ":kv", - nbytes=nbytes, - n_layers=len(kv_payload) if n_layers is None else n_layers, + nbytes=self._stored_payload_nbytes(stored_payload), + n_layers=len(stored_payload) if n_layers is None else n_layers, ready=False, ) node.blob = handle @@ -202,7 +266,7 @@ def ingest( def _publish() -> None: handle.ready = True - self.store.put_async(handle.locator, kv_payload, tier=self.cfg.commit_tier, on_ready=_publish) + self.store.put_async(handle.locator, stored_payload, tier=self.cfg.commit_tier, on_ready=_publish) return node # ---------------------------------------------------------------- Eviction @@ -237,3 +301,75 @@ def _window_path(self, node: TrieNode) -> list[TrieNode]: sinks.append(n) n = n.parent return list(reversed(sinks)) + list(reversed(recent)) + + def _encode_kv_payload(self, kv_payload: Sequence[Any]) -> list[Any]: + """Encode per-layer KV payloads when quantization is enabled.""" + + if self._kv_codec is None: + return list(kv_payload) + + encoded: list[Any] = [] + for layer, payload in enumerate(kv_payload): + if not isinstance(payload, (tuple, list)) or len(payload) != 2: + raise TypeError( + "quantized WorldKV ingest expects each layer payload to be a " + f"(key, value) pair; layer {layer} got {type(payload).__name__}" + ) + encoded.append(self._kv_codec.encode_layer(payload[0], payload[1])) + return encoded + + def _decode_layer_payload( + self, + payload: Any, + *, + device: Any | None = None, + ) -> Any: + """Decode one stored layer payload back to runtime KV form. + + Args: + payload: One layer payload returned by the tiered store. + device: Destination device for restored KV tensors. For quantized + payloads, the codec moves the compact representation to this + device before dequantization. Ignored when quantization is + disabled. + """ + + if self._kv_codec is None: + return payload + if not isinstance(payload, KVQuantizedLayer): + raise TypeError( + "quantized WorldKV materialize expected KVQuantizedLayer from the store, " + f"got {type(payload).__name__}" + ) + return self._kv_codec.decode_layer(payload, device=device) + + def _stored_payload_nbytes(self, payload: Sequence[Any]) -> int: + """Best-effort byte size for accounting after optional quantization.""" + + total = 0 + for layer_payload in payload: + total += self._object_nbytes(layer_payload) + return total + + def _object_nbytes(self, value: Any) -> int: + if isinstance(value, KVQuantizedLayer): + return ( + self._object_nbytes(value.key.qdata) + + self._object_nbytes(value.key.scale) + + self._object_nbytes(value.key.offset) + + self._object_nbytes(value.value.qdata) + + self._object_nbytes(value.value.scale) + + self._object_nbytes(value.value.offset) + ) + if isinstance(value, (tuple, list)): + return sum(self._object_nbytes(v) for v in value) + if value is None: + return 0 + nbytes = getattr(value, "nbytes", None) + if nbytes is not None: + return int(nbytes) + element_size = getattr(value, "element_size", None) + numel = getattr(value, "numel", None) + if callable(element_size) and callable(numel): + return int(element_size() * numel()) + return 0 diff --git a/cacheseek/reuse/exact_prefix/strategy.py b/cacheseek/reuse/exact_prefix/strategy.py index d743244..fd0681d 100644 --- a/cacheseek/reuse/exact_prefix/strategy.py +++ b/cacheseek/reuse/exact_prefix/strategy.py @@ -64,7 +64,7 @@ async def save(self, query: CacheQuery, outputs: Any = None, ctx: Any = None) -> Unlike request-level save, exact_prefix writes per chunk: the chunk's KV payload, latent, and mount point are passed via ``ctx`` (keys ns, parent, - action, node_key, depth, payload, latent, optional nbytes); ``outputs`` is + action, node_key, depth, payload, latent); ``outputs`` is unused. The newly created trie node is written back to ``ctx["node"]`` so the caller can advance its cursor. @@ -74,6 +74,6 @@ async def save(self, query: CacheQuery, outputs: Any = None, ctx: Any = None) -> assert isinstance(ctx, dict) and "ns" in ctx, "exact_prefix save requires chunk ctx" node = self.mgr.ingest( ctx["ns"], ctx["parent"], ctx["action"], ctx["node_key"], ctx["depth"], - ctx["payload"], ctx["latent"], nbytes=ctx.get("nbytes", 0), + ctx["payload"], ctx["latent"], ) ctx["node"] = node diff --git a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py index f66553f..77bfcf6 100644 --- a/cacheseek/reuse/exact_prefix/telefuser_lingbot.py +++ b/cacheseek/reuse/exact_prefix/telefuser_lingbot.py @@ -40,6 +40,7 @@ import numpy as np import torch +from cacheseek.quant import quant_fingerprint_from_config from cacheseek.service.query import CacheQuery from .config import WorldKVConfig @@ -50,28 +51,48 @@ from .types import ActionKey -def make_full_kv_config(*, break_even_k: int = 1) -> WorldKVConfig: - """Full-length KV (local_attn_size=-1) => materialization needs the entire - prefix.""" - return WorldKVConfig( - window_chunks=1_000_000, sink_chunks=1, break_even_k=break_even_k - ) - - -def make_rolling_config( +def make_world_kv_config( *, local_attn_size: int = 7, # latent frames, incl. sink (matches TeleFuser pipeline config) sink_size: int = 3, # latent frames chunk_size: int = 3, # latent frames per chunk break_even_k: int = 1, + quant: str = "none", + group_size: int = 64, + kv_layout: str = "B,T,H,D", + key_group_axis: str | int = "T", + value_group_axis: str | int = "D", + scale_dtype: str = "float32", + offset_dtype: str = "float32", ) -> WorldKVConfig: - """Rolling window: materialization needs only the sink chunk plus the chunks - covering the most recent (L-S) frames -- O(W) rather than O(K).""" - recent_frames = max(local_attn_size - sink_size, 1) + """Build WorldKVConfig from TeleFuser LingBot KV geometry. + + Uses local_attn_size=-1 for full-length KV. Positive values for + rolling KV, where materialization needs only the sink chunk plus the chunks covering the most recent (L-S) frames -- O(W) rather than O(K). + """ + if local_attn_size == -1: + window_chunks = 1_000_000 + sink_chunks = 1 + elif local_attn_size > 0: + recent_frames = max(local_attn_size - sink_size, 1) + + window_chunks=-(-recent_frames // chunk_size) # ceil + sink_chunks=-(-sink_size // chunk_size) + else: + raise ValueError(f"local_attn_size must be -1 for full KV and positive " + f"for rolling, got {local_attn_size}") + return WorldKVConfig( - window_chunks=-(-recent_frames // chunk_size), # ceil - sink_chunks=-(-sink_size // chunk_size), + window_chunks=window_chunks, + sink_chunks=sink_chunks, break_even_k=break_even_k, + quant=quant, + group_size=group_size, + kv_layout=kv_layout, + key_group_axis=key_group_axis, + value_group_axis=value_group_axis, + scale_dtype=scale_dtype, + offset_dtype=offset_dtype, ) @@ -88,7 +109,12 @@ def make_rolling_config( ) -def session_root_hash(session_config: Any, *, model_fingerprint: bytes) -> bytes: +def session_root_hash( + session_config: Any, + *, + model_fingerprint: bytes, + quant_config: Any | None = None, +) -> bytes: """Compute the namespace root_hash for a TeleFuser session. Combines an image fingerprint (mode, size, raw pixel bytes), a normalized @@ -100,6 +126,8 @@ def session_root_hash(session_config: Any, *, model_fingerprint: bytes) -> bytes session_config: TeleFuser session config (duck-typed; needs image, prompt, and the SESSION_KEY_FIELDS attributes). model_fingerprint: Weights/version fingerprint folded into config_blob_hash. + quant_config: Optional config-like object whose quantization fields are + folded into config_blob_hash to isolate incompatible KV payloads. Returns: The root_hash identifying this world. @@ -113,6 +141,8 @@ def session_root_hash(session_config: Any, *, model_fingerprint: bytes) -> bytes ) prompt_fp = sha256(b"prompt", session_config.prompt.strip().encode("utf-8")) blob = {f: getattr(session_config, f, None) for f in SESSION_KEY_FIELDS} + if quant_config is not None: + blob["kv_quant"] = quant_fingerprint_from_config(quant_config) cfg_hash = config_blob_hash(blob, weights_fingerprint=model_fingerprint) return root_hash(image_fp=image_fp, prompt_fp=prompt_fp, config_blob_hash=cfg_hash) @@ -173,6 +203,35 @@ def __init__(self, runtime: Any) -> None: self._rt = runtime self._local_end_tokens = 0 + @property + def device(self) -> torch.device: + """Return the device on which restored runtime KV must reside. + + All self-attention K/V cache tensors are expected to live on the same + device. The manager passes this device to the codec so compressed + qdata/scale/offset tensors are transferred before GPU dequantization. + """ + caches = self._rt.self_kv_cache + if not caches: + raise RuntimeError( + "runtime.self_kv_cache is empty; cannot determine KV device" + ) + + device = caches[0]["k"].device + + for layer, kv in enumerate(caches): + key_device = kv["k"].device + value_device = kv["v"].device + + if key_device != device or value_device != device: + raise RuntimeError( + "all runtime KV-cache tensors must share one device: " + f"expected {device}, but layer {layer} has " + f"key={key_device}, value={value_device}" + ) + + return device + def _frames_to_seed(self, layer_kv: dict, depth: int) -> list[int]: rt = self._rt ft, c = rt.frame_tokens, rt.chunk_size @@ -298,7 +357,9 @@ def on_runtime_created(self, runtime: Any, session_config: Any) -> None: from .keys import build_action_chain root = session_root_hash( - session_config, model_fingerprint=self.model_fingerprint + session_config, + model_fingerprint=self.model_fingerprint, + quant_config=self.mgr.cfg, ) self._ns = self.forest.get_or_create_namespace(root, root) self._actions = chunk_action_keys(runtime) @@ -321,6 +382,10 @@ def on_runtime_created(self, runtime: Any, session_config: Any) -> None: self.last_fast_forward = 0 return hint = res.resume_hint + + # _RingKVWindow exposes the actual runtime KV-cache device. The manager + # forwards it to the codec so compact quantized payloads are transferred + # before dequantization instead of restoring full KV on CPU first. if not self.mgr.materialize(hint.node, _RingKVWindow(runtime)): self._parent = self._ns.root # incomplete window -> fall back to cold run self.last_fast_forward = 0 @@ -380,8 +445,8 @@ def on_chunk_finalized( s = e - ct payload.append( ( - kv["k"][:, s:e].detach().to("cpu").clone(), - kv["v"][:, s:e].detach().to("cpu").clone(), + kv["k"][:, s:e].detach().contiguous(), + kv["v"][:, s:e].detach().contiguous(), ) ) latent = denoised.detach().to("cpu").clone() @@ -395,7 +460,6 @@ def on_chunk_finalized( "depth": idx, "payload": payload, "latent": latent, - "nbytes": sum(k.nbytes + v.nbytes for k, v in payload), } asyncio.run(self.strategy.save(self._query, None, ctx)) self._parent = ctx["node"] diff --git a/cacheseek/stores/tier.py b/cacheseek/stores/tier.py index 2d85263..26f3dac 100644 --- a/cacheseek/stores/tier.py +++ b/cacheseek/stores/tier.py @@ -17,11 +17,22 @@ import queue import threading from collections.abc import Callable, Sequence +from dataclasses import asdict from pathlib import Path from typing import Any from loguru import logger +from cacheseek.quant import ( + KVQuantizedLayer, + QuantDType, + QuantScheme, + QuantTensor, + QuantTensorSpec, + TensorRole, + TensorSpec, +) + from .base import BlobHandle, Tier @@ -137,6 +148,15 @@ class TensorStoreTierStore: def __init__(self, tensor_store: Any, *, async_put: bool = False, max_pending_chunks: int = 4) -> None: self._ts = tensor_store self._specs: dict[str, tuple[tuple[int, ...], str]] = {} # key -> (shape, dtype name) + # Payload-level sidecars describe composite logical layer formats. + # Plain KV layers do not need this: their presence is inferred from + # ``:k`` / ``:v`` tensor keys. Quantized layers do need an explicit + # marker because one logical layer expands to multiple tensors + # (qdata/scale/offset) plus codec metadata. The sidecar key + # ``:L:quant`` is therefore both the format discriminator + # and the data needed to rebuild a KVQuantizedLayer after process + # restart. + self._payload_meta: dict[str, dict[str, Any]] = {} self._async = async_put if async_put: self._q: queue.Queue = queue.Queue(maxsize=max_pending_chunks) @@ -169,10 +189,109 @@ def _get_one(self, key: str) -> Any: import torch return self._ts.get_tensor(key, shape=spec[0], dtype=getattr(torch, spec[1])) + def _put_meta(self, key: str, meta: dict[str, Any]) -> None: + """Store JSON metadata for a composite payload. + + The in-memory dict is the fast same-process path. When the backing + tensor store also exposes bytes ``put/get`` (Fluxon and local disk do), + persist the same metadata as a sidecar so mixed old/new cache contents + remain self-describing across processes. + """ + self._payload_meta[key] = meta + if hasattr(self._ts, "put"): + with contextlib.suppress(Exception): + self._ts.put(key, json.dumps(meta, sort_keys=True, separators=(",", ":")).encode()) + + def _get_meta(self, key: str) -> dict[str, Any] | None: + """Load composite-payload metadata, returning None for legacy/plain layers.""" + meta = self._payload_meta.get(key) + if meta is not None: + return meta + if hasattr(self._ts, "get"): + raw = self._ts.get(key) + if raw is not None: + meta = json.loads(raw) + self._payload_meta[key] = meta + return meta + return None + + @staticmethod + def _quant_tensor_meta(tensor: QuantTensor) -> dict[str, Any]: + tensor_meta = asdict(tensor.tensor) + quant_meta = asdict(tensor.quant) + quant_meta["role"] = tensor.quant.role.value + quant_meta["scheme"] = tensor.quant.scheme.value + quant_meta["storage_dtype"] = tensor.quant.storage_dtype.value + return { + "tensor": tensor_meta, + "quant": quant_meta, + "has_offset": tensor.offset is not None, + } + + def _put_quant_layer(self, locator: str, layer: int, payload: KVQuantizedLayer) -> None: + base = self._layer_key(locator, layer) + self._put_one(base + ":k:qdata", payload.key.qdata) + self._put_one(base + ":k:scale", payload.key.scale) + if payload.key.offset is not None: + self._put_one(base + ":k:offset", payload.key.offset) + self._put_one(base + ":v:qdata", payload.value.qdata) + self._put_one(base + ":v:scale", payload.value.scale) + if payload.value.offset is not None: + self._put_one(base + ":v:offset", payload.value.offset) + self._put_meta( + base + ":quant", + { + "kind": "kv_quantized_layer", + "key": self._quant_tensor_meta(payload.key), + "value": self._quant_tensor_meta(payload.value), + }, + ) + + def _get_quant_layer(self, locator: str, layer: int) -> KVQuantizedLayer | None: + base = self._layer_key(locator, layer) + meta = self._get_meta(base + ":quant") + if meta is None: + return None + return KVQuantizedLayer( + key=self._get_quant_tensor(base + ":k", meta["key"]), + value=self._get_quant_tensor(base + ":v", meta["value"]), + ) + + def _get_quant_tensor(self, key_prefix: str, meta: dict[str, Any]) -> QuantTensor: + tensor_meta = meta["tensor"] + quant_meta = meta["quant"] + offset = self._get_one(key_prefix + ":offset") if meta.get("has_offset") else None + return QuantTensor( + tensor=TensorSpec( + shape=tuple(tensor_meta["shape"]), + dtype=tensor_meta["dtype"], + layout=tensor_meta["layout"], + ), + quant=QuantTensorSpec( + role=TensorRole(quant_meta["role"]), + scheme=QuantScheme(quant_meta["scheme"]), + bits=quant_meta["bits"], + storage_dtype=QuantDType(quant_meta["storage_dtype"]), + group_size=quant_meta["group_size"], + group_axis=quant_meta["group_axis"], + padded_shape=tuple(quant_meta["padded_shape"]) if quant_meta["padded_shape"] is not None else None, + pack_order=quant_meta["pack_order"], + scale_dtype=quant_meta["scale_dtype"], + offset_dtype=quant_meta["offset_dtype"], + offset_kind=quant_meta["offset_kind"], + symmetric=quant_meta["symmetric"], + ), + qdata=self._get_one(key_prefix + ":qdata"), + scale=self._get_one(key_prefix + ":scale"), + offset=offset, + ) + # ----------------------------------------------------------------- async write def _do_put_payload(self, locator: str, payload: Sequence[Any]) -> None: for i, p in enumerate(payload): - if isinstance(p, (tuple, list)) and len(p) == 2: + if isinstance(p, KVQuantizedLayer): + self._put_quant_layer(locator, i, p) + elif isinstance(p, (tuple, list)) and len(p) == 2: # per-layer payload = (k, v) tensor pair -> two keys (put_tensor # takes a single tensor) self._put_one(self._layer_key(locator, i) + ":k", p[0]) @@ -222,7 +341,17 @@ def put_async( on_ready() def get_layer(self, handle: BlobHandle, layer: int) -> Any: - """Read back one layer, returning a ``(k, v)`` pair if stored as one, else a single tensor.""" + """Read back one logical layer. + + Format detection is sidecar-first: + - ``:L:quant`` present means the layer is quantized and must + be rebuilt as KVQuantizedLayer from qdata/scale/offset tensors. + - no quant sidecar means legacy/plain storage; fall back to the original + ``:k`` / ``:v`` tensor-pair layout. + """ + quantized = self._get_quant_layer(handle.locator, layer) + if quantized is not None: + return quantized k = self._get_one(self._layer_key(handle.locator, layer) + ":k") if k is not None: v = self._get_one(self._layer_key(handle.locator, layer) + ":v") diff --git a/docs/design/image/kv-quant/exact-reuse-kivi-development.png b/docs/design/image/kv-quant/exact-reuse-kivi-development.png new file mode 100644 index 0000000..ceb1a77 Binary files /dev/null and b/docs/design/image/kv-quant/exact-reuse-kivi-development.png differ diff --git a/docs/design/image/kv-quant/exact-reuse-kivi-logic.png b/docs/design/image/kv-quant/exact-reuse-kivi-logic.png new file mode 100644 index 0000000..b96dc78 Binary files /dev/null and b/docs/design/image/kv-quant/exact-reuse-kivi-logic.png differ diff --git a/docs/design/image/kv-quant/exact-reuse-kivi-physical.png b/docs/design/image/kv-quant/exact-reuse-kivi-physical.png new file mode 100644 index 0000000..5c455fa Binary files /dev/null and b/docs/design/image/kv-quant/exact-reuse-kivi-physical.png differ diff --git a/docs/design/image/kv-quant/exact-reuse-kivi-process.png b/docs/design/image/kv-quant/exact-reuse-kivi-process.png new file mode 100644 index 0000000..66e9a2d Binary files /dev/null and b/docs/design/image/kv-quant/exact-reuse-kivi-process.png differ diff --git a/docs/design/image/kv-quant/exact-reuse-kivi-scenario.png b/docs/design/image/kv-quant/exact-reuse-kivi-scenario.png new file mode 100644 index 0000000..b6f57ab Binary files /dev/null and b/docs/design/image/kv-quant/exact-reuse-kivi-scenario.png differ diff --git a/docs/design/kv-quant.md b/docs/design/kv-quant.md new file mode 100644 index 0000000..2b68973 --- /dev/null +++ b/docs/design/kv-quant.md @@ -0,0 +1,208 @@ +# KV Cache Quantization Design for CacheSeek + +CacheSeek provides KV cache reuse strategies such as exact prefix reuse and approximate reuse. As the cached context grows, storing KV cache in raw precision can introduce significant memory and storage overhead. This becomes more important in long-context, multi-session, or multi-chunk generation scenarios. + +To address this, we introduce a reusable KV cache quantization component under `cacheseek/quant`. The quantization component is designed as a common capability, and each reuse strategy can decide whether to enable it through its own configuration. + +## Quantization Methods + +CacheSeek supports 3 quantization methods: + +| Method | Description | Use Case | +|--------|-------------|----------| +| `none` | The KV cache is stored in its original precision without any quantization. | When memory and storage overhead are not a concern, or when the highest precision is required. | +| `kivi-int8` | The KV cache is quantized to 8 bits using the KIVI quantization method. | When memory and storage overhead is a concern, and a balance between precision and efficiency is desired. | +| `kivi-int4` | The KV cache is quantized to 4 bits using the KIVI quantization method. | When memory and storage overhead is very high, and a lower precision is acceptable. | + +Below is a summary of the supported quantization methods for each reuse strategy, along with their default settings: + +| Strategy | Supported Quantization Methods | Default | +|----------|----------------------------------|---------| +| Exact Prefix Reuse | `none`, `kivi-int8`, `kivi-int4` | `none` | +| Approximate Reuse | `none`| `none` | + +### KIVI Quantization + +KIVI quantization is a common quantization method, CacheSeek provides two variants: `kivi-int8` and `kivi-int4`. + +#### How KIVI Quantization Works + +KIVI stores each key/value tensor as small integer data plus per-group metadata. +For a group of floating-point values `x`, CacheSeek uses asymmetric min-offset +quantization: + +```text +levels = 2^bits - 1 +offset = min(x) +scale = (max(x) - offset) / levels +q = clamp(round((x - offset) / scale), 0, levels) +x' = q * scale + offset +``` + +If all values in a group are equal, the stored scale is `0` and decoding returns +the shared `offset` value. For `kivi-int8`, `q` is stored as `uint8`. For +`kivi-int4`, eight 4-bit values are packed into `uint32`. + +For example, with `bits=4`, a group whose values range from `-1.0` to `2.0` +uses `levels=15`, `offset=-1.0`, and `scale=0.2`. A value of `0.6` is stored as +`round((0.6 - (-1.0)) / 0.2) = 8` and decoded back to `8 * 0.2 - 1.0 = 0.6`. + +#### Key Parameters + +| Parameter | Type | Description | Default | +|-----------|------|-------------|---------| +| `bits` | int | The number of bits to quantize the KV cache to. | 8 for `kivi-int8`, 4 for `kivi-int4` | +| `group_size` | int | The size of the group for quantization. | 64 | +| `layout` | str | The layout of the input KV cache | `H,T,D` | +| `value_group_axis` | str/int | The axis along which the values are grouped for quantization. | `D` | +| `key_group_axis` | str/int | The axis along which the keys are grouped for quantization. | `T` | +| `scale_dtype` | str | The data type for the scale values used in quantization. | `float32` | +| `offset_dtype` | str | The data type for the offset values used in quantization. | `float32` | + +Detailed implementation of KIVI quantization can be found in the `cacheseek/quant/kivi.py` file. + +## How to Integrate Quantization into a Reuse Strategy + +Quantization is integrated at the reuse-strategy boundary, not in the model +runtime and not in the generic KV store. A strategy decides whether a codec is +enabled from its own config, writes quantized payloads when a chunk is committed, +and decodes them only when a cache hit is materialized back into the active KV +window. + + +### Example: Integrating KIVI into Exact-Prefix Reuse + +Using `cacheseek/reuse/exact_prefix/` as the reference implementation: + +1. Add quantization fields to the strategy config. + + `WorldKVConfig` carries the codec-facing knobs: + + ```python + quant: str = "none" + group_size: int = 64 + kv_layout: str = "H,T,D" + key_group_axis: str | int = "T" + value_group_axis: str | int = "D" + scale_dtype: str = "float32" + offset_dtype: str = "float32" + ``` + + The same fields are exposed through `quant_fingerprint()`. Include this + fingerprint in the namespace/root hash so caches written with different + quantization settings cannot be reused together. + +2. Initialize one codec when the strategy manager is constructed. + + `WorldKVManager.__init__` calls: + + ```python + self._kv_codec = build_kv_codec_from_config(cfg) + ``` + + The factory returns `None` for `quant="none"` and a `KIVICodec` for + `kivi-int8` or `kivi-int4`. Build the codec once per manager/config, then + reuse it for every chunk. Do not rebuild it per layer or per lookup. + +3. Encode before committing a finalized chunk. + + On the exact-prefix write path, `ingest()` receives `kv_payload` as a + sequence of per-layer `(key, value)` pairs. Before creating the `BlobHandle` + and calling `store.put_async`, the manager runs: + + ```python + stored_payload = self._encode_kv_payload(kv_payload) + ``` + + `_encode_kv_payload()` leaves the payload unchanged when `_kv_codec is None`. + Otherwise it calls `encode_layer(key, value)` for each layer and stores a + `KVQuantizedLayer` containing quantized key/value tensors plus scale and + offset metadata. The chunk commit then writes that encoded payload through + the tier store. The tier store remains opaque: it stores the returned Python + objects and does not need quantization-specific branching. + +4. Account for the stored representation. + + Size accounting should measure the post-quantization payload, not the + original runtime tensors. Exact-prefix reuse uses `_stored_payload_nbytes()` + to walk the stored payload and sum `qdata`, `scale`, and `offset` tensors for + each `KVQuantizedLayer`. This keeps eviction and tiering decisions aligned + with the actual cache footprint. + +5. Decode while materializing a cache hit. + + On a hit, `materialize()` walks the matched prefix path, fetches one stored + layer payload from every chunk, and converts it back to runtime KV form: + + ```python + payload = self.store.get_layer(n.blob, layer) + runtime_kv = self._decode_layer_payload(payload) + ``` + + `_decode_layer_payload()` returns raw payloads directly when quantization is + disabled. When a codec is enabled, it requires the stored object to be a + `KVQuantizedLayer` and calls `decode_layer(payload)` to reconstruct the + runtime `(key, value)` tensors. The decoded tensors are passed to + `window.seed_layer(...)`, so the runtime window receives the same logical KV + shape regardless of how the cache was stored. + +#### Illustration + +
+Process View + +![Exact-prefix KIVI process](image/kv-quant/exact-reuse-kivi-process.png) + +
+ +
+Logic View + +![Exact-prefix KIVI logic](image/kv-quant/exact-reuse-kivi-logic.png) + +
+ +
+Physical View + +![Exact-prefix KIVI physical](image/kv-quant/exact-reuse-kivi-physical.png) + +
+ +
+Development View + +![Exact-prefix KIVI development](image/kv-quant/exact-reuse-kivi-development.png) + +
+ +
+Scenario View + +The scenario view describes the end-to-end request flow with exact prefix reuse and optional KV quantization. + +![Exact-prefix KIVI scenario](image/kv-quant/exact-reuse-kivi-scenario.png) + +
+ + + +### New Reuse Strategies + +For a new reuse strategy, keep the same boundary: + +- Put quantization options in the strategy config. +- Add the quantization fingerprint to any cache namespace or compatibility hash. +- Build the codec once from config. +- Encode exactly once on the write path, immediately before storage. +- Decode exactly once on the read path, immediately before seeding runtime KV. +- Keep storage backends unaware of quantization internals unless a backend adds a + specialized tensor transport optimization. + +## Performance Considerations + +Tests/Experiments wait to be added here. + +## References + +- The KIVI quantization method is described in the paper "KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache" (https://arxiv.org/abs/2402.02750). diff --git a/examples/exact_prefix_reuse/README.md b/examples/exact_prefix_reuse/README.md index dc95132..357a381 100644 --- a/examples/exact_prefix_reuse/README.md +++ b/examples/exact_prefix_reuse/README.md @@ -86,6 +86,8 @@ and assertion outcomes. - `--frame-num 81 --prefix-chunks 4` — longer run, exercises multi-round rolling eviction - `--store localdisk --disk-root /path/on/local/nvme` — local-disk blob backend - `--store fluxon --fluxon-config /path/to/external_config.yaml` — Fluxon backend (stack setup: [`../../scripts/fluxon/README.md`](../../scripts/fluxon/README.md)) +- `--quant kivi_int8` or `--quant kivi_int4` — enable KIVI KV quantization for + exact-prefix payloads. - omit `--image-path` / `--action-path` to fall back to a synthetic image + trajectory > If the `cacheseek` (or `telefuser`) in your venv is an editable install pointing diff --git a/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py b/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py index 0a4f2e2..cd85d38 100644 --- a/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py +++ b/examples/exact_prefix_reuse/e2e_telefuser_lingbot.py @@ -69,8 +69,7 @@ from cacheseek.reuse.exact_prefix import NamespaceForest, WorldKVManager from cacheseek.reuse.exact_prefix.telefuser_lingbot import ( LingBotWorldKVBinding, - make_full_kv_config, - make_rolling_config, + make_world_kv_config, ) from cacheseek.stores import InMemoryTierStore @@ -88,6 +87,14 @@ ORIG_H, ORIG_W = 480, 832 +def parse_group_axis(axis: str) -> str | int: + """Accept layout labels such as T/D or integer axis indexes from CLI.""" + axis = axis.strip() + if axis.lstrip("-").isdigit(): + return int(axis) + return axis + + def load_image(image_path: str) -> Image.Image: """Prefer a real reference image (lingbot-world asset, same input as production); fall back to a deterministic synthetic image when none is given.""" @@ -148,6 +155,7 @@ def run_request(pipeline, binding, frame_num: int, poses, intrinsics, *, image_p control_mode="cam", frame_num=frame_num, seed=SEED, + sammple_shift=10, # better generation quality poses=poses, intrinsics=intrinsics, world_kv_binding=binding, @@ -200,6 +208,20 @@ def main() -> int: ap.add_argument("--disk-root", type=str, default="/tmp/worldkv_diskstore", help="root dir for the localdisk backend (should be real local disk)") ap.add_argument("--repeat", type=int, default=1, help="number of A/B/C/D loop iterations (weights loaded once; for soak testing, each round uses an independent cache stack)") + ap.add_argument("--quant", type=str, default="none", choices=["none", "kivi_int8", "kivi_int4"], + help="optional KIVI KV quantization for exact-prefix world-kv payloads") + ap.add_argument("--group-size", type=int, default=64, + help="KIVI quantization group size when --quant is enabled") + ap.add_argument("--kv-layout", type=str, default="B,T,H,D", + help="logical layout of TeleFuser self-attention KV tensors; LingBot uses (batch, tokens, heads, head_dim)") + ap.add_argument("--key-group-axis", type=str, default="T", + help="KIVI grouped axis for keys, as a layout label or integer axis") + ap.add_argument("--value-group-axis", type=str, default="D", + help="KIVI grouped axis for values, as a layout label or integer axis") + ap.add_argument("--scale-dtype", type=str, default="float32", + help="dtype used to store KIVI scale tensors") + ap.add_argument("--offset-dtype", type=str, default="float32", + help="dtype used to store KIVI offset tensors") ap.add_argument("--aux-device", type=str, default="", help="two-GPU usage: place the T5 text encoder (~11GB) + VAE on this device (e.g. cuda:1), keep the DiT on the main device. " "Note: the LingBot pipeline's DiT itself has no tensor parallelism — two GPUs split memory placement, not compute") @@ -242,16 +264,25 @@ def make_store(): # collide); the empty-cache cold run is guaranteed by a fresh forest (trie index) -- # old payloads in the store are unreachable without an index. shared_store = make_store() + key_group_axis = parse_group_axis(args.key_group_axis) + value_group_axis = parse_group_axis(args.value_group_axis) def fresh_stack(): forest = NamespaceForest() - world_kv_cfg = ( - make_full_kv_config() - if int(cfg.local_attn_size) == -1 - else make_rolling_config(local_attn_size=cfg.local_attn_size, sink_size=cfg.sink_size, chunk_size=3) - ) mgr = WorldKVManager( - forest, shared_store, world_kv_cfg, + forest, shared_store, + make_world_kv_config( + local_attn_size=cfg.local_attn_size, + sink_size=cfg.sink_size, + chunk_size=3, + quant=args.quant, + group_size=args.group_size, + kv_layout=args.kv_layout, + key_group_axis=key_group_axis, + value_group_axis=value_group_axis, + scale_dtype=args.scale_dtype, + offset_dtype=args.offset_dtype, + ), ) return forest, mgr @@ -288,6 +319,15 @@ def fresh_stack(): manifest = { "frame_num": args.frame_num, "prefix_chunks": args.prefix_chunks, "seed": SEED, "store": args.store, "pipeline": {"local_attn_size": cfg.local_attn_size, "sink_size": cfg.sink_size}, + "world_kv_quant": { + "quant": args.quant, + "group_size": args.group_size, + "kv_layout": args.kv_layout, + "key_group_axis": key_group_axis, + "value_group_axis": value_group_axis, + "scale_dtype": args.scale_dtype, + "offset_dtype": args.offset_dtype, + }, "torch": torch.__version__, "device": torch.cuda.get_device_name(0), "results": results, "checks": checks, "all_pass": all(checks.values()), } diff --git a/tests/test_reuse_strategy_protocol.py b/tests/test_reuse_strategy_protocol.py index 3cf6063..dc75830 100644 --- a/tests/test_reuse_strategy_protocol.py +++ b/tests/test_reuse_strategy_protocol.py @@ -52,7 +52,7 @@ def _ingest_chain(strategy, forest, root, actions): for i, a in enumerate(actions): ctx = { "ns": ns, "parent": parent, "action": a, "node_key": chain[i], "depth": i, - "payload": [(f"k{i}", f"v{i}")], "latent": f"x0@{i}", "nbytes": 0, + "payload": [(f"k{i}", f"v{i}")], "latent": f"x0@{i}", } asyncio.run(strategy.save(q, None, ctx)) parent = ctx["node"]