Add TurboQuant KV cache compression to serving pipeline

examples/model/qwen3_14b/npu_generate.py

@@ -36,7 +36,7 @@ def _bootstrap_package_root() -> None:
 from python.core.parallel import ParallelConfig, parse_device_ids
 from python.profile import get_profiler, merge_profile, profile_span
 from examples.model.qwen3_14b.runner.npu_executor import Qwen314BPyptoExecutor as PyptoExecutor
-from python.core.types import LoadedModel
+from python.core.types import KvQuantConfig, LoadedModel
 import dataclasses
 
 
@@ -410,6 +410,12 @@ def build_parser() -> argparse.ArgumentParser:
         help="Implies --profile. Also dump per-layer prefill times and "
              "per-decode-step layer breakdowns.",
     )
+    parser.add_argument(
+        "--tq",
+        action="store_true",
+        dest="tq_mode",
+        help="Enable TurboQuant KV cache compression (4-bit quantized K/V cache).",
+    )
     return parser
 
 
@@ -445,6 +451,7 @@ def main() -> None:
         device_ids=device_ids,
         save_kernels_dir=args.save_kernels_dir,
         l3_trace=args.profile_verbose,
+        tq_mode=args.tq_mode,
     )
     engine = LLMEngine(
         kv_cache_manager=kv_cache_manager,
@@ -468,12 +475,9 @@ def main() -> None:
                 device="cpu",
                 kv_dtype=args.kv_cache_dtype,
                 weight_dtype=args.dtype,
+                kv_quant_config=KvQuantConfig(enabled=True) if args.tq_mode else None,
                 npu_memory_utilization=args.npu_memory_utilization,
                 max_num_batched_tokens=args.max_num_batched_tokens,
-                # Conservative default — the decode kernel is compiled
-                # with this baked-in shape and cannot be resized later.
-                # 200 pages x 128 tokens = 25 600 tokens total capacity.
-                total_kv_pages=200,
             ),
         )
         if collector is not None:

examples/model/qwen3_14b/runner/npu_runner.py

@@ -107,6 +107,10 @@ class _CompiledKernels:
     decode_token_ids_buffer: torch.Tensor
     decode_sampled_ids_buffer: torch.Tensor
     decode_next_hidden_buffer: torch.Tensor
+    # TurboQuant (TQ) artifacts. Populated only when tq_mode=True.
+    tq_mode: bool = False
+    rot_matrices: torch.Tensor | None = None
+    tq_codebook: torch.Tensor | None = None
 
 
 @dataclass
@@ -164,6 +168,9 @@ class _StaticKernelArgs:
     padded_lm_head_weight: _StaticDeviceTensor
     padded_embed_weight: _StaticDeviceTensor
     decode_weights: dict[str, _StaticDeviceTensor]
+    # TurboQuant static artifacts (None unless tq_mode).
+    rot_matrices: _StaticDeviceTensor | None = None
+    tq_codebook: _StaticDeviceTensor | None = None
 
 
 class Qwen314BModelRunner(ModelRunner):
@@ -174,10 +181,12 @@ def __init__(
         *,
         compiled: _CompiledKernels,
         device_id: int = 0,
+        tq_mode: bool = False,
     ) -> None:
         super().__init__()
         self._compiled = compiled
         self._device_id = device_id
+        self._tq_mode = tq_mode
         self._l3_worker: Any | None = None
         self._l3_static_tensors: dict[tuple[int, tuple[int, ...], torch.dtype], object] = {}
         self._static_args: _StaticKernelArgs | None = None
@@ -245,11 +254,7 @@ def _alloc_kv_cache_with_retry(
             try:
                 logger.info(f"[init_kv_cache] num_pages={num_pages}, allocating …")
                 ModelRunner.init_kv_cache(self, model_id, config, runtime, num_pages=num_pages)
-                bytes_per_page = (
-                    config.num_hidden_layers * 2 * config.num_key_value_heads
-                    * runtime.page_size * config.head_dim
-                    * getattr(torch, runtime.kv_dtype).itemsize
-                )
+                bytes_per_page = Qwen314BModelRunner._kv_bytes_per_page(config, runtime)
                 logger.info(
                     f"[init_kv_cache] allocated {num_pages} pages "
                     f"(requested {requested}, downgraded after OOM): "
@@ -270,6 +275,22 @@ def _alloc_kv_cache_with_retry(
             f"KV cache allocation failed even at floor {floor} pages"
         )
 
+    @staticmethod
+    def _kv_bytes_per_page(config: ModelConfig, runtime: RuntimeConfig) -> int:
+        """Device HBM bytes per KV page, accounting for TurboQuant compression.
+
+        Non-TQ: bf16 K+V (``2 * head_dim * dtype_bytes`` per row).
+        TQ: uint8 quantized K+V (1 byte/elem) + one fp32 scale per K and per V
+        row (4 bytes each) — roughly half of bf16, so the same HBM budget fits
+        ~2x the pages.
+        """
+        per_page_rows = config.num_hidden_layers * config.num_key_value_heads * runtime.page_size
+        qcfg = getattr(runtime, "kv_quant_config", None)
+        if qcfg is not None and qcfg.enabled:
+            return per_page_rows * (2 * config.head_dim + 2 * 4)
+        dtype_bytes = getattr(torch, runtime.kv_dtype).itemsize
+        return per_page_rows * 2 * config.head_dim * dtype_bytes
+
     @staticmethod
     def _compute_kv_cache_pages(config: ModelConfig, runtime: RuntimeConfig, device_id: int = 0) -> int:
         """Compute KV cache pages, vLLM-style: total x utilization − peak_non_kv.
@@ -282,19 +303,17 @@ def _compute_kv_cache_pages(config: ModelConfig, runtime: RuntimeConfig, device_
         than ``free x fraction`` whose headroom shrinks when free is small).
         """
         free_bytes, total_bytes = torch.npu.mem_get_info(f"npu:{device_id}")
-        dtype_bytes = getattr(torch, runtime.kv_dtype).itemsize
-        bytes_per_page = (
-            config.num_hidden_layers * 2 * config.num_key_value_heads
-            * runtime.page_size * config.head_dim * dtype_bytes
-        )
+        bytes_per_page = Qwen314BModelRunner._kv_bytes_per_page(config, runtime)
+        qcfg = getattr(runtime, "kv_quant_config", None)
+        tq = qcfg is not None and qcfg.enabled
         utilization = getattr(runtime, "npu_memory_utilization", 0.90)
         peak_non_kv = total_bytes - free_bytes
         kv_budget = int(total_bytes * utilization - peak_non_kv)
         num_pages = max(kv_budget // bytes_per_page, 1)
         logger.info(
-            "KV cache sizing (vLLM-style): total=%.2f GB, utilization=%.2f, "
+            "KV cache sizing (vLLM-style%s): total=%.2f GB, utilization=%.2f, "
             "peak_non_kv=%.2f GB, kv_budget=%.2f GB, requested_pages=%d (%.1f MB/page)",
-            total_bytes / 1e9, utilization, peak_non_kv / 1e9,
+            ", TQ uint8" if tq else "", total_bytes / 1e9, utilization, peak_non_kv / 1e9,
             kv_budget / 1e9, num_pages, bytes_per_page / 1e6,
         )
         return num_pages
@@ -303,6 +322,7 @@ def _compute_kv_cache_pages(config: ModelConfig, runtime: RuntimeConfig, device_
     def _print_memory_breakdown(
         label: str, config: ModelConfig, runtime: RuntimeConfig, num_pages: int,
         device_id: int = 0,
+        tq_mode: bool = False,
     ) -> None:
         """Print a per-component NPU memory breakdown at ``label``.
 
@@ -331,10 +351,7 @@ def _print_memory_breakdown(
         weight_bytes = int(wt_params * dtype_bytes)
 
         # KV cache — exact (num_pages already reflects the real allocation).
-        bytes_per_page = (
-            config.num_hidden_layers * 2 * config.num_key_value_heads
-            * runtime.page_size * config.head_dim * dtype_bytes
-        )
+        bytes_per_page = Qwen314BModelRunner._kv_bytes_per_page(config, runtime)
         kv_bytes = num_pages * bytes_per_page
 
         # Simpler ring-heap arena — from env (matches _compute_kv_cache_pages).
@@ -439,7 +456,7 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
         self._run_distributed_program(
             compiled.prefill,
             *self._prefill_kernel_args(
-                prefill_inputs, kv_cache.key_pages, kv_cache.value_pages,
+                prefill_inputs, kv_cache,
                 compiled.prefill_logits_buffer,
             ),
         )
@@ -469,7 +486,7 @@ def _warmup_dispatch(self, runtime: RuntimeConfig) -> None:
         t0 = time.perf_counter()
         self._run_distributed_program(
             compiled.decode,
-            *self._decode_kernel_args(decode_kernel_inputs, kv_cache.key_pages, kv_cache.value_pages),
+            *self._decode_kernel_args(decode_kernel_inputs, kv_cache),
         )
         logger.info(f"[warmup] decode done ({time.perf_counter() - t0:.2f} s)")
 
@@ -569,6 +586,16 @@ def _build_static_kernel_args(self) -> _StaticKernelArgs:
                 name: self._static_device_tensor(tensor)
                 for name, tensor in compiled.decode_weights.items()
             },
+            rot_matrices=(
+                self._static_device_tensor(compiled.rot_matrices)
+                if compiled.rot_matrices is not None
+                else None
+            ),
+            tq_codebook=(
+                self._static_device_tensor(compiled.tq_codebook)
+                if compiled.tq_codebook is not None
+                else None
+            ),
         )
 
     def _require_static_args(self) -> _StaticKernelArgs:
@@ -585,13 +612,12 @@ def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult
         logits_padded = compiled.prefill_logits_buffer
 
         kv_cache = self._materialize_kv_cache(model)
-        k_cache = kv_cache.key_pages
-        v_cache = kv_cache.value_pages
-        self._validate_kv_cache_bounds(model, prefill_inputs.block_table, prefill_inputs.slot_mapping, k_cache)
+        bound_cache = kv_cache.quant_k_pages if self._compiled.tq_mode else kv_cache.key_pages
+        self._validate_kv_cache_bounds(model, prefill_inputs.block_table, prefill_inputs.slot_mapping, bound_cache)
 
         self._run_distributed_program(
             compiled.prefill,
-            *self._prefill_kernel_args(prefill_inputs, k_cache, v_cache, logits_padded),
+            *self._prefill_kernel_args(prefill_inputs, kv_cache, logits_padded),
         )
 
         for batch_idx, alloc in enumerate(batch.kv_allocations):
@@ -634,18 +660,17 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult:
         kv_cache = self._kv_caches.get(model_id)
         if kv_cache is None:
             raise RuntimeError(f"KV cache for model {model_id!r} is not initialized")
-        k_cache = kv_cache.key_pages
-        v_cache = kv_cache.value_pages
+        bound_cache = kv_cache.quant_k_pages if self._compiled.tq_mode else kv_cache.key_pages
 
         kernel_inputs = self._pad_decode_inputs(model, decode_inputs)
 
         # Padded block_table / slot_mapping only ever reference row 0's
         # already-valid pages, so bound-check exactly what the kernel will read.
-        self._validate_kv_cache_bounds(model, kernel_inputs.block_table, kernel_inputs.slot_mapping, k_cache)
+        self._validate_kv_cache_bounds(model, kernel_inputs.block_table, kernel_inputs.slot_mapping, bound_cache)
 
         self._run_distributed_program(
             compiled.decode,
-            *self._decode_kernel_args(kernel_inputs, k_cache, v_cache),
+            *self._decode_kernel_args(kernel_inputs, kv_cache),
         )
         for batch_idx, alloc in enumerate(batch.kv_allocations):
             alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item()))
@@ -715,13 +740,44 @@ def _maybe_run_sample_embed(
     def _prefill_kernel_args(
         self,
         inputs: _PrefillInputs,
-        k_cache: DeviceTensor,
-        v_cache: DeviceTensor,
+        kv_cache: Any,
         logits: torch.Tensor,
     ) -> tuple[Any, ...]:
-        """Return arguments in ``qwen3_prefill_host`` signature order."""
+        """Return arguments in ``qwen3_prefill_host`` (or ``_tq_host``) order."""
         static = self._require_static_args()
         weights = static.decode_weights
+        if self._compiled.tq_mode:
+            # qwen3_prefill_tq_host: 26 params, no chunk_lens/offsets; uint8
+            # quantized K/V caches + (cache_rows,1) FP32 scales; rotation
+            # matrices + codebook; weight order wo, post_rms, w_gate, w_up, w_down.
+            return (
+                inputs.hidden,
+                inputs.seq_lens,
+                weights["decode_input_rms_weight"],
+                weights["decode_wq"],
+                weights["decode_wk"],
+                weights["decode_wv"],
+                weights["decode_q_norm_weight"],
+                weights["decode_k_norm_weight"],
+                static.rope_cos,
+                static.rope_sin,
+                inputs.block_table,
+                inputs.slot_mapping,
+                kv_cache.quant_k_pages,
+                kv_cache.quant_v_pages,
+                kv_cache.k_scales_pages,
+                kv_cache.v_scales_pages,
+                static.rot_matrices,
+                static.tq_codebook,
+                weights["decode_wo"],
+                weights["decode_post_rms_weight"],
+                weights["decode_w_gate"],
+                weights["decode_w_up"],
+                weights["decode_w_down"],
+                static.final_norm_weight,
+                static.padded_lm_head_weight,
+                logits,
+            )
         return (
             inputs.hidden,
             inputs.seq_lens,
@@ -737,8 +793,8 @@ def _prefill_kernel_args(
             static.rope_sin,
             inputs.block_table,
             inputs.slot_mapping,
-            k_cache,
-            v_cache,
+            kv_cache.key_pages,
+            kv_cache.value_pages,
             weights["decode_wo"],
             weights["decode_w_gate"],
             weights["decode_w_up"],
@@ -752,12 +808,43 @@ def _prefill_kernel_args(
     def _decode_kernel_args(
         self,
         inputs: _DecodeKernelInputs,
-        k_cache: DeviceTensor,
-        v_cache: DeviceTensor,
+        kv_cache: Any,
     ) -> tuple[Any, ...]:
-        """Return arguments in ``qwen3_decode_host`` signature order."""
+        """Return arguments in ``qwen3_decode_host`` (or ``_tq_host``) order."""
         static = self._require_static_args()
         weights = static.decode_weights
+        if self._compiled.tq_mode:
+            # qwen3_decode_tq_host: 26 params. No device-side embed/token_ids/
+            # sampled buffers (TQ samples on CPU); uint8 quantized caches +
+            # scales + rotation/codebook; weight order wo, post_rms, w_gate...
+            return (
+                inputs.hidden,
+                weights["decode_input_rms_weight"],
+                weights["decode_wq"],
+                weights["decode_wk"],
+                weights["decode_wv"],
+                weights["decode_q_norm_weight"],
+                weights["decode_k_norm_weight"],
+                inputs.seq_lens,
+                inputs.block_table,
+                inputs.slot_mapping,
+                static.rope_cos,
+                static.rope_sin,
+                kv_cache.quant_k_pages,
+                kv_cache.quant_v_pages,
+                kv_cache.k_scales_pages,
+                kv_cache.v_scales_pages,
+                static.rot_matrices,
+                static.tq_codebook,
+                weights["decode_wo"],
+                weights["decode_post_rms_weight"],
+                weights["decode_w_gate"],
+                weights["decode_w_up"],
+                weights["decode_w_down"],
+                static.final_norm_weight,
+                static.padded_lm_head_weight,
+                inputs.logits,
+            )
         return (
             inputs.hidden,
             weights["decode_input_rms_weight"],
@@ -771,8 +858,8 @@ def _decode_kernel_args(
             inputs.slot_mapping,
             static.rope_cos,
             static.rope_sin,
-            k_cache,
-            v_cache,
+            kv_cache.key_pages,
+            kv_cache.value_pages,
             weights["decode_wo"],
             weights["decode_w_gate"],
             weights["decode_w_up"],
@@ -886,12 +973,18 @@ def _shared_l3_worker(self) -> Any:
         if worker is None:
             from pypto.runtime import DistributedWorker  # noqa: PLC0415
 
-            worker = DistributedWorker([
+            # greedy_sample / token_embed are None in TQ mode (TQ samples on
+            # CPU); worker.run dispatches by compiled-program reference, so the
+            # program set may omit them without changing prefill/decode indexing.
+            programs = [
                 self._compiled.prefill.compiled,
                 self._compiled.decode.compiled,
-                self._compiled.greedy_sample.compiled,
-                self._compiled.token_embed.compiled,
-            ])
+            ]
+            if self._compiled.greedy_sample is not None:
+                programs.append(self._compiled.greedy_sample.compiled)
+            if self._compiled.token_embed is not None:
+                programs.append(self._compiled.token_embed.compiled)
+            worker = DistributedWorker(programs)
             self._l3_worker = worker
         return worker