From 0d524e36eef599214995b0696f67bc1c443de8d4 Mon Sep 17 00:00:00 2001 From: huangzhuo Date: Tue, 19 May 2026 14:42:20 +0800 Subject: [PATCH 01/17] [Feature] Add TurboQuant --- examples/model/qwen3_14b/runner/npu_runner.py | 21 +- python/core/kv_cache.py | 163 ++++++++++++++ python/core/turboquant/__init__.py | 5 + python/core/turboquant/compressor.py | 207 ++++++++++++++++++ python/core/turboquant/lloyd_max.py | 79 +++++++ python/core/types.py | 13 ++ 6 files changed, 487 insertions(+), 1 deletion(-) create mode 100644 python/core/turboquant/__init__.py create mode 100644 python/core/turboquant/compressor.py create mode 100644 python/core/turboquant/lloyd_max.py diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index 5bc3bb4..3f993e9 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -215,6 +215,13 @@ def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult last_hidden_rows.append(hidden[batch_idx, seq_len - 1].float()) last_hidden = torch.stack(last_hidden_rows) logits = self._project_logits(model, last_hidden) + + # ── KV quantization: compress old tokens after prefill ── + model_id = model.config.model_id + if self._kv_cache_manager.is_quantization_enabled(model_id): + for alloc in batch.kv_allocations: + self._kv_cache_manager.compress_old_tokens(model_id, alloc) + return PrefillResult(last_hidden=last_hidden, logits=logits) def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: @@ -229,8 +236,14 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: hidden = decode_inputs.hidden dw = compiled.decode_weights + # ── KV quantization: restore compressed old tokens before kernel ── + model_id = model.config.model_id + if self._kv_cache_manager.is_quantization_enabled(model_id): + for alloc in batch.kv_allocations: + self._kv_cache_manager.restore_compressed_tokens(model_id, alloc) + k_cache, v_cache = self._kv_cache_manager.materialize_full_layer_cache( - model.config.model_id, + model_id, ) refresh_kv_cache = model.config.model_id in self._l2_dirty_kv_models out = torch.zeros_like(hidden) @@ -290,6 +303,12 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: logits = self._project_logits(model, final_hidden) for batch_idx, alloc in enumerate(batch.kv_allocations): alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item())) + + # ── KV quantization: compress old tokens after kernel ── + if self._kv_cache_manager.is_quantization_enabled(model_id): + for alloc in batch.kv_allocations: + self._kv_cache_manager.compress_old_tokens(model_id, alloc) + return DecodeResult(hidden_states=final_hidden, logits=logits) def _project_logits(self, model: RuntimeModel, hidden: torch.Tensor) -> torch.Tensor: diff --git a/python/core/kv_cache.py b/python/core/kv_cache.py index 481373b..7af385f 100644 --- a/python/core/kv_cache.py +++ b/python/core/kv_cache.py @@ -16,6 +16,27 @@ from .types import KvAllocation, ModelConfig, RuntimeConfig +# Lazy import to avoid hard dependency when quantization is disabled. +_KVCOMPRESSOR = None + + +def _get_kv_compressor_cls(): + global _KVCOMPRESSOR + if _KVCOMPRESSOR is None: + from .turboquant.compressor import KVCompressor + _KVCOMPRESSOR = KVCompressor + return _KVCOMPRESSOR + + +@dataclass +class _CompressedSegment: + """Compressed KV data for a contiguous range of old tokens in one layer.""" + + token_start: int + token_count: int + compressed_k: dict + compressed_v: dict + @dataclass class _CachePool: @@ -29,6 +50,11 @@ class _CachePool: key_pages: torch.Tensor value_pages: torch.Tensor free_pages: list[int] + # KV quantization (None when disabled) + kv_compressor: object = None # KVCompressor | None + kv_quant_config: object = None # KvQuantConfig | None + # Compressed segments: {request_id: {layer_idx: _CompressedSegment}} + compressed_segments: dict = None class KvCacheManager: @@ -57,6 +83,20 @@ def register_model(self, model_id: str, config: ModelConfig, runtime: RuntimeCon device=runtime.device, ) value_pages = torch.zeros_like(key_pages) + # Initialize KV quantization if configured + kv_compressor = None + kv_quant_config = runtime.kv_quant_config + compressed_segments = None + if kv_quant_config is not None and kv_quant_config.enabled: + KVCompressor = _get_kv_compressor_cls() + kv_compressor = KVCompressor( + head_dim=config.head_dim, + num_layers=config.num_hidden_layers, + config=kv_quant_config, + device=runtime.device, + ) + compressed_segments = {} + self._pools[model_id] = _CachePool( page_size=runtime.page_size, num_layers=config.num_hidden_layers, @@ -66,6 +106,9 @@ def register_model(self, model_id: str, config: ModelConfig, runtime: RuntimeCon key_pages=key_pages, value_pages=value_pages, free_pages=list(range(num_pages - 1, -1, -1)), + kv_compressor=kv_compressor, + kv_quant_config=kv_quant_config, + compressed_segments=compressed_segments, ) def allocate_for_prompt(self, model_id: str, request_id: str, prompt_len: int) -> KvAllocation: @@ -255,6 +298,126 @@ def free(self, alloc: KvAllocation) -> None: alloc.page_ids.clear() alloc.tokens_capacity = 0 alloc.tokens_used = 0 + # Also clean up compressed segments for this request + if pool.compressed_segments is not None: + pool.compressed_segments.pop(alloc.request_id, None) + + # ── KV cache quantization ── + + def is_quantization_enabled(self, model_id: str) -> bool: + """Return whether KV cache quantization is enabled for a model.""" + pool = self._pool(model_id) + return pool.kv_compressor is not None + + def restore_compressed_tokens(self, model_id: str, alloc: KvAllocation) -> None: + """Decompress old tokens back into bf16 cache pages before kernel execution. + + This writes the previously compressed old tokens back into their original + page locations in key_pages/value_pages so the kernel can read them. + """ + pool = self._pool(model_id) + if pool.compressed_segments is None or pool.kv_compressor is None: + return + req_segments = pool.compressed_segments.get(alloc.request_id) + if not req_segments: + return + + for layer_idx, segment in req_segments.items(): + # Decompress + keys, values = pool.kv_compressor.decompress_layer( + layer_idx, segment.compressed_k, segment.compressed_v, + ) + # keys, values: (S_old, H, D) — write back into paged cache + self._write_tokens_from_compressed( + pool, layer_idx, alloc, segment.token_start, keys, values, + ) + + def compress_old_tokens(self, model_id: str, alloc: KvAllocation) -> None: + """Compress old tokens (beyond residual_window) and store compressed data. + + Tokens with index < (tokens_used - residual_window) are compressed. + The bf16 cache pages are NOT freed (they are reused in-place on next restore). + """ + pool = self._pool(model_id) + if pool.kv_compressor is None or pool.kv_quant_config is None: + return + + residual_window = pool.kv_quant_config.residual_window + tokens_used = alloc.tokens_used + if tokens_used <= residual_window: + return # Not enough tokens to compress + + old_token_count = tokens_used - residual_window + + # Initialize compressed segments dict for this request + if pool.compressed_segments is None: + pool.compressed_segments = {} + if alloc.request_id not in pool.compressed_segments: + pool.compressed_segments[alloc.request_id] = {} + + for layer_idx in range(pool.num_layers): + # Read old tokens from bf16 cache + old_keys, old_values = self.read_context( + layer_idx, alloc, upto_tokens=old_token_count, + ) + # old_keys, old_values: (S_old, H, D) + + # Compress + compressed_k, compressed_v = pool.kv_compressor.compress_layer( + layer_idx, old_keys, old_values, + ) + + # Store + pool.compressed_segments[alloc.request_id][layer_idx] = _CompressedSegment( + token_start=0, + token_count=old_token_count, + compressed_k=compressed_k, + compressed_v=compressed_v, + ) + + def _write_tokens_from_compressed( + self, + pool: _CachePool, + layer_idx: int, + alloc: KvAllocation, + start_token_index: int, + keys: torch.Tensor, + values: torch.Tensor, + ) -> None: + """Write decompressed tokens back into paged cache. + + keys, values: (S, H, D) — will be cast to the cache dtype. + """ + cache_dtype = pool.key_pages.dtype + for row in range(keys.shape[0]): + token_index = start_token_index + row + page_idx = token_index // pool.page_size + offset = token_index % pool.page_size + physical_page = alloc.page_ids[page_idx] + pool.key_pages[layer_idx, physical_page, :, offset, :] = keys[row].to(cache_dtype) + pool.value_pages[layer_idx, physical_page, :, offset, :] = values[row].to(cache_dtype) + + def quantization_stats(self, model_id: str) -> dict: + """Return memory usage stats for compressed KV cache.""" + pool = self._pool(model_id) + if pool.compressed_segments is None: + return {"enabled": False} + total_compressed_bytes = 0 + total_segments = 0 + for req_segments in pool.compressed_segments.values(): + for segment in req_segments.values(): + for tensor in segment.compressed_k.values(): + if isinstance(tensor, torch.Tensor): + total_compressed_bytes += tensor.numel() * tensor.element_size() + for tensor in segment.compressed_v.values(): + if isinstance(tensor, torch.Tensor): + total_compressed_bytes += tensor.numel() * tensor.element_size() + total_segments += 1 + return { + "enabled": True, + "total_segments": total_segments, + "compressed_bytes": total_compressed_bytes, + } def _pool(self, model_id: str) -> _CachePool: """Return the registered cache pool for a model.""" diff --git a/python/core/turboquant/__init__.py b/python/core/turboquant/__init__.py new file mode 100644 index 0000000..e0e4e78 --- /dev/null +++ b/python/core/turboquant/__init__.py @@ -0,0 +1,5 @@ +"""TurboQuant KV cache compression module.""" + +from .compressor import KVCompressor, KvQuantConfig, MSECompressor + +__all__ = ["KVCompressor", "KvQuantConfig", "MSECompressor"] diff --git a/python/core/turboquant/compressor.py b/python/core/turboquant/compressor.py new file mode 100644 index 0000000..9ba6c24 --- /dev/null +++ b/python/core/turboquant/compressor.py @@ -0,0 +1,207 @@ +"""MSE-optimal KV cache compressor using random rotation + Lloyd-Max quantization. + +Adapted from TurboQuant V3 (MSE-only, no QJL). +Provides per-layer compressors with asymmetric key/value bit-widths. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from .lloyd_max import LloydMaxCodebook + + +def _generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor: + """Generate a random orthogonal rotation matrix via QR decomposition.""" + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + G = torch.randn(d, d, generator=gen) + Q, R = torch.linalg.qr(G) + diag_sign = torch.sign(torch.diag(R)) + diag_sign[diag_sign == 0] = 1.0 + Q = Q * diag_sign.unsqueeze(0) + return Q.to(device) + + +class MSECompressor: + """Single-stage MSE-optimal compressor for one side (keys or values). + + Compress normalizes to unit sphere, quantizes with Lloyd-Max, + and stores bit-packed indices + norms. + """ + + def __init__(self, head_dim: int, bits: int, seed: int, device: str = "cpu"): + self.head_dim = head_dim + self.bits = bits + self.device = device + + self.Pi = _generate_rotation_matrix(head_dim, seed=seed, device=device) + self.centroids = LloydMaxCodebook(head_dim, bits).centroids.to(device) + + @torch.no_grad() + def compress(self, states: torch.Tensor) -> dict: + """Compress (B, H, S, D) -> dict with bit-packed indices + norms.""" + B, H, S, D = states.shape + N = B * H * S + flat = states.reshape(N, D).float() + + # Normalize to unit sphere, store norms + vec_norms = torch.norm(flat, dim=-1) # (N,) + flat_norm = flat / (vec_norms.unsqueeze(-1) + 1e-8) + + # Rotate + quantize + rotated = flat_norm @ self.Pi.T + diffs = rotated.unsqueeze(-1) - self.centroids # (N, D, levels) + indices = diffs.abs().argmin(dim=-1).to(torch.uint8) # (N, D) + + # Bit-pack indices + indices_per_byte = 8 // self.bits + idx_pad = (indices_per_byte - D % indices_per_byte) % indices_per_byte + idx_flat = indices.long() + if idx_pad: + idx_flat = F.pad(idx_flat, (0, idx_pad)) + n_groups = idx_flat.shape[-1] // indices_per_byte + idx_powers = torch.tensor( + [2 ** (self.bits * i) for i in range(indices_per_byte - 1, -1, -1)], + dtype=torch.long, + device=idx_flat.device, + ) + idx_bytes = ( + (idx_flat.reshape(N, n_groups, indices_per_byte) * idx_powers) + .sum(-1) + .to(torch.uint8) + ) + + return { + "idx_bytes": idx_bytes.reshape(B, H, S, n_groups), + "vec_norms": vec_norms.to(torch.float16).reshape(B, H, S), + "shape": (B, H, S, D), + "idx_pad": idx_pad, + } + + @torch.no_grad() + def decompress(self, compressed: dict) -> torch.Tensor: + """Decompress back to (B, H, S, D) tensor.""" + B, H, S, D = compressed["shape"] + N = B * H * S + idx_bytes = compressed["idx_bytes"].reshape(N, -1) + vec_norms = compressed["vec_norms"].reshape(N, 1).float() + idx_pad = compressed["idx_pad"] + + # Unpack indices + indices_per_byte = 8 // self.bits + mask = (1 << self.bits) - 1 + idx_shifts = torch.tensor( + [self.bits * i for i in range(indices_per_byte - 1, -1, -1)], + dtype=torch.long, + device=idx_bytes.device, + ) + indices = ( + (idx_bytes.long().unsqueeze(-1) >> idx_shifts) & mask + ).reshape(N, -1) + if idx_pad: + indices = indices[:, :D] + + # Reconstruct + reconstructed = (self.centroids[indices] @ self.Pi) * vec_norms + return reconstructed.reshape(B, H, S, D) + + +@dataclass +class KvQuantConfig: + """Configuration for KV cache quantization.""" + + enabled: bool = False + key_bits: int = 4 + value_bits: int = 2 + residual_window: int = 128 + protected_layers: int = 4 + protected_bits: int = 8 + + +class KVCompressor: + """Per-layer KV cache compressor with asymmetric key/value bit-widths. + + Each layer gets its own compressor instance with a unique seed, + enabling layer-adaptive precision (protected layers get more bits). + """ + + def __init__( + self, + head_dim: int, + num_layers: int, + config: KvQuantConfig, + seed: int = 42, + device: str = "cpu", + ): + self.head_dim = head_dim + self.config = config + + self.key_compressors: list[MSECompressor] = [] + self.val_compressors: list[MSECompressor] = [] + + for layer_idx in range(num_layers): + is_protected = ( + layer_idx < config.protected_layers + or layer_idx >= (num_layers - config.protected_layers) + ) + effective_key_bits = config.protected_bits if is_protected else config.key_bits + effective_val_bits = config.protected_bits if is_protected else config.value_bits + effective_key_bits = min(effective_key_bits, 8) + effective_val_bits = min(effective_val_bits, 8) + + seed_base = seed + layer_idx * 1000 + self.key_compressors.append( + MSECompressor(head_dim, effective_key_bits, seed=seed_base, device=device) + ) + self.val_compressors.append( + MSECompressor(head_dim, effective_val_bits, seed=seed_base + 500, device=device) + ) + + @torch.no_grad() + def compress_layer( + self, + layer_idx: int, + keys: torch.Tensor, + values: torch.Tensor, + ) -> tuple[dict, dict]: + """Compress keys/values for one layer. + + Args: + keys: (S, H, D) — old tokens to compress + values: (S, H, D) — old tokens to compress + + Returns: + (compressed_k, compressed_v) dicts + """ + # Reshape (S, H, D) -> (1, H, S, D) for MSECompressor + keys_4d = keys.permute(1, 0, 2).unsqueeze(0) + values_4d = values.permute(1, 0, 2).unsqueeze(0) + + compressed_k = self.key_compressors[layer_idx].compress(keys_4d) + compressed_v = self.val_compressors[layer_idx].compress(values_4d) + return compressed_k, compressed_v + + @torch.no_grad() + def decompress_layer( + self, + layer_idx: int, + compressed_k: dict, + compressed_v: dict, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Decompress keys/values for one layer. + + Returns: + keys: (S, H, D), values: (S, H, D) + """ + keys_4d = self.key_compressors[layer_idx].decompress(compressed_k) + values_4d = self.val_compressors[layer_idx].decompress(compressed_v) + + # (1, H, S, D) -> (S, H, D) + keys = keys_4d.squeeze(0).permute(1, 0, 2) + values = values_4d.squeeze(0).permute(1, 0, 2) + return keys, values diff --git a/python/core/turboquant/lloyd_max.py b/python/core/turboquant/lloyd_max.py new file mode 100644 index 0000000..36db0b2 --- /dev/null +++ b/python/core/turboquant/lloyd_max.py @@ -0,0 +1,79 @@ +"""Lloyd-Max optimal scalar quantizer for the Gaussian distribution. + +After rotating a d-dimensional unit vector by a random orthogonal matrix, +each coordinate follows approximately N(0, 1/d) for d >= 64. +We solve the Lloyd-Max conditions to find optimal centroids. +""" + +import math + +import torch + + +def solve_lloyd_max(d: int, bits: int, max_iter: int = 200, tol: float = 1e-10): + """Solve Lloyd-Max optimal quantizer for N(0, 1/d). + + Returns: + centroids: sorted tensor of 2^bits optimal centroids + boundaries: sorted tensor of 2^bits - 1 boundaries + """ + n_levels = 2 ** bits + sigma = 1.0 / math.sqrt(d) + + # Initialize centroids uniformly in [-3.5*sigma, 3.5*sigma] + lo, hi = -3.5 * sigma, 3.5 * sigma + centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)] + + # Gaussian PDF: N(0, sigma^2) + def pdf(x): + return (1.0 / math.sqrt(2 * math.pi * sigma * sigma)) * math.exp( + -x * x / (2 * sigma * sigma) + ) + + for _ in range(max_iter): + # Step 1: Compute boundaries (midpoints) + boundaries = [ + (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1) + ] + + # Step 2: Update centroids as conditional expectations + edges = [lo * 3] + boundaries + [hi * 3] + new_centroids = [] + for i in range(n_levels): + a, b = edges[i], edges[i + 1] + # Numerical integration via sampling (avoids scipy dependency) + n_samples = 2048 + xs = torch.linspace(a, b, n_samples) + pdf_vals = torch.tensor([pdf(x) for x in xs]) + weighted = xs * pdf_vals + if pdf_vals.sum() > 1e-15: + new_centroids.append((weighted.sum() / pdf_vals.sum()).item()) + else: + new_centroids.append(centroids[i]) + + # Check convergence + max_shift = max( + abs(new_centroids[i] - centroids[i]) for i in range(n_levels) + ) + centroids = new_centroids + if max_shift < tol: + break + + boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)] + return ( + torch.tensor(centroids, dtype=torch.float32), + torch.tensor(boundaries, dtype=torch.float32), + ) + + +class LloydMaxCodebook: + """Precomputed Lloyd-Max codebook for a given dimension and bit-width.""" + + def __init__(self, d: int, bits: int): + self.d = d + self.bits = bits + self.n_levels = 2 ** bits + self.centroids, self.boundaries = solve_lloyd_max(d, bits) + + def __repr__(self): + return f"LloydMaxCodebook(d={self.d}, bits={self.bits}, levels={self.n_levels})" diff --git a/python/core/types.py b/python/core/types.py index c26bc4f..5d59046 100644 --- a/python/core/types.py +++ b/python/core/types.py @@ -17,6 +17,18 @@ from .tokenizer import TokenizerAdapter +@dataclass(frozen=True) +class KvQuantConfig: + """Configuration for TurboQuant KV cache compression.""" + + enabled: bool = False + key_bits: int = 4 + value_bits: int = 2 + residual_window: int = 128 + protected_layers: int = 4 + protected_bits: int = 8 + + @dataclass(frozen=True) class GenerateConfig: """User-facing options that control text generation.""" @@ -67,6 +79,7 @@ class RuntimeConfig: # runs this many decode steps (sample_and_prepare no-ops after EOS or when # the per-call max_new_tokens cap is reached). max_new_tokens: int = 256 + kv_quant_config: KvQuantConfig | None = None @dataclass(frozen=True) From dbad9fb7612c1ac717271ca96148b1916f5da62d Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 16 Jun 2026 09:19:52 +0800 Subject: [PATCH 02/17] [Feature] Add TurboQuant --- examples/model/qwen3_14b/cpu_generate.py | 220 ++++++++-- examples/model/qwen3_14b/npu_generate.py | 12 +- .../model/qwen3_14b/runner/cpu_executor.py | 219 +++++++++- .../model/qwen3_14b/runner/npu_executor.py | 260 ++++++++++-- examples/model/qwen3_14b/runner/npu_runner.py | 197 ++++++++- .../qwen3_14b/runner/qwen3_l3_dispatch.py | 122 ++++++ pypto-lib | 2 +- python/cli/main.py | 22 +- python/core/kv_cache.py | 389 +++++++++++++++++- python/core/model_runner.py | 64 ++- python/core/turboquant/__init__.py | 15 + python/core/turboquant/compressor.py | 163 ++++++++ python/core/turboquant/lloyd_max.py | 114 +++++ python/core/types.py | 16 +- python/runtime/worker.py | 8 +- tests/test_turboquant.py | 343 +++++++++++++++ 16 files changed, 2066 insertions(+), 100 deletions(-) create mode 100644 python/core/turboquant/__init__.py create mode 100644 python/core/turboquant/compressor.py create mode 100644 python/core/turboquant/lloyd_max.py create mode 100644 tests/test_turboquant.py diff --git a/examples/model/qwen3_14b/cpu_generate.py b/examples/model/qwen3_14b/cpu_generate.py index a381a77..31c02bf 100644 --- a/examples/model/qwen3_14b/cpu_generate.py +++ b/examples/model/qwen3_14b/cpu_generate.py @@ -6,11 +6,42 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- +"""CPU-only Qwen3-14B generation (unified entry point). + +Runs full-precision (FP) by default through the reference executor +(`CpuModelExecutor`). Pass ``--tq`` to enable TurboQuant KV cache +compression (`CpuTqModelExecutor`), or ``--compare`` to run both and +print side-by-side output. + +Usage: + # Single prompt (FP, default) + python cpu_generate.py \\ + --model-dir /path/to/Qwen3-14B \\ + --prompt "Hello, please introduce yourself" \\ + --max-new-tokens 128 + + # Interactive chat + python cpu_generate.py \\ + --model-dir /path/to/Qwen3-14B \\ + --interactive + + # TurboQuant KV cache compression + python cpu_generate.py \\ + --model-dir /path/to/Qwen3-14B \\ + --prompt "What is machine learning?" --tq + + # Compare TQ vs FP + python cpu_generate.py \\ + --model-dir /path/to/Qwen3-14B \\ + --prompt "What is machine learning?" \\ + --compare +""" from __future__ import annotations import argparse import sys +import time from pathlib import Path @@ -29,49 +60,168 @@ def _bootstrap_package_root() -> None: from python.core import GenerateConfig, LLMEngine, RuntimeConfig from python.core.kv_cache import KvCacheManager -from examples.model.qwen3_14b.runner.cpu_executor import CpuModelExecutor +from python.core.types import KvQuantConfig +from examples.model.qwen3_14b.runner.cpu_executor import CpuTqModelExecutor, CpuModelExecutor def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser( - description="Run CPU-only Qwen3-14B generation with the reference executor." + description="CPU-only Qwen3-14B generation with TurboQuant KV cache compression.", ) - parser.add_argument("--model-dir", required=True, help="Local model directory, e.g. a Hugging Face snapshot.") - parser.add_argument("--prompt", required=True, help="Prompt text.") - parser.add_argument("--model-id", default="qwen3-14b-cpu-ref") + parser.add_argument("--model-dir", required=True, help="Local model directory (HuggingFace safetensors).") + parser.add_argument("--prompt", default=None, help="Prompt text. Omit for interactive mode.") + parser.add_argument("--interactive", action="store_true", help="Enter interactive chat mode.") + parser.add_argument("--model-id", default="qwen3-14b-tq-cpu") parser.add_argument("--max-seq-len", type=int, default=4096) - parser.add_argument("--max-new-tokens", type=int, default=32) + parser.add_argument("--max-new-tokens", type=int, default=256) parser.add_argument("--temperature", type=float, default=0.0) - parser.add_argument("--top-p", type=float, default=1.0) + parser.add_argument("--top-p", type=float, default=0.95) parser.add_argument("--top-k", type=int, default=None) parser.add_argument("--stream", action="store_true", default=False) + parser.add_argument("--tq", action="store_true", + help="Enable TurboQuant KV cache compression.") + parser.add_argument("--compare", action="store_true", + help="Run both TQ and FP, print side-by-side comparison.") + parser.add_argument("--dump-dir", default=None, + help="Directory to save per-layer intermediate tensors for debugging.") + parser.add_argument("--num-layers-override", type=int, default=None, + help="Only run first N transformer layers.") return parser -def main() -> None: - args = build_parser().parse_args() - model_dir = Path(args.model_dir).resolve() - if not model_dir.is_dir(): - raise FileNotFoundError(f"Model directory does not exist: {model_dir}") +def run_single( + engine: LLMEngine, + model_id: str, + prompt: str, + config: GenerateConfig, +) -> None: + """Run a single prompt and print the result.""" + t0 = time.perf_counter() + result = engine.generate_result(model_id, prompt, config) + dt = time.perf_counter() - t0 + n_tok = len(result.token_ids) + print(f"\n{'=' * 60}") + print(f"Response ({n_tok} tokens, {dt:.2f}s, {n_tok / dt:.1f} tok/s):") + print(f"{'=' * 60}") + print(result.text) + print(f"{'=' * 60}") + print(f"token_ids: {result.token_ids}") + print(f"finish_reason: {result.finish_reason}") + + +def run_interactive( + engine: LLMEngine, + model_id: str, + config: GenerateConfig, +) -> None: + """Interactive chat loop.""" + print("\n" + "=" * 60) + print("Qwen3-14B CPU Chat with TurboQuant (type 'quit' to exit)") + print("=" * 60) + + while True: + try: + prompt = input("\nYou: ").strip() + except (EOFError, KeyboardInterrupt): + print("\nBye!") + break + if not prompt or prompt.lower() in ("quit", "exit", "q"): + print("Bye!") + break + + t0 = time.perf_counter() + result = engine.generate_result(model_id, prompt, config) + dt = time.perf_counter() - t0 + n_tok = len(result.token_ids) + + print(f"\nAssistant ({n_tok} tokens, {dt:.2f}s):") + print(result.text) + + +def run_compare( + engine_tq: LLMEngine, + engine_fp: LLMEngine, + model_id: str, + prompt: str, + config: GenerateConfig, +) -> None: + """Run same prompt on both TQ and FP, print comparison.""" + print(f"\nPrompt: {prompt}") + print("-" * 60) + + t0 = time.perf_counter() + result_tq = engine_tq.generate_result(model_id, prompt, config) + dt_tq = time.perf_counter() - t0 + + t0 = time.perf_counter() + result_fp = engine_fp.generate_result(model_id, prompt, config) + dt_fp = time.perf_counter() - t0 + + print(f"\n[TQ] ({len(result_tq.token_ids)} tokens, {dt_tq:.2f}s):") + print(result_tq.text) + print(f"\n[FP] ({len(result_fp.token_ids)} tokens, {dt_fp:.2f}s):") + print(result_fp.text) + # Token overlap + common = len(set(result_tq.token_ids) & set(result_fp.token_ids)) + total = max(len(result_tq.token_ids), len(result_fp.token_ids)) + print(f"\nToken overlap: {common}/{total} ({common / total * 100:.1f}%)") + + +def create_engine( + model_dir: str, + model_id: str, + max_seq_len: int, + max_new_tokens: int, + use_tq: bool = True, + num_layers_override: int | None = None, +) -> LLMEngine: + """Create an LLMEngine with the appropriate executor.""" kv_cache_manager = KvCacheManager() + + if use_tq: + executor = CpuTqModelExecutor(kv_cache_manager, num_layers_override=num_layers_override) + else: + executor = CpuModelExecutor(kv_cache_manager) + + kv_quant_config = None + if use_tq: + kv_quant_config = KvQuantConfig( + enabled=True, + key_bits=4, + value_bits=4, + ) + engine = LLMEngine( kv_cache_manager=kv_cache_manager, - executor=CpuModelExecutor(kv_cache_manager), + executor=executor, ) + engine.init_model( - model_id=args.model_id, - model_dir=str(model_dir), + model_id=model_id, + model_dir=str(Path(model_dir).resolve()), model_format="huggingface", runtime_config=RuntimeConfig( - page_size=64, + page_size=128, max_batch_size=1, - max_seq_len=args.max_seq_len, + max_seq_len=max_seq_len, + max_new_tokens=max_new_tokens, device="cpu", kv_dtype="bfloat16", weight_dtype="float32", + kv_quant_config=kv_quant_config, ), ) + + return engine + + +def main() -> None: + args = build_parser().parse_args() + model_dir = Path(args.model_dir).resolve() + if not model_dir.is_dir(): + raise FileNotFoundError(f"Model directory does not exist: {model_dir}") + config = GenerateConfig( max_new_tokens=args.max_new_tokens, temperature=args.temperature, @@ -79,16 +229,34 @@ def main() -> None: top_k=args.top_k, stream=args.stream, ) - if args.stream: - result = engine.generate(args.model_id, args.prompt, config) - for chunk in result: - print(chunk, end="", flush=True) - print() + + if args.compare: + print("Loading model for TQ ...") + engine_tq = create_engine(args.model_dir, args.model_id, + args.max_seq_len, args.max_new_tokens, use_tq=True, + num_layers_override=args.num_layers_override) + print("Loading model for FP ...") + engine_fp = create_engine(args.model_dir, "qwen3-14b-fp-cpu", + args.max_seq_len, args.max_new_tokens, use_tq=False, + num_layers_override=args.num_layers_override) + + prompt = args.prompt or "Hello, how are you?" + run_compare(engine_tq, engine_fp, args.model_id, prompt, config) + return + + # Single engine mode + use_tq = args.tq + label = "TQ" if use_tq else "FP" + print(f"Loading model ({label}) ...") + engine = create_engine(args.model_dir, args.model_id, + args.max_seq_len, args.max_new_tokens, use_tq=use_tq, + num_layers_override=args.num_layers_override) + print(f"Model loaded ({label}).") + + if args.interactive or args.prompt is None: + run_interactive(engine, args.model_id, config) else: - result = engine.generate_result(args.model_id, args.prompt, config) - print(f"text: {result.text}") - print(f"token_ids: {result.token_ids}") - print(f"finish_reason: {result.finish_reason}") + run_single(engine, args.model_id, args.prompt, config) if __name__ == "__main__": diff --git a/examples/model/qwen3_14b/npu_generate.py b/examples/model/qwen3_14b/npu_generate.py index 526acc4..f98f96c 100644 --- a/examples/model/qwen3_14b/npu_generate.py +++ b/examples/model/qwen3_14b/npu_generate.py @@ -35,7 +35,7 @@ def _bootstrap_package_root() -> None: from python.core.kv_cache import KvCacheManager 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 @@ -372,6 +372,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 @@ -392,6 +398,7 @@ def main() -> None: device_id=args.device_id, 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, @@ -414,7 +421,8 @@ def main() -> None: max_new_tokens=args.max_new_tokens, device="cpu", kv_dtype="bfloat16", - weight_dtype="float32", + weight_dtype="bfloat16", + kv_quant_config=KvQuantConfig(enabled=True) if args.tq_mode else None, ), ) if collector is not None: diff --git a/examples/model/qwen3_14b/runner/cpu_executor.py b/examples/model/qwen3_14b/runner/cpu_executor.py index aa725b7..ebf0769 100644 --- a/examples/model/qwen3_14b/runner/cpu_executor.py +++ b/examples/model/qwen3_14b/runner/cpu_executor.py @@ -40,8 +40,9 @@ class CpuModelExecutor(ModelExecutor): """Reference CPU executor for functional generation and small tests.""" - def __init__(self, kv_cache_manager: KvCacheManager) -> None: + def __init__(self, kv_cache_manager: KvCacheManager, *, num_layers_override: int | None = None) -> None: super().__init__(kv_cache_manager) + self._num_layers_override = num_layers_override def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult: """Run each prompt through all transformer layers on CPU.""" @@ -54,6 +55,8 @@ def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult positions = torch.arange(seq_len, device=model.runtime.device, dtype=torch.long) for layer_idx, layer in enumerate(model.layers): + if self._num_layers_override is not None and layer_idx >= self._num_layers_override: + break hidden = self._layer_prefill( model=model, layer_idx=layer_idx, @@ -77,6 +80,8 @@ def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: position = int(batch.seq_lens[batch_idx].item()) - 1 for layer_idx, layer in enumerate(model.layers): + if self._num_layers_override is not None and layer_idx >= self._num_layers_override: + break hidden = self._layer_decode( model=model, layer_idx=layer_idx, @@ -231,3 +236,215 @@ def _attention_decode( scores = torch.matmul(q_heads, k_rep.transpose(-1, -2)).squeeze(1) / math.sqrt(q.shape[-1]) attn = torch.softmax(scores, dim=-1) return torch.matmul(attn.unsqueeze(1), v_rep).squeeze(1) + +# --------------------------------------------------------------------------- +# TurboQuant variant +# --------------------------------------------------------------------------- + +try: + from python.core.turboquant.compressor import generate_rotation_matrix + from python.core.turboquant.lloyd_max import LloydMaxCodebook +except ImportError: + from python.core.turboquant.compressor import generate_rotation_matrix + from python.core.turboquant.lloyd_max import LloydMaxCodebook + + +HEAD_DIM_TQ = 128 +EPS_TQ = 1e-6 + + +class CpuTqModelExecutor(CpuModelExecutor): + """CPU executor with TurboQuant KV cache compression (no dumps).""" + + def __init__(self, kv_cache_manager, *, num_layers_override=None): + super().__init__(kv_cache_manager, num_layers_override=num_layers_override) + self._codebook = None + self._rot_matrices = [] + + def _ensure_codebook(self, head_dim): + if self._codebook is None: + self._codebook = LloydMaxCodebook(head_dim, bits=4) + + def _ensure_rot_matrices(self, num_layers, head_dim): + if len(self._rot_matrices) == num_layers: + return + self._rot_matrices = [ + generate_rotation_matrix(head_dim, seed=42 + l * 1000).bfloat16() + for l in range(num_layers) + ] + + def _tq_quantize(self, x, rot_matrix): + cb = self._codebook.centroids + bnd = self._codebook.boundaries + l2 = torch.sqrt(x.float().pow(2).sum(dim=-1, keepdim=True) + EPS_TQ) + xr = torch.matmul((x.float() / l2).bfloat16(), rot_matrix).float() + idx = torch.searchsorted(bnd, xr).to(torch.int32).to(torch.float16).to(torch.uint8) + return idx, l2.float() + + def _tq_dequantize(self, qi, sc, rm): + cb = self._codebook.centroids + yh = cb[qi.to(torch.float16).to(torch.int32).long()] + yh = yh.float() / torch.sqrt(yh.float().pow(2).sum(dim=-1, keepdim=True) + EPS_TQ) + return torch.matmul(yh, rm.float().T) * sc + + def run_prefill(self, model, batch): + cfg = model.config + self._ensure_codebook(cfg.head_dim) + self._ensure_rot_matrices(cfg.num_hidden_layers, cfg.head_dim) + last_hidden_rows, logits_rows = [], [] + for bi, alloc in enumerate(batch.kv_allocations): + h = batch.input_embeddings[bi].to(model.runtime.device).float() + seq_len = int(batch.seq_lens[bi].item()) + h = h[:seq_len] + pos = torch.arange(seq_len, device=model.runtime.device, dtype=torch.long) + for li, layer in enumerate(model.layers): + if self._num_layers_override is not None and li >= self._num_layers_override: + break + h = self._layer_prefill_tq(model, li, layer, h, pos, alloc) + last_hidden_rows.append(h[-1]) + logits_rows.append(self._project_logits(model, h[-1])) + return PrefillResult(last_hidden=torch.stack(last_hidden_rows), logits=torch.stack(logits_rows)) + + def run_decode(self, model, batch): + hidden_rows, logits_rows = [], [] + for bi, alloc in enumerate(batch.kv_allocations): + h = batch.hidden_states[bi].to(model.runtime.device).float() + pos = int(batch.seq_lens[bi].item()) - 1 + for li, layer in enumerate(model.layers): + if self._num_layers_override is not None and li >= self._num_layers_override: + break + h = self._layer_decode_tq(model, li, layer, h, pos, alloc) + hidden_rows.append(h) + logits_rows.append(self._project_logits(model, h)) + return DecodeResult(hidden_states=torch.stack(hidden_rows), logits=torch.stack(logits_rows)) + + def _layer_prefill_tq(self, model, layer_idx, layer, hidden_states, positions, alloc): + cfg = model.config + nkh, nh, hd = cfg.num_key_value_heads, cfg.num_attention_heads, cfg.head_dim + rm = self._rot_matrices[layer_idx] + sl = hidden_states.shape[0] + normed = self._rms_norm(hidden_states, layer.input_rms_weight, cfg.rms_norm_eps) + q = self._linear(normed, layer.wq).view(-1, nh, hd) + k = self._linear(normed, layer.wk).view(-1, nkh, hd) + v = self._linear(normed, layer.wv).view(-1, nkh, hd) + q = self._per_head_rms_norm(q, layer.q_norm_weight, cfg.rms_norm_eps) + k = self._per_head_rms_norm(k, layer.k_norm_weight, cfg.rms_norm_eps) + q = self._apply_rope(q, positions, cfg.rope_theta) + k = self._apply_rope(k, positions, cfg.rope_theta) + qk_all = torch.zeros(sl, nkh, hd, dtype=torch.uint8) + qv_all = torch.zeros(sl, nkh, hd, dtype=torch.uint8) + ks_all = torch.zeros(sl, nkh, 1) + vs_all = torch.zeros(sl, nkh, 1) + for hx in range(nkh): + qi, ks = self._tq_quantize(k[:, hx, :].reshape(-1, hd), rm) + qk_all[:, hx, :] = qi; ks_all[:, hx, :] = ks + qi, vs = self._tq_quantize(v[:, hx, :].reshape(-1, hd), rm) + qv_all[:, hx, :] = qi; vs_all[:, hx, :] = vs + pool = self._kv_cache_manager._pool(model.config.model_id) + if pool.quant_key_indices is not None: + self._write_quant_cache(pool, layer_idx, alloc, 0, qk_all, qv_all, ks_all, vs_all) + ao = self._attention_prefill_tq(q, qk_all, qv_all, ks_all, vs_all, nh, nkh, rm) + ar = hidden_states + self._linear(ao.reshape(sl, -1), layer.wo) + mn = self._rms_norm(ar, layer.post_rms_weight, cfg.rms_norm_eps) + g, u = self._linear(mn, layer.w_gate), self._linear(mn, layer.w_up) + return ar + self._linear(torch.nn.functional.silu(g) * u, layer.w_down) + + def _layer_decode_tq(self, model, layer_idx, layer, hidden_state, position, alloc): + cfg = model.config + nkh, nh, hd = cfg.num_key_value_heads, cfg.num_attention_heads, cfg.head_dim + rm = self._rot_matrices[layer_idx] + normed = self._rms_norm(hidden_state.unsqueeze(0), layer.input_rms_weight, cfg.rms_norm_eps) + q = self._linear(normed, layer.wq).view(nh, hd) + k = self._linear(normed, layer.wk).view(nkh, hd) + v = self._linear(normed, layer.wv).view(nkh, hd) + q = self._per_head_rms_norm(q.unsqueeze(0), layer.q_norm_weight, cfg.rms_norm_eps).squeeze(0) + k = self._per_head_rms_norm(k.unsqueeze(0), layer.k_norm_weight, cfg.rms_norm_eps).squeeze(0) + pt = torch.tensor([position], dtype=torch.long) + q = self._apply_rope(q.unsqueeze(0), pt, cfg.rope_theta).squeeze(0) + k = self._apply_rope(k.unsqueeze(0), pt, cfg.rope_theta).squeeze(0) + nqk, nqv = torch.zeros(nkh, hd, dtype=torch.uint8), torch.zeros(nkh, hd, dtype=torch.uint8) + nks, nvs = torch.zeros(nkh, 1), torch.zeros(nkh, 1) + for hx in range(nkh): + qi, ks = self._tq_quantize(k[hx:hx+1], rm); nqk[hx]=qi[0]; nks[hx]=ks[0] + qi, vs = self._tq_quantize(v[hx:hx+1], rm); nqv[hx]=qi[0]; nvs[hx]=vs[0] + pool = self._kv_cache_manager._pool(model.config.model_id) + if pool.quant_key_indices is not None: + self._write_quant_cache_single(pool, layer_idx, alloc, position, nqk, nqv, nks, nvs) + qkf, qvf, ksf, vsf = self._read_full_quant_cache(pool, layer_idx, alloc, nkh, hd) + ao = self._attention_decode_tq(q, qkf, qvf, ksf, vsf, nh, nkh, rm) + ar = hidden_state + self._linear(ao.reshape(1, -1), layer.wo).squeeze(0) + mn = self._rms_norm(ar.unsqueeze(0), layer.post_rms_weight, cfg.rms_norm_eps) + g, u = self._linear(mn, layer.w_gate), self._linear(mn, layer.w_up) + return ar + self._linear(torch.nn.functional.silu(g) * u, layer.w_down).squeeze(0) + + def _attention_prefill_tq(self, q, qk, qv, ks, vs, nh, nkh, rm): + sl, qpk = q.shape[0], nh // nkh + ctx = torch.zeros(sl, nh, HEAD_DIM_TQ, dtype=torch.float32) + for kvh in range(nkh): + kd = self._tq_dequantize(qk[:, kvh, :], ks[:, kvh, :], rm) + vd = self._tq_dequantize(qv[:, kvh, :], vs[:, kvh, :], rm) + qs = kvh * qpk + qh = q[:, qs:qs + qpk, :] + for i in range(sl): + sc = torch.matmul(qh[i], kd[:i+1].T) / (HEAD_DIM_TQ ** 0.5) + ctx[i, qs:qs+qpk, :] = torch.matmul(torch.softmax(sc, dim=-1), vd[:i+1]) + return ctx + + def _attention_decode_tq(self, q, qk, qv, ks, vs, nh, nkh, rm): + qpk = nh // nkh + parts = [] + for kvh in range(nkh): + kd = self._tq_dequantize(qk[:, kvh, :], ks[:, kvh, :], rm) + vd = self._tq_dequantize(qv[:, kvh, :], vs[:, kvh, :], rm) + qh = q[kvh * qpk:(kvh + 1) * qpk] + sc = torch.matmul(qh, kd.T) / (HEAD_DIM_TQ ** 0.5) + parts.append(torch.matmul(torch.softmax(sc, dim=-1), vd)) + return torch.cat(parts, dim=0) + + @staticmethod + def _write_quant_cache(pool, layer_idx, alloc, start_token, qk, qv, ks, vs): + sl, nkh, hd = qk.shape + ps = pool.page_size + for to in range(sl): + ti = start_token + to + pg = ti // ps; off = ti % ps + if pg >= len(alloc.page_ids): break + pp = alloc.page_ids[pg] + for hx in range(nkh): + pool.quant_key_indices[layer_idx, pp, hx, off, :] = qk[to, hx, :] + pool.quant_val_indices[layer_idx, pp, hx, off, :] = qv[to, hx, :] + pool.quant_key_norms[layer_idx, pp, hx, off, 0] = ks[to, hx, 0].to(torch.bfloat16) + pool.quant_val_norms[layer_idx, pp, hx, off, 0] = vs[to, hx, 0].to(torch.bfloat16) + alloc.tokens_used = max(alloc.tokens_used, start_token + sl) + + @staticmethod + def _write_quant_cache_single(pool, layer_idx, alloc, pos, qk, qv, ks, vs): + nkh, hd = qk.shape + ps = pool.page_size + pg, off = pos // ps, pos % ps + if pg >= len(alloc.page_ids): return + pp = alloc.page_ids[pg] + for hx in range(nkh): + pool.quant_key_indices[layer_idx, pp, hx, off, :] = qk[hx] + pool.quant_val_indices[layer_idx, pp, hx, off, :] = qv[hx] + pool.quant_key_norms[layer_idx, pp, hx, off, 0] = ks[hx, 0].to(torch.bfloat16) + pool.quant_val_norms[layer_idx, pp, hx, off, 0] = vs[hx, 0].to(torch.bfloat16) + alloc.tokens_used = max(alloc.tokens_used, pos + 1) + + @staticmethod + def _read_full_quant_cache(pool, layer_idx, alloc, nkh, hd): + tu, ps = alloc.tokens_used, pool.page_size + qk = torch.zeros(tu, nkh, hd, dtype=torch.uint8) + qv = torch.zeros(tu, nkh, hd, dtype=torch.uint8) + ks = torch.zeros(tu, nkh, 1) + vs = torch.zeros(tu, nkh, 1) + for tok in range(tu): + pg, off = tok // ps, tok % ps + if pg >= len(alloc.page_ids): break + pp = alloc.page_ids[pg] + for hx in range(nkh): + qk[tok, hx, :] = pool.quant_key_indices[layer_idx, pp, hx, off, :] + qv[tok, hx, :] = pool.quant_val_indices[layer_idx, pp, hx, off, :] + ks[tok, hx, 0] = pool.quant_key_norms[layer_idx, pp, hx, off, 0].float() + vs[tok, hx, 0] = pool.quant_val_norms[layer_idx, pp, hx, off, 0].float() + return qk, qv, ks, vs diff --git a/examples/model/qwen3_14b/runner/npu_executor.py b/examples/model/qwen3_14b/runner/npu_executor.py index c9b0ba5..0d81ad7 100644 --- a/examples/model/qwen3_14b/runner/npu_executor.py +++ b/examples/model/qwen3_14b/runner/npu_executor.py @@ -107,6 +107,7 @@ def __init__( device_id: int = 0, save_kernels_dir: str | None = None, l3_trace: bool = False, + tq_mode: bool = False, ) -> None: super().__init__( kv_cache_manager, @@ -115,6 +116,7 @@ def __init__( save_kernels_dir=save_kernels_dir, ) self._l3_trace = l3_trace + self._tq_mode = tq_mode @property def profile_verbose(self) -> bool: @@ -127,6 +129,7 @@ def _create_runner(self, model_id: str, compiled: object) -> ModelRunner: raise TypeError("Qwen314BPyptoExecutor requires Qwen3-14B compiled kernels.") return Qwen314BModelRunner( compiled=compiled, + tq_mode=self._tq_mode, ) def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: @@ -140,15 +143,23 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: def _mark(label: str) -> None: timer.mark(label) - qwen3_prefill_fwd = _load_pypto_lib_qwen14b_module("prefill_fwd") # The fused all-layer decode lives in decode_layer.decode_fwd (the # standalone decode_fwd.py module was removed in pypto-lib). It is now # PAGED: it consumes block_table + slot_mapping and reads/writes the SAME # device-resident paged KV pool prefill writes (self._kv_caches), so no # contiguous bridge / MAX_SEQ env is needed. - qwen3_decode_layer = _load_pypto_lib_qwen14b_module("decode_layer") - qwen3_l3_dispatch.prefill_fwd = qwen3_prefill_fwd.prefill_fwd - qwen3_l3_dispatch.decode_fwd = qwen3_decode_layer.decode_fwd + if self._tq_mode: + # TurboQuant: load the quantized prefill/decode kernels and bind + # them to the TQ host wrappers in qwen3_l3_dispatch. + qwen3_prefill_fwd = _load_pypto_lib_qwen14b_module("prefill_tq") + qwen3_decode_layer = _load_pypto_lib_qwen14b_module("decode_tq") + qwen3_l3_dispatch.prefill_fwd_tq = qwen3_prefill_fwd.prefill_fwd_tq + qwen3_l3_dispatch.decode_fwd_tq = qwen3_decode_layer.decode_fwd_tq + else: + qwen3_prefill_fwd = _load_pypto_lib_qwen14b_module("prefill_fwd") + qwen3_decode_layer = _load_pypto_lib_qwen14b_module("decode_layer") + qwen3_l3_dispatch.prefill_fwd = qwen3_prefill_fwd.prefill_fwd + qwen3_l3_dispatch.decode_fwd = qwen3_decode_layer.decode_fwd _mark("imports") @@ -181,36 +192,68 @@ def _mark(label: str) -> None: ) page_size = model.runtime.page_size max_blocks_per_seq = (model.runtime.max_seq_len + page_size - 1) // page_size - prefill = self._compile_prefill_fwd_callable( - qwen3_l3_dispatch.qwen3_prefill_host, - batch=kernel_batch, - max_seq=model.runtime.max_seq_len, - hidden_size=model.config.hidden_size, - intermediate_size=model.config.intermediate_size, - num_heads=model.config.num_attention_heads, - num_kv_heads=model.config.num_key_value_heads, - head_dim=model.config.head_dim, - num_layers=model.config.num_hidden_layers, - vocab_size=padded_vocab, - block_table_stride=max_blocks_per_seq, - page_size=page_size, - ) - _mark("compile_prefill") - decode = self._compile_decode_fwd_callable( - qwen3_l3_dispatch.qwen3_decode_host, - batch=kernel_batch, - max_seq=model.runtime.max_seq_len, - block_table_stride=max_blocks_per_seq, - hidden_size=model.config.hidden_size, - intermediate_size=model.config.intermediate_size, - num_heads=model.config.num_attention_heads, - num_kv_heads=model.config.num_key_value_heads, - head_dim=model.config.head_dim, - num_layers=model.config.num_hidden_layers, - vocab_size=padded_vocab, - page_size=page_size, - ) - _mark("compile_decode") + if self._tq_mode: + prefill = self._compile_prefill_fwd_tq_callable( + qwen3_l3_dispatch.qwen3_prefill_tq_host, + batch=kernel_batch, + max_seq=model.runtime.max_seq_len, + hidden_size=model.config.hidden_size, + intermediate_size=model.config.intermediate_size, + num_heads=model.config.num_attention_heads, + num_kv_heads=model.config.num_key_value_heads, + head_dim=model.config.head_dim, + num_layers=model.config.num_hidden_layers, + vocab_size=padded_vocab, + block_table_stride=max_blocks_per_seq, + page_size=page_size, + ) + _mark("compile_prefill_tq") + decode = self._compile_decode_fwd_tq_callable( + qwen3_l3_dispatch.qwen3_decode_tq_host, + batch=kernel_batch, + max_seq=model.runtime.max_seq_len, + block_table_stride=max_blocks_per_seq, + hidden_size=model.config.hidden_size, + intermediate_size=model.config.intermediate_size, + num_heads=model.config.num_attention_heads, + num_kv_heads=model.config.num_key_value_heads, + head_dim=model.config.head_dim, + num_layers=model.config.num_hidden_layers, + vocab_size=padded_vocab, + page_size=page_size, + ) + _mark("compile_decode_tq") + else: + prefill = self._compile_prefill_fwd_callable( + qwen3_l3_dispatch.qwen3_prefill_host, + batch=kernel_batch, + max_seq=model.runtime.max_seq_len, + hidden_size=model.config.hidden_size, + intermediate_size=model.config.intermediate_size, + num_heads=model.config.num_attention_heads, + num_kv_heads=model.config.num_key_value_heads, + head_dim=model.config.head_dim, + num_layers=model.config.num_hidden_layers, + vocab_size=padded_vocab, + block_table_stride=max_blocks_per_seq, + page_size=page_size, + ) + _mark("compile_prefill") + decode = self._compile_decode_fwd_callable( + qwen3_l3_dispatch.qwen3_decode_host, + batch=kernel_batch, + max_seq=model.runtime.max_seq_len, + block_table_stride=max_blocks_per_seq, + hidden_size=model.config.hidden_size, + intermediate_size=model.config.intermediate_size, + num_heads=model.config.num_attention_heads, + num_kv_heads=model.config.num_key_value_heads, + head_dim=model.config.head_dim, + num_layers=model.config.num_hidden_layers, + vocab_size=padded_vocab, + page_size=page_size, + ) + _mark("compile_decode") rope_cos_raw, rope_sin_raw = rope_tables( model.runtime.max_seq_len, model.config.head_dim, @@ -221,6 +264,30 @@ def _mark(label: str) -> None: _mark("rope_tables") + # TurboQuant artifacts: per-layer random orthogonal rotation matrices + # + a Lloyd-Max 4-bit codebook. Built only when tq_mode is enabled. + rot_matrices: torch.Tensor | None = None + tq_codebook: torch.Tensor | None = None + if self._tq_mode: + head_dim = model.config.head_dim + num_layers = model.config.num_hidden_layers + rot_list = [] + for layer_idx in range(num_layers): + gen = torch.Generator(device="cpu").manual_seed(42 + layer_idx * 1000) + q_rot, r_rot = torch.linalg.qr(torch.randn(head_dim, head_dim, generator=gen)) + diag_sign = torch.sign(torch.diag(r_rot)) + diag_sign[diag_sign == 0] = 1.0 + rot_list.append(q_rot * diag_sign.unsqueeze(0)) + rot_matrices = self._shared_tensor( + torch.cat(rot_list, dim=0).to(torch.bfloat16).contiguous().cpu() + ) + tq_mod = _load_pypto_lib_qwen14b_module("turboquant_kv") + centroids, _ = tq_mod.solve_lloyd_max(head_dim, bits=4) + tq_codebook = self._shared_tensor( + centroids.float().unsqueeze(0).contiguous().cpu() # [1, 16] + ) + _mark("tq_artifacts") + lm_head_weight = model.lm_head if padded_vocab != lm_head_weight.shape[0]: pad_rows = padded_vocab - lm_head_weight.shape[0] @@ -303,6 +370,9 @@ def _mark(label: str) -> None: decode_block_table_buffer=decode_block_table_buffer, decode_slot_mapping_buffer=decode_slot_mapping_buffer, decode_logits_buffer=decode_logits_buffer, + tq_mode=self._tq_mode, + rot_matrices=rot_matrices, + tq_codebook=tq_codebook, ) def _compile_prefill_fwd_callable( @@ -415,6 +485,128 @@ def _compile_decode_fwd_callable( ] return self._compile_jit_fwd_callable("decode_fwd", jit_fn, dummy_args) + def _compile_prefill_fwd_tq_callable( + self, + jit_fn: object, + *, + batch: int, + max_seq: int, + block_table_stride: int, + hidden_size: int, + intermediate_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + num_layers: int, + vocab_size: int, + page_size: int, + ) -> _L3Callable: + """Compile the TurboQuant ``prefill_fwd_tq`` HOST wrapper into an L3 callable. + + The TQ prefill signature (26 args) differs from non-TQ (24 args): + chunk_lens/chunk_offsets and the BF16 k_cache/v_cache are replaced by + UINT8 quant_k/quant_v caches, FP32 per-row scales, the BF16 rotation + matrices, and the FP32 Lloyd-Max codebook. + """ + kv_hidden = num_kv_heads * head_dim + total_tokens = batch * max_seq + runtime_cache_blocks = (max_seq + page_size - 1) // page_size + cache_rows = batch * runtime_cache_blocks * num_layers * num_kv_heads * page_size + rot_rows = num_layers * head_dim + dummy_args = [ + # args 0-1 + torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16), # hidden_states + torch.empty((batch,), dtype=torch.int32), # seq_lens + # args 2-7: weights (NO chunk_lens/chunk_offsets in TQ) + torch.empty((num_layers, hidden_size), dtype=torch.float32), # input_rms_weight + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wq + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wk + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wv + torch.empty((num_layers, head_dim), dtype=torch.float32), # q_norm_weight + torch.empty((num_layers, head_dim), dtype=torch.float32), # k_norm_weight + # args 8-11: rope + addressing + torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_cos + torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_sin + torch.empty((batch * block_table_stride,), dtype=torch.int32), # block_table + torch.empty((total_tokens,), dtype=torch.int32), # slot_mapping + # args 12-17: TQ quantized caches + scales + rotation + codebook + torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_k_cache + torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_v_cache + torch.empty((cache_rows, 1), dtype=torch.float32), # quant_k_scales + torch.empty((cache_rows, 1), dtype=torch.float32), # quant_v_scales + torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), # rot_matrices + torch.empty((1, 16), dtype=torch.float32), # tq_codebook + # args 18-25: standard weights + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wo + torch.empty((num_layers, hidden_size), dtype=torch.float32), # post_rms_weight + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_gate + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_up + torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), # w_down + torch.empty((1, hidden_size), dtype=torch.float32), # final_norm_weight + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), # lm_head_weight + torch.empty((batch, vocab_size), dtype=torch.float32), # out + ] + return self._compile_jit_fwd_callable("prefill_fwd_tq", jit_fn, dummy_args) + + def _compile_decode_fwd_tq_callable( + self, + jit_fn: object, + *, + batch: int, + max_seq: int, + block_table_stride: int, + hidden_size: int, + intermediate_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + num_layers: int, + vocab_size: int, + page_size: int, + ) -> _L3Callable: + """Compile the TurboQuant ``decode_fwd_tq`` HOST wrapper into an L3 callable. + + TQ decode replaces the BF16 k_cache/v_cache with the same six TQ tensors + (quant caches + scales + rotation + codebook). Weight order is + wo, post_rms, w_gate, w_up, w_down. Total: 26 args (vs 22 for non-TQ). + """ + kv_hidden = num_kv_heads * head_dim + runtime_cache_blocks = (max_seq + page_size - 1) // page_size + cache_rows = num_layers * batch * runtime_cache_blocks * num_kv_heads * page_size + rot_rows = num_layers * head_dim + dummy_args = [ + # args 0-11: same as non-TQ decode + torch.empty((batch, hidden_size), dtype=torch.bfloat16), # hidden_states + torch.empty((num_layers, hidden_size), dtype=torch.float32), # input_rms_weight + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wq + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wk + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wv + torch.empty((num_layers, head_dim), dtype=torch.float32), # q_norm_weight + torch.empty((num_layers, head_dim), dtype=torch.float32), # k_norm_weight + torch.empty((batch,), dtype=torch.int32), # seq_lens + torch.empty((batch * block_table_stride,), dtype=torch.int32), # block_table + torch.empty((batch,), dtype=torch.int32), # slot_mapping + torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_cos + torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_sin + # args 12-17: TQ quantized caches + scales + rotation + codebook + torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_k_cache + torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_v_cache + torch.empty((cache_rows, 1), dtype=torch.float32), # quant_k_scales + torch.empty((cache_rows, 1), dtype=torch.float32), # quant_v_scales + torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), # rot_matrices + torch.empty((1, 16), dtype=torch.float32), # tq_codebook + # args 18-25: weights (TQ order: wo, post_rms, w_gate, w_up, w_down) + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wo + torch.empty((num_layers, hidden_size), dtype=torch.float32), # post_rms_weight + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_gate + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_up + torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), # w_down + torch.empty((1, hidden_size), dtype=torch.float32), # final_norm_weight + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), # lm_head_weight + torch.empty((batch, vocab_size), dtype=torch.float32), # out + ] + return self._compile_jit_fwd_callable("decode_fwd_tq", jit_fn, dummy_args) + def _compile_jit_fwd_callable( self, name: str, diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index 56dc597..3c801a2 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -93,6 +93,10 @@ class _CompiledKernels: decode_block_table_buffer: torch.Tensor decode_slot_mapping_buffer: torch.Tensor decode_logits_buffer: torch.Tensor + # TurboQuant (TQ) artifacts. Populated only when tq_mode=True. + tq_mode: bool = False + rot_matrices: torch.Tensor | None = None # [num_layers*head_dim, head_dim] BF16 + tq_codebook: torch.Tensor | None = None # [1, 16] FP32 Lloyd-Max centroids @dataclass @@ -147,6 +151,9 @@ class _StaticKernelArgs: rope_sin: _StaticDeviceTensor padded_lm_head_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): @@ -156,9 +163,11 @@ def __init__( self, *, compiled: _CompiledKernels, + tq_mode: bool = False, ) -> None: super().__init__() self._compiled = compiled + 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 @@ -234,7 +243,7 @@ def _share_static_kernel_tensors(self) -> None: def _iter_static_host_tensors(self) -> tuple[torch.Tensor, ...]: """Return host tensors that must be shared before the worker forks.""" compiled = self._compiled - return ( + tensors: tuple[torch.Tensor, ...] = ( compiled.final_norm_weight, compiled.rope_cos, compiled.rope_sin, @@ -253,6 +262,12 @@ def _iter_static_host_tensors(self) -> tuple[torch.Tensor, ...]: compiled.decode_slot_mapping_buffer, compiled.decode_logits_buffer, ) + # TurboQuant rotation matrices + codebook are also static kernel inputs. + if compiled.rot_matrices is not None: + tensors += (compiled.rot_matrices,) + if compiled.tq_codebook is not None: + tensors += (compiled.tq_codebook,) + return tensors def _build_static_kernel_args(self) -> _StaticKernelArgs: """Create static device-upload markers once per runner.""" @@ -266,6 +281,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: @@ -282,14 +307,32 @@ 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) - - self._run_distributed_program( - compiled.prefill, - *self._prefill_kernel_args(prefill_inputs, k_cache, v_cache, logits_padded), - ) + if self._tq_mode: + quant_k = kv_cache.quant_k_pages + quant_v = kv_cache.quant_v_pages + k_scales = kv_cache.k_scales_pages + v_scales = kv_cache.v_scales_pages + if quant_k is None or quant_v is None or k_scales is None or v_scales is None: + raise RuntimeError("TurboQuant KV caches are not initialized") + self._validate_kv_cache_bounds( + model, prefill_inputs.block_table, prefill_inputs.slot_mapping, quant_k + ) + self._run_distributed_program( + compiled.prefill, + *self._prefill_tq_kernel_args( + prefill_inputs, quant_k, quant_v, k_scales, v_scales, logits_padded + ), + ) + else: + 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 + ) + self._run_distributed_program( + compiled.prefill, + *self._prefill_kernel_args(prefill_inputs, k_cache, v_cache, logits_padded), + ) for batch_idx, alloc in enumerate(batch.kv_allocations): seq_len = int(batch.seq_lens[batch_idx].item()) @@ -322,19 +365,37 @@ 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 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._run_distributed_program( - compiled.decode, - *self._decode_kernel_args(kernel_inputs, k_cache, v_cache), - ) + if self._tq_mode: + quant_k = kv_cache.quant_k_pages + quant_v = kv_cache.quant_v_pages + k_scales = kv_cache.k_scales_pages + v_scales = kv_cache.v_scales_pages + if quant_k is None or quant_v is None or k_scales is None or v_scales is None: + raise RuntimeError("TurboQuant KV caches are not initialized") + # 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, quant_k + ) + self._run_distributed_program( + compiled.decode, + *self._decode_tq_kernel_args(kernel_inputs, quant_k, quant_v, k_scales, v_scales), + ) + else: + k_cache = kv_cache.key_pages + v_cache = kv_cache.value_pages + # 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._run_distributed_program( + compiled.decode, + *self._decode_kernel_args(kernel_inputs, k_cache, v_cache), + ) for batch_idx, alloc in enumerate(batch.kv_allocations): alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item())) return DecodeResult( @@ -415,6 +476,99 @@ def _decode_kernel_args( inputs.logits, ) + def _prefill_tq_kernel_args( + self, + inputs: _PrefillInputs, + quant_k: DeviceTensor, + quant_v: DeviceTensor, + k_scales: DeviceTensor, + v_scales: DeviceTensor, + logits: torch.Tensor, + ) -> tuple[Any, ...]: + """Return arguments in ``qwen3_prefill_tq_host`` signature order. + + TQ prefill drops chunk_lens/chunk_offsets and replaces the BF16 k/v cache + with UINT8 quantized caches + FP32 scales + rotation matrices + codebook. + """ + static = self._require_static_args() + if static.rot_matrices is None or static.tq_codebook is None: + raise RuntimeError("TurboQuant rotation matrices / codebook are not compiled") + weights = static.decode_weights + 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, + quant_k, + quant_v, + k_scales, + v_scales, + 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, + ) + + def _decode_tq_kernel_args( + self, + inputs: _DecodeKernelInputs, + quant_k: DeviceTensor, + quant_v: DeviceTensor, + k_scales: DeviceTensor, + v_scales: DeviceTensor, + ) -> tuple[Any, ...]: + """Return arguments in ``qwen3_decode_tq_host`` signature order. + + TQ decode replaces the BF16 k/v cache with the six TQ tensors. Weight + order is wo, post_rms, w_gate, w_up, w_down (differs from non-TQ decode). + """ + static = self._require_static_args() + if static.rot_matrices is None or static.tq_codebook is None: + raise RuntimeError("TurboQuant rotation matrices / codebook are not compiled") + weights = static.decode_weights + 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, + quant_k, + quant_v, + k_scales, + v_scales, + 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, + ) + def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _DecodeKernelInputs: """Pad active decode rows to the fixed kernel batch. @@ -527,6 +681,11 @@ def _materialize_static_tensors(self) -> None: *static.decode_weights.values(), ): self._coerce_l3_arg(worker, arg) + # TurboQuant rotation matrices + codebook are also uploaded once. + if static.rot_matrices is not None: + self._coerce_l3_arg(worker, static.rot_matrices) + if static.tq_codebook is not None: + self._coerce_l3_arg(worker, static.tq_codebook) @staticmethod def _copy_replicated_rows( diff --git a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py index 1ae905c..ef03fe0 100644 --- a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py +++ b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py @@ -16,6 +16,10 @@ prefill_fwd = None decode_fwd = None +# TurboQuant (TQ) device kernels; assigned by the executor before compilation +# when tq_mode is enabled. The TQ wrappers below forward to these. +prefill_fwd_tq = None +decode_fwd_tq = None @pl.jit.host @@ -122,3 +126,121 @@ def qwen3_decode_host( lm_head_weight, out, ) + + +@pl.jit.host +def qwen3_prefill_tq_host( + hidden_states: pl.Tensor, + seq_lens: pl.Tensor, + input_rms_weight: pl.Tensor, + wq: pl.Tensor, + wk: pl.Tensor, + wv: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + block_table: pl.Tensor, + slot_mapping: pl.Tensor, + quant_k_cache: pl.Tensor, + quant_v_cache: pl.Tensor, + quant_k_scales: pl.Tensor, + quant_v_scales: pl.Tensor, + rot_matrices: pl.Tensor, + tq_codebook: pl.Tensor, + wo: pl.Tensor, + post_rms_weight: pl.Tensor, + w_gate: pl.Tensor, + w_up: pl.Tensor, + w_down: pl.Tensor, + final_norm_weight: pl.Tensor, + lm_head_weight: pl.Tensor, + out: pl.Out[pl.Tensor], +) -> pl.Tensor: + return prefill_fwd_tq( + hidden_states, + seq_lens, + input_rms_weight, + wq, + wk, + wv, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + block_table, + slot_mapping, + quant_k_cache, + quant_v_cache, + quant_k_scales, + quant_v_scales, + rot_matrices, + tq_codebook, + wo, + post_rms_weight, + w_gate, + w_up, + w_down, + final_norm_weight, + lm_head_weight, + out, + ) + + +@pl.jit.host +def qwen3_decode_tq_host( + hidden_states: pl.Tensor, + input_rms_weight: pl.Tensor, + wq: pl.Tensor, + wk: pl.Tensor, + wv: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + seq_lens: pl.Tensor, + block_table: pl.Tensor, + slot_mapping: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + quant_k_cache: pl.Tensor, + quant_v_cache: pl.Tensor, + quant_k_scales: pl.Tensor, + quant_v_scales: pl.Tensor, + rot_matrices: pl.Tensor, + tq_codebook: pl.Tensor, + wo: pl.Tensor, + post_rms_weight: pl.Tensor, + w_gate: pl.Tensor, + w_up: pl.Tensor, + w_down: pl.Tensor, + final_norm_weight: pl.Tensor, + lm_head_weight: pl.Tensor, + out: pl.Out[pl.Tensor], +) -> pl.Tensor: + return decode_fwd_tq( + hidden_states, + input_rms_weight, + wq, + wk, + wv, + q_norm_weight, + k_norm_weight, + seq_lens, + block_table, + slot_mapping, + rope_cos, + rope_sin, + quant_k_cache, + quant_v_cache, + quant_k_scales, + quant_v_scales, + rot_matrices, + tq_codebook, + wo, + post_rms_weight, + w_gate, + w_up, + w_down, + final_norm_weight, + lm_head_weight, + out, + ) diff --git a/pypto-lib b/pypto-lib index 65871d8..2ae5e30 160000 --- a/pypto-lib +++ b/pypto-lib @@ -1 +1 @@ -Subproject commit 65871d8d2e50e6c943e029272249aef933c2bbba +Subproject commit 2ae5e30ae6df8c58dbb71b92375eefd740f7de75 diff --git a/python/cli/main.py b/python/cli/main.py index 32342b9..46da6e6 100644 --- a/python/cli/main.py +++ b/python/cli/main.py @@ -86,6 +86,13 @@ def build_parser() -> argparse.ArgumentParser: action="store_true", help="Show model loading and kernel compilation logs. Startup logs are suppressed by default.", ) + # TurboQuant + parser.add_argument( + "--tq", + action="store_true", + dest="tq_mode", + help="Enable TurboQuant KV cache compression (4-bit quantized K/V cache).", + ) return parser @@ -99,7 +106,7 @@ def build_serving_engine_config(args: argparse.Namespace) -> EngineConfig: from python.core.async_engine import EngineConfig model_dir = str(Path(args.model).resolve()) - executor_kwargs = _build_executor_kwargs() + executor_kwargs = _build_executor_kwargs(args) return EngineConfig( model_id=args.served_model_name or Path(args.model).name, @@ -122,6 +129,14 @@ def _build_runtime_config(args: argparse.Namespace): if kv_dtype == "auto": kv_dtype = args.dtype + kv_quant_config = None + if getattr(args, "tq_mode", False): + try: + from ..core.types import KvQuantConfig + except ImportError: + from python.core.types import KvQuantConfig + kv_quant_config = KvQuantConfig(enabled=True) + return RuntimeConfig( page_size=args.block_size, max_batch_size=args.max_num_seqs, @@ -130,10 +145,11 @@ def _build_runtime_config(args: argparse.Namespace): kv_dtype=kv_dtype, weight_dtype=args.dtype, max_new_tokens=args.max_new_tokens, + kv_quant_config=kv_quant_config, ) -def _build_executor_kwargs() -> dict[str, object]: +def _build_executor_kwargs(args: argparse.Namespace) -> dict[str, object]: executor_kwargs: dict[str, object] = {} pypto_root = os.environ.get("PYPTO_ROOT") save_kernels_dir = os.environ.get("PYPTO_SAVE_KERNELS_DIR") @@ -141,6 +157,8 @@ def _build_executor_kwargs() -> dict[str, object]: executor_kwargs["pypto_root"] = pypto_root if save_kernels_dir: executor_kwargs["save_kernels_dir"] = save_kernels_dir + if getattr(args, "tq_mode", False): + executor_kwargs["tq_mode"] = True return executor_kwargs diff --git a/python/core/kv_cache.py b/python/core/kv_cache.py index 9d0096e..49cfb6f 100644 --- a/python/core/kv_cache.py +++ b/python/core/kv_cache.py @@ -14,7 +14,24 @@ import torch -from .types import KvAllocation, ModelConfig, RuntimeConfig +from .types import KvAllocation, KvQuantConfig, ModelConfig, RuntimeConfig + +# Lazy import to avoid hard dependency when quantization is disabled. +_KVCOMPRESSOR = None + + +def _get_kv_compressor_cls(): + global _KVCOMPRESSOR + if _KVCOMPRESSOR is None: + from .turboquant.compressor import KVCompressor + _KVCOMPRESSOR = KVCompressor + return _KVCOMPRESSOR + + +@dataclass +class _RequestQuantState: + """Per-request tracking for KV quantization.""" + quant_page_count: int = 0 # how many compressed pages stored in quant_* buffers NONE_HASH = hash(("__none__",)) @@ -112,6 +129,14 @@ class _CachePool: max_blocks_per_seq: int key_pages: torch.Tensor value_pages: torch.Tensor + # TurboQuant KV cache quantization fields (None when quantization is disabled). + kv_compressor: object | None = None + kv_quant_config: KvQuantConfig | None = None + quant_key_indices: torch.Tensor | None = None + quant_key_norms: torch.Tensor | None = None + quant_val_indices: torch.Tensor | None = None + quant_val_norms: torch.Tensor | None = None + request_quant_states: dict = field(default_factory=dict) class KvCacheManager: @@ -160,7 +185,32 @@ def register_model(self, model_id: str, config: ModelConfig, runtime: RuntimeCon if model_id in self._pools: return max_blocks_per_seq = math.ceil(runtime.max_seq_len / runtime.page_size) + # Determine number of bf16 pages. When TurboQuant is enabled we only + # need a small pool: one resident segment per batch slot (for the + # residual window) plus enough workspace pages for one full-length + # sequence (used temporarily during decode to hold decompressed old + # tokens). + # + # NOTE: The workspace is shared across batch slots — it is only needed + # during the restore→kernel→compress cycle. For batch=1 this formula + # is exact. For batch>1, the caller must serialize the restore/compress + # per request so that workspace pages are reused across slots, or set + # total_kv_pages explicitly to cover simultaneous workspace demand. + kv_quant_config = runtime.kv_quant_config num_pages = runtime.total_kv_pages + if num_pages is None and kv_quant_config is not None and kv_quant_config.enabled: + # NOTE: The fused NPU decode kernel has the per-layer cache stride + # (num_pages * num_kv_heads * page_size) baked in at compile time. + # Reducing the page count would cause incorrect layer offsets and a + # device hang. Keep the full page count; TurboQuant saves memory by + # compressing old tokens in-place (freeing bf16 pages back to the + # pool for reuse by *other* requests on the same model). + num_pages = runtime.max_batch_size * max_blocks_per_seq + print( + f"[KV-Quant] Using full bf16 pages ({num_pages}) to match " + f"compiled kernel stride. Compression frees pages for reuse.", + flush=True, + ) if num_pages is None: num_pages = runtime.max_batch_size * max_blocks_per_seq self._init_blocks(num_pages, runtime.page_size) @@ -175,6 +225,48 @@ def register_model(self, model_id: str, config: ModelConfig, runtime: RuntimeCon device=runtime.device, ) value_pages = torch.zeros_like(key_pages) + # Initialize KV quantization if configured + kv_compressor = None + quant_key_indices = None + quant_key_norms = None + quant_val_indices = None + quant_val_norms = None + if kv_quant_config is not None and kv_quant_config.enabled: + print(f"[KV-Quant] Creating KVCompressor for model {model_id} ...", flush=True) + KVCompressor = _get_kv_compressor_cls() + kv_compressor = KVCompressor( + head_dim=config.head_dim, + num_layers=config.num_hidden_layers, + config=kv_quant_config, + device=runtime.device, + ) + print(f"[KV-Quant] KVCompressor for model {model_id} created", flush=True) + + # Compressed storage: nibble-packed uint8, two 4-bit indices per byte + # (the NPU kernel packs byte c = idx[c+HALF]<<4 | idx[c]). Each head + # therefore occupies head_dim // 2 bytes, not head_dim. + # Shape: [num_layers, max_seq_pages, num_kv_heads, page_size, head_dim // 2] + quant_key_indices = torch.zeros( + config.num_hidden_layers, max_blocks_per_seq, + config.num_key_value_heads, runtime.page_size, config.head_dim // 2, + dtype=torch.uint8, device=runtime.device, + ) + quant_key_norms = torch.zeros( + config.num_hidden_layers, max_blocks_per_seq, + config.num_key_value_heads, runtime.page_size, 1, + dtype=torch.float32, device=runtime.device, + ) + quant_val_indices = torch.zeros( + config.num_hidden_layers, max_blocks_per_seq, + config.num_key_value_heads, runtime.page_size, config.head_dim // 2, + dtype=torch.uint8, device=runtime.device, + ) + quant_val_norms = torch.zeros( + config.num_hidden_layers, max_blocks_per_seq, + config.num_key_value_heads, runtime.page_size, 1, + dtype=torch.float32, device=runtime.device, + ) + self._pools[model_id] = _CachePool( page_size=runtime.page_size, num_layers=config.num_hidden_layers, @@ -183,6 +275,12 @@ def register_model(self, model_id: str, config: ModelConfig, runtime: RuntimeCon max_blocks_per_seq=max_blocks_per_seq, key_pages=key_pages, value_pages=value_pages, + kv_compressor=kv_compressor, + kv_quant_config=kv_quant_config, + quant_key_indices=quant_key_indices, + quant_key_norms=quant_key_norms, + quant_val_indices=quant_val_indices, + quant_val_norms=quant_val_norms, ) def allocate_for_prompt(self, model_id: str, request_id: str, prompt_len: int) -> KvAllocation: @@ -322,13 +420,15 @@ def ensure_one_more_slot(self, alloc: KvAllocation) -> int: alloc.tokens_capacity = len(alloc.page_ids) * pool.page_size return self.slot_mapping_for_request(alloc, alloc.tokens_used) + def block_table_for_batch(self, allocations: list[KvAllocation]) -> torch.Tensor: + """Return a dense ``[batch, max_blocks]`` block table for requests.""" + return block_table_from_block_ids([alloc.page_ids for alloc in allocations]) + def slot_mapping_for_request(self, alloc: KvAllocation, token_index: int | None = None) -> int: """Return the physical slot index for a request token.""" pool = self._pool(alloc.model_id) logical_index = alloc.tokens_used if token_index is None else token_index - page_idx = logical_index // pool.page_size - offset = logical_index % pool.page_size - return alloc.page_ids[page_idx] * pool.page_size + offset + return slot_mapping_for_decode(alloc.page_ids, logical_index, pool.page_size) def slot_mapping_for_batch(self, allocations: list[KvAllocation]) -> torch.Tensor: """Return current decode slot mappings for a batch.""" @@ -337,6 +437,17 @@ def slot_mapping_for_batch(self, allocations: list[KvAllocation]) -> torch.Tenso dtype=torch.int32, ) + def slot_mapping_for_positions( + self, + alloc: KvAllocation, + num_tokens: int, + *, + max_tokens: int | None = None, + ) -> torch.Tensor: + """Return per-position slot mappings, optionally padded with -1.""" + pool = self._pool(alloc.model_id) + return slot_mapping_for_positions(alloc.page_ids, num_tokens, pool.page_size, max_tokens=max_tokens) + def write_tokens( self, layer_idx: int, @@ -401,15 +512,285 @@ def materialize_single_layer_cache( pool.value_pages[layer_idx].reshape(-1, pool.head_dim), ) + def materialize_full_layer_cache(self, model_id: str) -> tuple[torch.Tensor, torch.Tensor]: + """Return flattened K/V cache views stacked across every model layer. + + Use this API for fused or L3 decode kernels that select layer i via + an arithmetic offset (layer_idx * cache_rows_per_layer) on a single + cache tensor. The pool is already laid out as + ``[num_layers, num_pages, num_kv_heads, page_size, head_dim]`` so the + flat view is zero-copy. + + Returns: + (key_cache_all, value_cache_all) each shaped + [num_layers * num_pages * num_kv_heads * page_size, head_dim]. + """ + pool = self._pool(model_id) + return ( + pool.key_pages.reshape(-1, pool.head_dim), + pool.value_pages.reshape(-1, pool.head_dim), + ) + + def materialize_single_layer_quant_cache(self, model_id: str, layer_idx: int): + """Return flattened quant K/V cache views for one layer. + + Returns (quant_k, quant_v, quant_k_norms, quant_v_norms) each as 2D + tensors suitable for kernel arguments. Returns (None, None, None, None) + if quantization is not enabled. + """ + pool = self._pool(model_id) + if pool.quant_key_indices is None: + return None, None, None, None + rows_per_layer = pool.max_blocks_per_seq * pool.num_kv_heads * pool.page_size + return ( + pool.quant_key_indices[layer_idx].reshape(-1, pool.head_dim // 2), + pool.quant_val_indices[layer_idx].reshape(-1, pool.head_dim // 2), + pool.quant_key_norms[layer_idx].reshape(-1, 1), + pool.quant_val_norms[layer_idx].reshape(-1, 1), + ) + + def materialize_full_layer_quant_cache(self, model_id: str): + """Return flattened quant K/V cache views across all layers. + + Returns (quant_k, quant_v, quant_k_norms, quant_v_norms) each as 2D + tensors. Returns (None, None, None, None) if quantization is not enabled. + """ + pool = self._pool(model_id) + if pool.quant_key_indices is None: + return None, None, None, None + return ( + pool.quant_key_indices.reshape(-1, pool.head_dim // 2), + pool.quant_val_indices.reshape(-1, pool.head_dim // 2), + pool.quant_key_norms.reshape(-1, 1), + pool.quant_val_norms.reshape(-1, 1), + ) + def free(self, alloc: KvAllocation) -> None: """Return an allocation's pages to the model pool.""" self.release_request(alloc.request_id) alloc.page_ids.clear() alloc.tokens_capacity = 0 alloc.tokens_used = 0 + # Clean up per-request quantization state. + self._pool(alloc.model_id).request_quant_states.pop(alloc.request_id, None) + + # ── KV cache quantization ── + + def is_quantization_enabled(self, model_id: str) -> bool: + """Return whether KV cache quantization is enabled for a model.""" + pool = self._pool(model_id) + return pool.kv_compressor is not None + + def get_compressed_block_info(self, model_id: str, batch) -> tuple[torch.Tensor, torch.Tensor]: + """Build cmp_block_table and cmp_num_blocks for the decode_tq kernel. + + The TQ decode kernel reads ALL context blocks from the compressed + (quantized) cache. cmp_block_table maps logical block positions to + physical page IDs in the quant cache (same layout as block_table). + cmp_num_blocks is the number of context blocks per request + (ceil(seq_len / page_size)). + + Returns: + cmp_block_table: [batch * max_blocks_per_seq] INT32 — physical + block IDs of compressed pages for each request. + cmp_num_blocks: [batch] INT32 — number of compressed blocks per + request. + """ + import torch + pool = self._pool(model_id) + max_blocks = pool.max_blocks_per_seq + actual_batch = len(batch.request_ids) if hasattr(batch, 'request_ids') else len(batch.kv_allocations) + cmp_block_table = torch.zeros(actual_batch * max_blocks, dtype=torch.int32) + cmp_num_blocks = torch.zeros(actual_batch, dtype=torch.int32) + + for i, alloc in enumerate(batch.kv_allocations): + # All blocks written by prefill_tq are compressed. + # Physical page IDs come from the allocation's page_ids. + n_pages = len(alloc.page_ids) + ctx_blocks = min(n_pages, max_blocks) + cmp_num_blocks[i] = ctx_blocks + for j in range(ctx_blocks): + cmp_block_table[i * max_blocks + j] = alloc.page_ids[j] + + return cmp_block_table.share_memory_(), cmp_num_blocks.share_memory_() + + def get_quant_state(self, model_id: str, request_id: str) -> _RequestQuantState: + """Get or create per-request quantization state.""" + pool = self._pool(model_id) + if request_id not in pool.request_quant_states: + pool.request_quant_states[request_id] = _RequestQuantState() + return pool.request_quant_states[request_id] + + def _layer_key_bits(self, pool: _CachePool, layer_idx: int) -> int: + """Return the effective key bits for a layer (respects protected layers).""" + cfg = pool.kv_quant_config + protected = cfg.protected_layers + if layer_idx < protected or layer_idx >= pool.num_layers - protected: + return cfg.protected_bits + return cfg.key_bits + + def _layer_val_bits(self, pool: _CachePool, layer_idx: int) -> int: + """Return the effective value bits for a layer (respects protected layers).""" + cfg = pool.kv_quant_config + protected = cfg.protected_layers + if layer_idx < protected or layer_idx >= pool.num_layers - protected: + return cfg.protected_bits + return cfg.value_bits + + def compress_to_quant(self, model_id: str, alloc: KvAllocation, npu_runner=None) -> list[int]: + """Quantize ALL tokens into quant_* buffers and free their BF16 pages. + + DEPRECATED: With prefill_tq writing directly to quant format, this method + is no longer needed for the prefill path. Kept for potential decode-time + residual window migration (not yet implemented). + """ + pool = self._pool(model_id) + if pool.kv_compressor is None or pool.kv_quant_config is None: + return [] + if pool.quant_key_indices is None: + return [] + return [] + + def print_kv_cache_memory(self, model_id: str) -> None: + """Calculate and print KV cache memory usage for a model.""" + pool = self._pool(model_id) + + single_tensor_bytes = pool.key_pages.numel() * pool.key_pages.element_size() + total_kv_bytes = single_tensor_bytes * 2 + + dtype_size = pool.key_pages.element_size() + per_token_bytes = 2 * pool.num_layers * pool.num_kv_heads * pool.head_dim * dtype_size + + total_pages = pool.key_pages.shape[1] + free_pages = self.num_free_blocks + used_pages = total_pages - free_pages + max_seq_len = pool.max_blocks_per_seq * pool.page_size + + def _fmt(b: int) -> str: + if b >= 1024 ** 3: + return f"{b / 1024 ** 3:.2f} GiB" + return f"{b / 1024 ** 2:.2f} MiB" + + print(f"[KV Cache] model_id = {model_id}", flush=True) + print(f"[KV Cache] dtype = {pool.key_pages.dtype}, element_size = {dtype_size} bytes", flush=True) + print(f"[KV Cache] shape per tensor = {list(pool.key_pages.shape)} " + f"[layers={pool.num_layers}, pages={total_pages}, " + f"kv_heads={pool.num_kv_heads}, page_size={pool.page_size}, head_dim={pool.head_dim}]", flush=True) + print(f"[KV Cache] key_pages = {_fmt(single_tensor_bytes)}", flush=True) + print(f"[KV Cache] value_pages = {_fmt(single_tensor_bytes)}", flush=True) + print(f"[KV Cache] total KV cache memory = {_fmt(total_kv_bytes)}", flush=True) + print(f"[KV Cache] per-token memory = {per_token_bytes} bytes ({per_token_bytes / 1024:.2f} KiB)", flush=True) + print(f"[KV Cache] pages: total={total_pages}, free={free_pages}, used={used_pages}", flush=True) + print(f"[KV Cache] max_seq_len per request = {max_seq_len} tokens " + f"(max_blocks_per_seq={pool.max_blocks_per_seq}, page_size={pool.page_size})", flush=True) + + if pool.kv_compressor is not None: + stats = self.quantization_stats(model_id) + quant_bytes = stats.get("quant_buffer_bytes", 0) + total_extra = quant_bytes + print(f"[KV Cache] quant buffers = {_fmt(quant_bytes)}", flush=True) + if total_kv_bytes > 0: + print(f"[KV Cache] total (bf16 + quant) = {_fmt(total_kv_bytes + total_extra)}", flush=True) + pages_per_seq_uncompressed = pool.max_blocks_per_seq + print( + f"[KV Cache] pages per request: pool={total_pages}, " + f"uncompressed would need={pages_per_seq_uncompressed} per request", + flush=True, + ) + + def quantization_stats(self, model_id: str) -> dict: + """Return memory usage stats for compressed KV cache.""" + pool = self._pool(model_id) + if pool.kv_compressor is None: + return {"enabled": False} + quant_bytes = 0 + if pool.quant_key_indices is not None: + quant_bytes += pool.quant_key_indices.numel() * pool.quant_key_indices.element_size() + if pool.quant_key_norms is not None: + quant_bytes += pool.quant_key_norms.numel() * pool.quant_key_norms.element_size() + if pool.quant_val_indices is not None: + quant_bytes += pool.quant_val_indices.numel() * pool.quant_val_indices.element_size() + if pool.quant_val_norms is not None: + quant_bytes += pool.quant_val_norms.numel() * pool.quant_val_norms.element_size() + return { + "enabled": True, + "quant_buffer_bytes": quant_bytes, + "quant_requests": len(pool.request_quant_states), + } def _pool(self, model_id: str) -> _CachePool: """Return the registered cache pool for a model.""" if model_id not in self._pools: raise KeyError(f"Model {model_id} is not registered with the KV cache manager.") return self._pools[model_id] + + +# --------------------------------------------------------------------------- +# Standalone utility functions for block-table / slot-mapping computation. +# These are pure math (no allocation state) and can be used directly by +# the serving worker without a KvCacheManager instance. +# --------------------------------------------------------------------------- + +def block_table_from_block_ids(block_ids_batch: list[list[int]]) -> torch.Tensor: + """Return a dense ``[batch, max_blocks]`` int32 block table from physical block IDs.""" + max_blocks = max((len(bids) for bids in block_ids_batch), default=0) + table = torch.full((len(block_ids_batch), max_blocks), -1, dtype=torch.int32) + for row, block_ids in enumerate(block_ids_batch): + if block_ids: + table[row, : len(block_ids)] = torch.tensor(block_ids, dtype=torch.int32) + return table + + +def slot_mapping_from_block_ids( + block_ids: list[int], + token_indices: range | list[int], + page_size: int, +) -> torch.Tensor: + """Return physical slot IDs for absolute token indices in one request.""" + slots = [] + for token_index in token_indices: + page_idx = token_index // page_size + offset = token_index % page_size + slots.append(block_ids[page_idx] * page_size + offset) + return torch.tensor(slots, dtype=torch.int32) + + +def slot_mapping_from_block_ids_batch( + block_ids_batch: list[list[int]], + token_indices_batch: list[range | list[int]], + page_size: int, +) -> torch.Tensor: + """Return a padded ``[batch, max_tokens]`` int32 slot mapping from block IDs.""" + mappings = [ + slot_mapping_from_block_ids(block_ids, token_indices, page_size) + for block_ids, token_indices in zip(block_ids_batch, token_indices_batch, strict=True) + ] + max_slots = max((mapping.numel() for mapping in mappings), default=0) + table = torch.full((len(mappings), max_slots), -1, dtype=torch.int32) + for row, mapping in enumerate(mappings): + if mapping.numel() > 0: + table[row, : mapping.numel()] = mapping + return table + + +def slot_mapping_for_decode(page_ids: list[int], tokens_used: int, page_size: int) -> int: + """Return the single physical slot for a decode step.""" + page_idx = tokens_used // page_size + offset = tokens_used % page_size + return page_ids[page_idx] * page_size + offset + + +def slot_mapping_for_positions( + page_ids: list[int], + num_tokens: int, + page_size: int, + max_tokens: int | None = None, +) -> torch.Tensor: + """Return per-position slot mappings, optionally padded with -1.""" + size = num_tokens if max_tokens is None else max_tokens + mapping = torch.full((size,), -1, dtype=torch.int32) + for token_index in range(num_tokens): + page_idx = token_index // page_size + offset = token_index % page_size + mapping[token_index] = page_ids[page_idx] * page_size + offset + return mapping diff --git a/python/core/model_runner.py b/python/core/model_runner.py index 28830cc..94789f8 100644 --- a/python/core/model_runner.py +++ b/python/core/model_runner.py @@ -30,10 +30,19 @@ @dataclass class _KvCachePool: - """Worker-resident flat all-layer KV cache for one model.""" + """Worker-resident flat all-layer KV cache for one model. - key_pages: DeviceTensor - value_pages: DeviceTensor + For TurboQuant (TQ) mode the quantized caches and scales are populated + instead of (or in addition to) the BF16 key_pages / value_pages. + """ + + key_pages: DeviceTensor | None = None + value_pages: DeviceTensor | None = None + # TurboQuant compressed caches. + quant_k_pages: DeviceTensor | None = None + quant_v_pages: DeviceTensor | None = None + k_scales_pages: DeviceTensor | None = None + v_scales_pages: DeviceTensor | None = None class ModelRunner(ABC): @@ -52,12 +61,45 @@ def init_kv_cache( num_pages = runtime.total_kv_pages if num_pages is None: num_pages = runtime.max_batch_size * max_blocks_per_seq - kv_dtype = getattr(torch, runtime.kv_dtype) cache_rows = config.num_hidden_layers * num_pages * config.num_key_value_heads * runtime.page_size cache_shape = ( cache_rows, config.head_dim, ) + + qcfg = runtime.kv_quant_config + if qcfg is not None and qcfg.enabled: + # TurboQuant: UINT8 quantized K/V caches + FP32 per-row scales, + # allocated in place of the BF16 key/value pages. + scales_shape = (cache_rows, 1) + quant_k = self._alloc_kv_cache_tensor(cache_shape, torch.uint8) + try: + quant_v = self._alloc_kv_cache_tensor(cache_shape, torch.uint8) + except Exception: + self._free_kv_cache_tensor(quant_k) + raise + try: + k_scales = self._alloc_kv_cache_tensor(scales_shape, torch.float32) + except Exception: + self._free_kv_cache_tensor(quant_k) + self._free_kv_cache_tensor(quant_v) + raise + try: + v_scales = self._alloc_kv_cache_tensor(scales_shape, torch.float32) + except Exception: + self._free_kv_cache_tensor(quant_k) + self._free_kv_cache_tensor(quant_v) + self._free_kv_cache_tensor(k_scales) + raise + self._kv_caches[model_id] = _KvCachePool( + quant_k_pages=quant_k, + quant_v_pages=quant_v, + k_scales_pages=k_scales, + v_scales_pages=v_scales, + ) + return + + kv_dtype = getattr(torch, runtime.kv_dtype) key_pages = self._alloc_kv_cache_tensor(cache_shape, kv_dtype) try: value_pages = self._alloc_kv_cache_tensor(cache_shape, kv_dtype) @@ -72,8 +114,18 @@ def init_kv_cache( def close_kv_cache(self) -> None: """Release all runner-owned KV cache tensors.""" for pool in list(self._kv_caches.values()): - self._free_kv_cache_tensor(pool.key_pages) - self._free_kv_cache_tensor(pool.value_pages) + if pool.key_pages is not None: + self._free_kv_cache_tensor(pool.key_pages) + if pool.value_pages is not None: + self._free_kv_cache_tensor(pool.value_pages) + if pool.quant_k_pages is not None: + self._free_kv_cache_tensor(pool.quant_k_pages) + if pool.quant_v_pages is not None: + self._free_kv_cache_tensor(pool.quant_v_pages) + if pool.k_scales_pages is not None: + self._free_kv_cache_tensor(pool.k_scales_pages) + if pool.v_scales_pages is not None: + self._free_kv_cache_tensor(pool.v_scales_pages) self._kv_caches.clear() @abstractmethod diff --git a/python/core/turboquant/__init__.py b/python/core/turboquant/__init__.py new file mode 100644 index 0000000..17e362c --- /dev/null +++ b/python/core/turboquant/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""KV cache quantization module.""" + +from ..types import KvQuantConfig +from .compressor import KVCompressor, TurboQuantCompressor + +__all__ = ["KVCompressor", "KvQuantConfig", "TurboQuantCompressor"] diff --git a/python/core/turboquant/compressor.py b/python/core/turboquant/compressor.py new file mode 100644 index 0000000..3768849 --- /dev/null +++ b/python/core/turboquant/compressor.py @@ -0,0 +1,163 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""TurboQuant KV cache compressor: rotation matrix generation + bit-packing utilities. + +The actual KV quantization is performed by NPU kernels (prefill_tq / decode_tq) +which write directly to compressed format. This module provides: + - Rotation matrix generation (seeded per-layer) + - Bit-packing / bit-unpacking for sub-8-bit storage + - KVCompressor for layer-adaptive precision configuration and rotation matrix export +""" + +from __future__ import annotations + +import math +import time + +import torch +import torch.nn.functional as F + +from ..types import KvQuantConfig + + +def generate_rotation_matrix(d: int, seed: int = 42, device: str = "cpu") -> torch.Tensor: + """Generate a random orthogonal rotation matrix via QR decomposition.""" + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + G = torch.randn(d, d, generator=gen) + Q, R = torch.linalg.qr(G) + diag_sign = torch.sign(torch.diag(R)) + diag_sign[diag_sign == 0] = 1.0 + return (Q * diag_sign.unsqueeze(0)).to(device) + + +class TurboQuantCompressor: + """Per-vector TurboQuant compressor configuration. + + Stores the rotation matrix and quantization parameters for one layer. + The actual compress/decompress is done by NPU kernels; this class + provides bit-packing utilities for host-side storage. + """ + + def __init__(self, head_dim: int, bits: int = 4, seed: int = 42, device: str = "cpu"): + self.head_dim = head_dim + self.bits = min(bits, 8) + self.device = device + self.n_levels = 2 ** self.bits + + # Rotation matrix (fixed at init). + self.Pi = generate_rotation_matrix(head_dim, seed=seed, device=device) + self.PiT = self.Pi.T.contiguous() + + # Quantization range for normalized vectors (~N(0, 1/d)). + sigma = 1.0 / math.sqrt(head_dim) + self.lo = -3.5 * sigma + self.hi = 3.5 * sigma + + # Uniform centroids (matching NPU kernel's uniform quantization). + self.uniform_centroids = torch.linspace(self.lo, self.hi, self.n_levels, dtype=torch.float32) + + def _bit_pack(self, indices: torch.Tensor, N: int, D: int) -> tuple[torch.Tensor, int]: + """Pack UINT8 indices into bit-packed bytes for sub-8-bit.""" + if self.bits >= 8: + return indices.reshape(N, D), 0 + + indices_per_byte = 8 // self.bits + idx_pad = (indices_per_byte - D % indices_per_byte) % indices_per_byte + idx_flat = indices.long() + if idx_pad: + idx_flat = F.pad(idx_flat, (0, idx_pad)) + n_groups = idx_flat.shape[-1] // indices_per_byte + idx_powers = torch.tensor( + [2 ** (self.bits * i) for i in range(indices_per_byte - 1, -1, -1)], + dtype=torch.long, + device=idx_flat.device, + ) + idx_bytes = ( + (idx_flat.reshape(N, n_groups, indices_per_byte) * idx_powers) + .sum(-1) + .to(torch.uint8) + ) + return idx_bytes, idx_pad + + def _bit_unpack(self, idx_bytes: torch.Tensor, N: int, D: int, idx_pad: int) -> torch.Tensor: + """Unpack bit-packed bytes back to UINT8 indices.""" + if self.bits >= 8: + return idx_bytes.reshape(N, D) + + indices_per_byte = 8 // self.bits + mask = (1 << self.bits) - 1 + idx_shifts = torch.tensor( + [self.bits * i for i in range(indices_per_byte - 1, -1, -1)], + dtype=torch.long, + device=idx_bytes.device, + ) + indices = ( + (idx_bytes.long().unsqueeze(-1) >> idx_shifts) & mask + ).reshape(N, -1) + if idx_pad: + indices = indices[:, :D] + return indices.to(torch.uint8) + + +class KVCompressor: + """Per-layer KV cache compressor configuration. + + Each layer gets its own TurboQuantCompressor with a unique rotation + matrix and configurable bit width, enabling layer-adaptive precision. + The NPU kernels (prefill_tq / decode_tq) handle all compression/decompression. + """ + + def __init__( + self, + head_dim: int, + num_layers: int, + config: KvQuantConfig, + seed: int = 42, + device: str = "cpu", + ): + self.head_dim = head_dim + self.config = config + + self.key_compressors: list[TurboQuantCompressor] = [] + self.val_compressors: list[TurboQuantCompressor] = [] + + print(f"[TQ] Initializing: {num_layers} layers, head_dim={head_dim}, " + f"key_bits={config.key_bits}, val_bits={config.value_bits}", flush=True) + t_total = time.perf_counter() + for layer_idx in range(num_layers): + is_protected = ( + layer_idx < config.protected_layers + or layer_idx >= (num_layers - config.protected_layers) + ) + effective_key_bits = config.protected_bits if is_protected else config.key_bits + effective_val_bits = config.protected_bits if is_protected else config.value_bits + effective_key_bits = min(effective_key_bits, 8) + effective_val_bits = min(effective_val_bits, 8) + + self.key_compressors.append( + TurboQuantCompressor(head_dim, effective_key_bits, + seed=seed + layer_idx * 1000, device=device) + ) + self.val_compressors.append( + TurboQuantCompressor(head_dim, effective_val_bits, + seed=seed + layer_idx * 1000, device=device) + ) + print(f"[TQ] layer {layer_idx}: key_bits={effective_key_bits}, " + f"val_bits={effective_val_bits}", flush=True) + dt_total = (time.perf_counter() - t_total) * 1000 + print(f"[TQ] Initialization complete: {dt_total:.1f} ms", flush=True) + + def get_rot_matrices(self, device: str = "cpu") -> torch.Tensor: + """Stack all per-layer rotation matrices for NPU upload. + + Returns: (num_layers, head_dim, head_dim) BF16 tensor. + """ + return torch.stack([c.Pi for c in self.key_compressors]).bfloat16().to(device) diff --git a/python/core/turboquant/lloyd_max.py b/python/core/turboquant/lloyd_max.py new file mode 100644 index 0000000..2de3530 --- /dev/null +++ b/python/core/turboquant/lloyd_max.py @@ -0,0 +1,114 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Lloyd-Max optimal scalar quantizer for the Gaussian distribution. + +After rotating a d-dimensional unit vector by a random orthogonal matrix, +each coordinate follows approximately N(0, 1/d) for d >= 64. +We solve the Lloyd-Max conditions to find optimal centroids. +""" + +import math +from functools import lru_cache + +import torch + + +def solve_lloyd_max(d: int, bits: int, max_iter: int = 200, tol: float = 1e-10): + """Solve Lloyd-Max optimal quantizer for N(0, 1/d). + + Uses fully vectorized operations: + - Analytical Gaussian CDF via torch.erf for centroid updates + - No per-element Python loops in the hot path + + Returns: + centroids: sorted tensor of 2^bits optimal centroids + boundaries: sorted tensor of 2^bits - 1 boundaries + """ + n_levels = 2 ** bits + sigma = 1.0 / math.sqrt(d) + + # Initialize centroids uniformly in [-3.5*sigma, 3.5*sigma] + lo, hi = -3.5 * sigma, 3.5 * sigma + centroids = torch.linspace(lo, hi, n_levels + 2)[1:-1] # n_levels points, excluding endpoints + + for _ in range(max_iter): + # Step 1: boundaries = midpoints between adjacent centroids + boundaries = (centroids[:-1] + centroids[1:]) / 2.0 + + # Step 2: update centroids via analytical Gaussian conditional mean. + # For N(0, sigma^2) in interval [a, b]: + # E[X | a < X < b] = sigma * (phi(a/sigma) - phi(b/sigma)) + # / (Phi(b/sigma) - Phi(a/sigma)) + # where phi = standard normal PDF, Phi = CDF. + # phi(x) = exp(-x^2/2) / sqrt(2pi), Phi(x) = 0.5*(1+erf(x/sqrt(2))) + edges = torch.zeros(n_levels + 1) + edges[0] = lo * 3 + edges[1:-1] = boundaries + edges[-1] = hi * 3 + + # Standardize edges + z = edges / sigma + + # phi(z) = exp(-z^2/2) / sqrt(2*pi) + neg_half_z2 = -0.5 * z * z + phi_z = neg_half_z2.exp() * (1.0 / math.sqrt(2.0 * math.pi)) + + # Phi(z) = 0.5 * (1 + erf(z / sqrt(2))) + cdf_z = 0.5 * (1.0 + torch.erf(z * (1.0 / math.sqrt(2.0)))) + + # Conditional means: sigma * (phi(z_a) - phi(z_b)) / (Phi(z_b) - Phi(z_a)) + phi_diff = phi_z[:-1] - phi_z[1:] # phi(z_a) - phi(z_b), shape (n_levels,) + cdf_diff = cdf_z[1:] - cdf_z[:-1] # Phi(z_b) - Phi(z_a), shape (n_levels,) + + # Guard against tiny intervals + valid = cdf_diff > 1e-15 + new_centroids = torch.where( + valid, + sigma * phi_diff / cdf_diff.clamp(min=1e-15), + centroids, + ) + + # Check convergence + max_shift = (new_centroids - centroids).abs().max().item() + centroids = new_centroids + if max_shift < tol: + break + + boundaries = (centroids[:-1] + centroids[1:]) / 2.0 + return centroids.clone(), boundaries.clone() + + +@lru_cache(maxsize=None) +def _cached_solve(d: int, bits: int): + """Cache (d, bits) -> codebook to avoid redundant computation.""" + return solve_lloyd_max(d, bits) + + +class LloydMaxCodebook: + """Precomputed Lloyd-Max codebook for a given dimension and bit-width. + + Results are cached by (d, bits) so that multiple layers with the same + head_dim and bit-width share one codebook. + """ + + def __init__(self, d: int, bits: int): + import time + self.d = d + self.bits = bits + self.n_levels = 2 ** bits + t0 = time.perf_counter() + self.centroids, self.boundaries = _cached_solve(d, bits) + dt = (time.perf_counter() - t0) * 1000 + tag = "(cached)" if dt < 0.1 else "" + print(f"[TurboQuant] Lloyd-Max codebook: d={d}, bits={bits}, levels={self.n_levels} " + f"{tag} {dt:.1f} ms", flush=True) + + def __repr__(self): + return f"LloydMaxCodebook(d={self.d}, bits={self.bits}, levels={self.n_levels})" diff --git a/python/core/types.py b/python/core/types.py index bb3343e..271b624 100644 --- a/python/core/types.py +++ b/python/core/types.py @@ -51,6 +51,16 @@ class ModelConfig: torch_dtype: str +@dataclass +class KvQuantConfig: + enabled: bool = False + key_bits: int = 4 + value_bits: int = 4 + residual_window: int = 128 + protected_layers: int = 0 + protected_bits: int = 4 + + @dataclass(frozen=True) class RuntimeConfig: """Runtime limits and device placement for one loaded model.""" @@ -62,7 +72,11 @@ class RuntimeConfig: kv_dtype: str = "float32" weight_dtype: str = "float32" total_kv_pages: int | None = None - # Compile-time generation limit used by model-specific runners. + kv_quant_config: KvQuantConfig | None = None + # Compile-time decode loop length for the host_orch pl.unroll in l3_generate. + # All generate() calls must use max_new_tokens <= this value; the loop always + # runs this many decode steps (sample_and_prepare no-ops after EOS or when + # the per-call max_new_tokens cap is reached). max_new_tokens: int = 256 diff --git a/python/runtime/worker.py b/python/runtime/worker.py index 1f69db8..b7f1b4b 100644 --- a/python/runtime/worker.py +++ b/python/runtime/worker.py @@ -23,7 +23,7 @@ from typing import Any import torch -from simpler.task_interface import ContinuousTensor, DataType +from simpler.task_interface import DataType, Tensor from simpler.worker import Worker as SimplerWorker @@ -119,9 +119,9 @@ def torch_dtype(self) -> torch.dtype: """Return the corresponding ``torch.dtype``.""" return _to_torch_dtype(self.dtype) - def to_continuous_tensor(self) -> ContinuousTensor: + def to_continuous_tensor(self) -> Tensor: """Return a Simpler tensor view that skips runtime malloc/free.""" - return ContinuousTensor.make(self.data_ptr, self.shape, self.dtype, child_memory=True) + return Tensor.make(self.data_ptr, self.shape, self.dtype, True) class Worker: @@ -133,7 +133,7 @@ class Worker: optional Python sub-workers. The wrapper is intentionally small: it adds LLM-oriented validation and ``WorkerTensor`` helpers, but preserves Simpler concepts such as callables, worker ids, ``CallConfig``, and - ``ContinuousTensor``. + ``Tensor`` (child-memory views). """ def __init__( diff --git a/tests/test_turboquant.py b/tests/test_turboquant.py new file mode 100644 index 0000000..6415b22 --- /dev/null +++ b/tests/test_turboquant.py @@ -0,0 +1,343 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- + +"""Host-side unit tests for the TurboQuant (TQ) NPU path. + +These cover the TQ wiring ported onto main's ``qwen3_l3_dispatch`` +architecture without an NPU: + * the quantized KV-cache allocation driven by ``kv_quant_config``; + * the TQ kernel-argument builders (26-arg signature + ordering); + * static upload of the TQ rotation matrices / codebook. + +They mirror ``test_batching.py``'s fake style and import the real runner, so +they need ``pypto``/``simpler`` importable (same as ``test_batching.py``) but no +NPU. The kernel-argument ordering is the highest-value check: a wrong order is +silent until the fused kernel runs on hardware. +""" + +from __future__ import annotations + +import math +from typing import Any + +import pytest +import torch + +from python.core.model_runner import ModelRunner +from python.core.types import KvQuantConfig, ModelConfig, RuntimeConfig +from examples.model.qwen3_14b.runner.npu_runner import ( + Qwen314BModelRunner, + _CompiledKernels, + _DecodeKernelInputs, + _L3Callable, + _PrefillInputs, + _StaticDeviceTensor, +) + + +def _config() -> ModelConfig: + return ModelConfig( + model_id="tq-test", + architecture="qwen3", + vocab_size=16, + hidden_size=8, + intermediate_size=16, + num_hidden_layers=2, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=4, + max_position_embeddings=128, + rms_norm_eps=1e-6, + rope_theta=10000.0, + bos_token_id=None, + eos_token_id=None, + pad_token_id=None, + torch_dtype="float32", + ) + + +def _runtime(*, kv_quant_config: KvQuantConfig | None = None) -> RuntimeConfig: + return RuntimeConfig( + page_size=64, + max_batch_size=1, + max_seq_len=128, + kv_dtype="bfloat16", + kv_quant_config=kv_quant_config, + ) + + +def _decode_weights(hidden_size: int, kv_hidden: int, head_dim: int, intermediate_size: int): + return { + "decode_input_rms_weight": torch.ones(1, hidden_size), + "decode_wq": torch.zeros(hidden_size, hidden_size), + "decode_wk": torch.zeros(hidden_size, kv_hidden), + "decode_wv": torch.zeros(hidden_size, kv_hidden), + "decode_q_norm_weight": torch.ones(1, head_dim), + "decode_k_norm_weight": torch.ones(1, head_dim), + "decode_wo": torch.zeros(hidden_size, hidden_size), + "decode_post_rms_weight": torch.ones(1, hidden_size), + "decode_w_gate": torch.zeros(hidden_size, intermediate_size), + "decode_w_up": torch.zeros(hidden_size, intermediate_size), + "decode_w_down": torch.zeros(intermediate_size, hidden_size), + } + + +def _compiled_kernels(config: ModelConfig, *, tq: bool = False) -> _CompiledKernels: + hidden_size = config.hidden_size + intermediate_size = config.intermediate_size + head_dim = config.head_dim + kv_hidden = config.num_key_value_heads * head_dim + kernel_batch = 1 + max_seq = 128 + max_blocks = math.ceil(max_seq / 64) + callable_ = _L3Callable(compiled=object(), name="fake", block_dim=1, aicpu_thread_num=1) + return _CompiledKernels( + prefill=callable_, + decode=callable_, + final_norm_weight=torch.ones(1, hidden_size), + rope_cos=torch.zeros(max_seq, head_dim), + rope_sin=torch.zeros(max_seq, head_dim), + padded_vocab=config.vocab_size, + padded_lm_head_weight=torch.zeros(config.vocab_size, hidden_size), + decode_weights=_decode_weights(hidden_size, kv_hidden, head_dim, intermediate_size), + prefill_hidden_buffer=torch.empty(kernel_batch * max_seq, hidden_size, dtype=torch.bfloat16), + prefill_seq_lens_buffer=torch.empty(kernel_batch, dtype=torch.int32), + prefill_chunk_lens_buffer=torch.empty(kernel_batch, dtype=torch.int32), + prefill_chunk_offsets_buffer=torch.empty(kernel_batch, dtype=torch.int32), + prefill_block_table_buffer=torch.empty(kernel_batch * max_blocks, dtype=torch.int32), + prefill_slot_mapping_buffer=torch.empty(kernel_batch * max_seq, dtype=torch.int32), + prefill_logits_buffer=torch.empty(kernel_batch, config.vocab_size), + decode_hidden_buffer=torch.zeros(kernel_batch, hidden_size, dtype=torch.bfloat16), + decode_seq_lens_buffer=torch.zeros(kernel_batch, dtype=torch.int32), + decode_block_table_buffer=torch.zeros(kernel_batch * max_blocks, dtype=torch.int32), + decode_slot_mapping_buffer=torch.zeros(kernel_batch, dtype=torch.int32), + decode_logits_buffer=torch.zeros(kernel_batch, config.vocab_size), + tq_mode=tq, + rot_matrices=( + torch.zeros(config.num_hidden_layers * head_dim, head_dim, dtype=torch.bfloat16) + if tq + else None + ), + tq_codebook=torch.zeros(1, 16, dtype=torch.float32) if tq else None, + ) + + +class _RecordingRunner(ModelRunner): + """Minimal concrete ModelRunner that records KV-cache allocations.""" + + def __init__(self) -> None: + super().__init__() + self.allocs: list[tuple[tuple[int, ...], torch.dtype]] = [] + + def _alloc_kv_cache_tensor(self, shape, dtype) -> Any: + record = (tuple(shape), dtype) + self.allocs.append(record) + return record + + def _free_kv_cache_tensor(self, tensor) -> None: # noqa: ARG002 + return None + + def run_prefill(self, model, batch): # noqa: ARG002 + raise NotImplementedError + + def run_decode(self, model, batch): # noqa: ARG002 + raise NotImplementedError + + +def _sentinel() -> Any: + """A unique object usable as a stand-in for a DeviceTensor/tensor arg.""" + return object() + + +def test_init_kv_cache_allocates_quant_pool_when_kv_quant_config_enabled(): + config = _config() + runtime = _runtime(kv_quant_config=KvQuantConfig(enabled=True)) + runner = _RecordingRunner() + runner.init_kv_cache("m", config, runtime) + pool = runner._kv_caches["m"] + + # BF16 key/value pages are NOT allocated in TQ mode. + assert pool.key_pages is None + assert pool.value_pages is None + # UINT8 quantized K/V caches + FP32 per-row scales are. + assert pool.quant_k_pages is not None + assert pool.quant_v_pages is not None + assert pool.k_scales_pages is not None + assert pool.v_scales_pages is not None + + num_pages = math.ceil(runtime.max_seq_len / runtime.page_size) + cache_rows = config.num_hidden_layers * num_pages * config.num_key_value_heads * runtime.page_size + # Quant indices are nibble-packed: 2x4-bit per byte -> head_dim // 2 wide. + cache_shape = (cache_rows, config.head_dim // 2) + scales_shape = (cache_rows, 1) + assert runner.allocs == [ + (cache_shape, torch.uint8), + (cache_shape, torch.uint8), + (scales_shape, torch.float32), + (scales_shape, torch.float32), + ] + + +def test_init_kv_cache_allocates_bf16_pool_by_default(): + config = _config() + runtime = _runtime() # kv_quant_config=None -> standard path + runner = _RecordingRunner() + runner.init_kv_cache("m", config, runtime) + pool = runner._kv_caches["m"] + + assert pool.key_pages is not None + assert pool.value_pages is not None + assert pool.quant_k_pages is None + assert pool.quant_v_pages is None + assert pool.k_scales_pages is None + assert pool.v_scales_pages is None + + num_pages = math.ceil(runtime.max_seq_len / runtime.page_size) + cache_rows = config.num_hidden_layers * num_pages * config.num_key_value_heads * runtime.page_size + cache_shape = (cache_rows, config.head_dim) + assert runner.allocs == [ + (cache_shape, torch.bfloat16), + (cache_shape, torch.bfloat16), + ] + + +def _tq_runner(config: ModelConfig) -> Qwen314BModelRunner: + return Qwen314BModelRunner(compiled=_compiled_kernels(config, tq=True), tq_mode=True) + + +def _prefill_inputs(hidden_size: int) -> _PrefillInputs: + return _PrefillInputs( + actual_batch=1, + hidden=torch.zeros(1, hidden_size, dtype=torch.bfloat16), + seq_lens=torch.zeros(1, dtype=torch.int32), + chunk_lens=torch.zeros(1, dtype=torch.int32), + chunk_offsets=torch.zeros(1, dtype=torch.int32), + block_table=torch.zeros(2, dtype=torch.int32), + slot_mapping=torch.zeros(1, dtype=torch.int32), + ) + + +def test_prefill_tq_kernel_args_match_host_wrapper_signature(): + config = _config() + runner = _tq_runner(config) + static = runner._require_static_args() + inputs = _prefill_inputs(config.hidden_size) + quant_k, quant_v = _sentinel(), _sentinel() + k_scales, v_scales = _sentinel(), _sentinel() + logits = _sentinel() + + args = runner._prefill_tq_kernel_args(inputs, quant_k, quant_v, k_scales, v_scales, logits) + + assert len(args) == 26 + # args 0-1: hidden, seq_lens (TQ prefill drops chunk_lens/chunk_offsets). + assert args[0] is inputs.hidden + assert args[1] is inputs.seq_lens + # args 8-11: rope + addressing. + assert args[8] is static.rope_cos + assert args[9] is static.rope_sin + assert args[10] is inputs.block_table + assert args[11] is inputs.slot_mapping + # args 12-17: TQ quant caches + scales + rotation + codebook. + assert args[12] is quant_k + assert args[13] is quant_v + assert args[14] is k_scales + assert args[15] is v_scales + assert args[16] is static.rot_matrices + assert args[17] is static.tq_codebook + # args 18-22: weight order wo, post_rms, w_gate, w_up, w_down. + assert args[18] is static.decode_weights["decode_wo"] + assert args[19] is static.decode_weights["decode_post_rms_weight"] + assert args[20] is static.decode_weights["decode_w_gate"] + assert args[21] is static.decode_weights["decode_w_up"] + assert args[22] is static.decode_weights["decode_w_down"] + # tail: final_norm, lm_head, out. + assert args[23] is static.final_norm_weight + assert args[24] is static.padded_lm_head_weight + assert args[25] is logits + + +def test_decode_tq_kernel_args_match_host_wrapper_signature(): + config = _config() + runner = _tq_runner(config) + static = runner._require_static_args() + inputs = _DecodeKernelInputs( + actual_batch=1, + hidden=torch.zeros(1, config.hidden_size, dtype=torch.bfloat16), + seq_lens=torch.zeros(1, dtype=torch.int32), + block_table=torch.zeros(2, dtype=torch.int32), + slot_mapping=torch.zeros(1, dtype=torch.int32), + logits=_sentinel(), + ) + quant_k, quant_v = _sentinel(), _sentinel() + k_scales, v_scales = _sentinel(), _sentinel() + + args = runner._decode_tq_kernel_args(inputs, quant_k, quant_v, k_scales, v_scales) + + assert len(args) == 26 + assert args[0] is inputs.hidden + # args 7-9: seq_lens, block_table, slot_mapping. + assert args[7] is inputs.seq_lens + assert args[8] is inputs.block_table + assert args[9] is inputs.slot_mapping + # args 12-17: TQ tensors. + assert args[12] is quant_k + assert args[13] is quant_v + assert args[14] is k_scales + assert args[15] is v_scales + assert args[16] is static.rot_matrices + assert args[17] is static.tq_codebook + # args 18-22: TQ decode weight order is wo, post_rms, w_gate, w_up, w_down + # (differs from non-TQ decode: wo, w_gate, w_up, w_down, post_rms). + assert args[18] is static.decode_weights["decode_wo"] + assert args[19] is static.decode_weights["decode_post_rms_weight"] + assert args[20] is static.decode_weights["decode_w_gate"] + assert args[21] is static.decode_weights["decode_w_up"] + assert args[22] is static.decode_weights["decode_w_down"] + assert args[25] is inputs.logits + + +def test_tq_static_tensors_include_rotation_and_codebook(): + config = _config() + runner = _tq_runner(config) + static = runner._require_static_args() + + assert isinstance(static.rot_matrices, _StaticDeviceTensor) + assert isinstance(static.tq_codebook, _StaticDeviceTensor) + # They must be shared before the worker forks. Use identity (tensors do + # not support boolean `in`). + shared = runner._iter_static_host_tensors() + assert any(runner._compiled.rot_matrices is t for t in shared) + assert any(runner._compiled.tq_codebook is t for t in shared) + + +def test_non_tq_compiled_kernels_have_no_tq_statics(): + config = _config() + tq_runner = _tq_runner(config) + fp_runner = Qwen314BModelRunner(compiled=_compiled_kernels(config, tq=False), tq_mode=False) + + fp_static = fp_runner._require_static_args() + assert fp_static.rot_matrices is None + assert fp_static.tq_codebook is None + + # TQ mode uploads exactly two extra static tensors (rot + codebook) on top + # of the same shared set the standard path uploads. + tq_shared = tq_runner._iter_static_host_tensors() + fp_shared = fp_runner._iter_static_host_tensors() + assert len(tq_shared) == len(fp_shared) + 2 + + +def test_tq_kernel_args_reject_missing_rotation_artifacts(): + config = _config() + # tq_mode runner built from a non-TQ compiled kernel set (no rot/codebook). + runner = Qwen314BModelRunner(compiled=_compiled_kernels(config, tq=False), tq_mode=True) + inputs = _prefill_inputs(config.hidden_size) + with pytest.raises(RuntimeError, match="rotation matrices"): + runner._prefill_tq_kernel_args( + inputs, _sentinel(), _sentinel(), _sentinel(), _sentinel(), _sentinel() + ) From be33939c52b03a9c8a750ccddc193e4c80b8b609 Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 12:08:43 +0800 Subject: [PATCH 03/17] Merge hw-native-sys/main: L3 serving + device sampling + TurboQuant --- .../model/qwen3_14b/runner/npu_executor.py | 403 ++++++++---- examples/model/qwen3_14b/runner/npu_runner.py | 604 ++++++++++++------ .../qwen3_14b/runner/qwen3_l3_dispatch.py | 117 ++++ 3 files changed, 806 insertions(+), 318 deletions(-) diff --git a/examples/model/qwen3_14b/runner/npu_executor.py b/examples/model/qwen3_14b/runner/npu_executor.py index 0d81ad7..c615ad1 100644 --- a/examples/model/qwen3_14b/runner/npu_executor.py +++ b/examples/model/qwen3_14b/runner/npu_executor.py @@ -11,6 +11,7 @@ import importlib.util import sys +from collections.abc import Sequence from dataclasses import dataclass from pathlib import Path from typing import Any @@ -104,7 +105,7 @@ def __init__( kv_cache_manager=None, *, platform: str = "a2a3sim", - device_id: int = 0, + device_ids: Sequence[int] = (0,), save_kernels_dir: str | None = None, l3_trace: bool = False, tq_mode: bool = False, @@ -112,7 +113,7 @@ def __init__( super().__init__( kv_cache_manager, platform=platform, - device_id=device_id, + device_ids=device_ids, save_kernels_dir=save_kernels_dir, ) self._l3_trace = l3_trace @@ -123,12 +124,23 @@ def profile_verbose(self) -> bool: """Return whether compile and L3 execution timing logs are enabled.""" return self._l3_trace + @property + def supports_device_sampling(self) -> bool: + """Qwen3 NPU runner can return greedy sampled token ids.""" + return True + + @property + def supports_device_embedding(self) -> bool: + """Qwen3 NPU decode embeds greedy token ids inside the device kernel.""" + return True + def _create_runner(self, model_id: str, compiled: object) -> ModelRunner: """Create the Qwen3-14B runtime runner for compiled kernels.""" if not isinstance(compiled, _CompiledKernels): raise TypeError("Qwen314BPyptoExecutor requires Qwen3-14B compiled kernels.") return Qwen314BModelRunner( compiled=compiled, + device_id=self._device_ids[0], tq_mode=self._tq_mode, ) @@ -143,28 +155,38 @@ def _compile_model(self, model: RuntimeModel) -> _CompiledKernels: def _mark(label: str) -> None: timer.mark(label) - # The fused all-layer decode lives in decode_layer.decode_fwd (the - # standalone decode_fwd.py module was removed in pypto-lib). It is now - # PAGED: it consumes block_table + slot_mapping and reads/writes the SAME - # device-resident paged KV pool prefill writes (self._kv_caches), so no - # contiguous bridge / MAX_SEQ env is needed. if self._tq_mode: - # TurboQuant: load the quantized prefill/decode kernels and bind - # them to the TQ host wrappers in qwen3_l3_dispatch. + # TurboQuant: load quantized prefill/decode kernels. qwen3_prefill_fwd = _load_pypto_lib_qwen14b_module("prefill_tq") qwen3_decode_layer = _load_pypto_lib_qwen14b_module("decode_tq") qwen3_l3_dispatch.prefill_fwd_tq = qwen3_prefill_fwd.prefill_fwd_tq qwen3_l3_dispatch.decode_fwd_tq = qwen3_decode_layer.decode_fwd_tq else: qwen3_prefill_fwd = _load_pypto_lib_qwen14b_module("prefill_fwd") + # The fused all-layer decode lives in decode_layer.decode_fwd. qwen3_decode_layer = _load_pypto_lib_qwen14b_module("decode_layer") + qwen3_greedy_sample = _load_pypto_lib_qwen14b_module("greedy_sample") + qwen3_token_embed = _load_pypto_lib_qwen14b_module("token_embed") qwen3_l3_dispatch.prefill_fwd = qwen3_prefill_fwd.prefill_fwd qwen3_l3_dispatch.decode_fwd = qwen3_decode_layer.decode_fwd + qwen3_l3_dispatch.greedy_sample_fwd = qwen3_greedy_sample.greedy_sample_fwd + qwen3_l3_dispatch.token_embed_fwd = qwen3_token_embed.token_embed_fwd _mark("imports") self._validate_supported_shape(model) kernel_batch = model.runtime.max_batch_size + + kernel_max_seq = int(getattr(qwen3_decode_layer, "MAX_SEQ", 4096)) + if model.runtime.max_seq_len > kernel_max_seq: + raise ValueError( + f"max_model_len {model.runtime.max_seq_len} exceeds the kernel's " + f"compile-time MAX_SEQ {kernel_max_seq} (config.py). The decode/prefill " + "kernels precompute MAX_CTX_BLOCKS, NUM_PAGES, and rope table sizes from " + "MAX_SEQ; a larger runtime value silently produces wrong attention and " + "out-of-bounds rope reads. Rebuild the kernel with a larger MAX_SEQ." + ) + if int(qwen3_decode_layer.BATCH) != kernel_batch: raise ValueError( "decode_layer.decode_fwd is compiled for a fixed kernel BATCH of " @@ -190,6 +212,42 @@ def _mark(label: str) -> None: f"{padded_vocab} (round_up({model.config.vocab_size}, {_VOCAB_PAD_MULTIPLE})); " "they must match for the decode logits buffer / lm_head weight to line up." ) + if model.config.vocab_size != int(qwen3_decode_layer.REAL_VOCAB): + raise ValueError( + "decode_layer.decode_fwd hard-codes REAL_VOCAB for padded-token masking, " + f"but the runtime model vocab_size is {model.config.vocab_size}; expected " + f"{int(qwen3_decode_layer.REAL_VOCAB)}." + ) + if int(qwen3_greedy_sample.BATCH) != kernel_batch: + raise ValueError( + "greedy_sample_fwd is compiled for a fixed kernel BATCH of " + f"{int(qwen3_greedy_sample.BATCH)}, but runtime max_batch_size is {kernel_batch}." + ) + if int(qwen3_greedy_sample.VOCAB) != padded_vocab: + raise ValueError( + "greedy_sample_fwd VOCAB must match the padded logits vocab: " + f"{int(qwen3_greedy_sample.VOCAB)} != {padded_vocab}." + ) + if int(qwen3_token_embed.BATCH) != kernel_batch: + raise ValueError( + "token_embed_fwd is compiled for a fixed kernel BATCH of " + f"{int(qwen3_token_embed.BATCH)}, but runtime max_batch_size is {kernel_batch}." + ) + if int(qwen3_token_embed.VOCAB) != padded_vocab: + raise ValueError( + "token_embed_fwd VOCAB must match the padded embedding vocab: " + f"{int(qwen3_token_embed.VOCAB)} != {padded_vocab}." + ) + if int(qwen3_token_embed.HIDDEN) != model.config.hidden_size: + raise ValueError( + "token_embed_fwd HIDDEN must match model hidden_size: " + f"{int(qwen3_token_embed.HIDDEN)} != {model.config.hidden_size}." + ) + sampled_ids_width = int( + getattr(qwen3_decode_layer, "SAMPLED_IDS_PAD", getattr(qwen3_greedy_sample, "SAMPLED_IDS_PAD", 1)) + ) + else: + sampled_ids_width = 1 page_size = model.runtime.page_size max_blocks_per_seq = (model.runtime.max_seq_len + page_size - 1) // page_size if self._tq_mode: @@ -237,6 +295,7 @@ def _mark(label: str) -> None: vocab_size=padded_vocab, block_table_stride=max_blocks_per_seq, page_size=page_size, + sampled_ids_width=sampled_ids_width, ) _mark("compile_prefill") decode = self._compile_decode_fwd_callable( @@ -252,8 +311,24 @@ def _mark(label: str) -> None: num_layers=model.config.num_hidden_layers, vocab_size=padded_vocab, page_size=page_size, + sampled_ids_width=sampled_ids_width, ) _mark("compile_decode") + greedy_sample = self._compile_greedy_sample_callable( + qwen3_l3_dispatch.qwen3_greedy_sample_host, + batch=kernel_batch, + sampled_ids_width=sampled_ids_width, + vocab_size=padded_vocab, + ) + _mark("compile_greedy_sample") + token_embed = self._compile_token_embed_callable( + qwen3_l3_dispatch.qwen3_token_embed_host, + batch=kernel_batch, + hidden_size=model.config.hidden_size, + sampled_ids_width=sampled_ids_width, + vocab_size=padded_vocab, + ) + _mark("compile_token_embed") rope_cos_raw, rope_sin_raw = rope_tables( model.runtime.max_seq_len, model.config.head_dim, @@ -266,39 +341,42 @@ def _mark(label: str) -> None: # TurboQuant artifacts: per-layer random orthogonal rotation matrices # + a Lloyd-Max 4-bit codebook. Built only when tq_mode is enabled. - rot_matrices: torch.Tensor | None = None - tq_codebook: torch.Tensor | None = None + rot_matrices = None + tq_codebook = None if self._tq_mode: - head_dim = model.config.head_dim - num_layers = model.config.num_hidden_layers + hd_tq = model.config.head_dim + nl_tq = model.config.num_hidden_layers rot_list = [] - for layer_idx in range(num_layers): - gen = torch.Generator(device="cpu").manual_seed(42 + layer_idx * 1000) - q_rot, r_rot = torch.linalg.qr(torch.randn(head_dim, head_dim, generator=gen)) - diag_sign = torch.sign(torch.diag(r_rot)) - diag_sign[diag_sign == 0] = 1.0 - rot_list.append(q_rot * diag_sign.unsqueeze(0)) - rot_matrices = self._shared_tensor( - torch.cat(rot_list, dim=0).to(torch.bfloat16).contiguous().cpu() - ) + for _ in range(nl_tq): + rot = torch.randn(hd_tq, hd_tq) + q, _ = torch.linalg.qr(rot) + rot_list.append(q.to(torch.bfloat16)) + rot_matrices = self._shared_tensor(torch.cat(rot_list, dim=0).contiguous().cpu()) tq_mod = _load_pypto_lib_qwen14b_module("turboquant_kv") - centroids, _ = tq_mod.solve_lloyd_max(head_dim, bits=4) + centroids, _ = tq_mod.solve_lloyd_max(hd_tq, bits=4) tq_codebook = self._shared_tensor( - centroids.float().unsqueeze(0).contiguous().cpu() # [1, 16] + torch.tensor(centroids, dtype=torch.float32).view(1, -1).cpu() ) _mark("tq_artifacts") lm_head_weight = model.lm_head if padded_vocab != lm_head_weight.shape[0]: pad_rows = padded_vocab - lm_head_weight.shape[0] - padding = torch.zeros( - (pad_rows, lm_head_weight.shape[1]), - dtype=lm_head_weight.dtype, - device=lm_head_weight.device, - ) + padding = lm_head_weight[:1].expand(pad_rows, -1).clone() lm_head_weight = torch.cat([lm_head_weight, padding], dim=0) padded_lm_head_weight = self._shared_tensor(lm_head_weight.to(torch.bfloat16).contiguous().cpu()) _mark("pad_lm_head") + embed_weight = model.embed_tokens + if padded_vocab != embed_weight.shape[0]: + pad_rows = padded_vocab - embed_weight.shape[0] + padding = torch.zeros( + (pad_rows, embed_weight.shape[1]), + dtype=embed_weight.dtype, + device=embed_weight.device, + ) + embed_weight = torch.cat([embed_weight, padding], dim=0) + padded_embed_weight = self._shared_tensor(embed_weight.to(torch.bfloat16).contiguous().cpu()) + _mark("pad_embed") layers = [] for layer in model.layers: layers.append(self._kernel_layer_weights(layer)) @@ -330,6 +408,14 @@ def _mark(label: str) -> None: (kernel_batch, padded_vocab), dtype=torch.float32, ).share_memory_() + prefill_sampled_ids_buffer = torch.empty( + (kernel_batch, sampled_ids_width), + dtype=torch.int32, + ).share_memory_() + prefill_next_hidden_buffer = torch.empty( + (kernel_batch, model.config.hidden_size), + dtype=torch.bfloat16, + ).share_memory_() _mark("prefill_buffers") decode_logits_buffer = torch.empty( (kernel_batch, padded_vocab), @@ -345,18 +431,36 @@ def _mark(label: str) -> None: dtype=torch.int32, ).share_memory_() decode_slot_mapping_buffer = torch.empty((kernel_batch,), dtype=torch.int32).share_memory_() + decode_token_ids_buffer = torch.empty( + (kernel_batch, sampled_ids_width), + dtype=torch.int32, + ).share_memory_() + decode_sampled_ids_buffer = torch.empty( + (kernel_batch, sampled_ids_width), + dtype=torch.int32, + ).share_memory_() + decode_next_hidden_buffer = torch.empty( + (kernel_batch, model.config.hidden_size), + dtype=torch.bfloat16, + ).share_memory_() _mark("decode_logits_buffer") timer.report() + if self._tq_mode: + greedy_sample = None + token_embed = None return _CompiledKernels( prefill=prefill, decode=decode, + greedy_sample=greedy_sample, + token_embed=token_embed, final_norm_weight=final_norm_weight, rope_cos=rope_cos, rope_sin=rope_sin, padded_vocab=padded_vocab, padded_lm_head_weight=padded_lm_head_weight, + padded_embed_weight=padded_embed_weight, decode_weights=decode_weights, prefill_hidden_buffer=prefill_hidden_buffer, prefill_seq_lens_buffer=prefill_seq_lens_buffer, @@ -365,11 +469,16 @@ def _mark(label: str) -> None: prefill_block_table_buffer=prefill_block_table_buffer, prefill_slot_mapping_buffer=prefill_slot_mapping_buffer, prefill_logits_buffer=prefill_logits_buffer, + prefill_sampled_ids_buffer=prefill_sampled_ids_buffer, + prefill_next_hidden_buffer=prefill_next_hidden_buffer, decode_hidden_buffer=decode_hidden_buffer, decode_seq_lens_buffer=decode_seq_lens_buffer, decode_block_table_buffer=decode_block_table_buffer, decode_slot_mapping_buffer=decode_slot_mapping_buffer, decode_logits_buffer=decode_logits_buffer, + decode_token_ids_buffer=decode_token_ids_buffer, + decode_sampled_ids_buffer=decode_sampled_ids_buffer, + decode_next_hidden_buffer=decode_next_hidden_buffer, tq_mode=self._tq_mode, rot_matrices=rot_matrices, tq_codebook=tq_codebook, @@ -390,6 +499,7 @@ def _compile_prefill_fwd_callable( num_layers: int, vocab_size: int, page_size: int, + sampled_ids_width: int, ) -> _L3Callable: """Compile the prefill HOST wrapper into a distributed program.""" kv_hidden = num_kv_heads * head_dim @@ -414,10 +524,10 @@ def _compile_prefill_fwd_callable( torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), - torch.empty((num_layers, hidden_size), dtype=torch.float32), torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers, hidden_size), dtype=torch.float32), torch.empty((1, hidden_size), dtype=torch.float32), torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), torch.empty((batch, vocab_size), dtype=torch.float32), @@ -439,6 +549,7 @@ def _compile_decode_fwd_callable( num_layers: int, vocab_size: int, page_size: int, + sampled_ids_width: int, ) -> _L3Callable: """Compile the fused all-layer PAGED decode HOST wrapper into a distributed program. @@ -482,130 +593,44 @@ def _compile_decode_fwd_callable( torch.empty((1, hidden_size), dtype=torch.float32), # final_norm_weight torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), # lm_head_weight torch.empty((batch, vocab_size), dtype=torch.float32), # out + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), # embed_weight + torch.empty((batch, sampled_ids_width), dtype=torch.int32), # sampled_ids_in + torch.empty((batch, sampled_ids_width), dtype=torch.int32), # sampled_ids_out + torch.empty((batch, hidden_size), dtype=torch.bfloat16), # next_hidden ] return self._compile_jit_fwd_callable("decode_fwd", jit_fn, dummy_args) - def _compile_prefill_fwd_tq_callable( + def _compile_greedy_sample_callable( self, jit_fn: object, *, batch: int, - max_seq: int, - block_table_stride: int, - hidden_size: int, - intermediate_size: int, - num_heads: int, - num_kv_heads: int, - head_dim: int, - num_layers: int, + sampled_ids_width: int, vocab_size: int, - page_size: int, ) -> _L3Callable: - """Compile the TurboQuant ``prefill_fwd_tq`` HOST wrapper into an L3 callable. - - The TQ prefill signature (26 args) differs from non-TQ (24 args): - chunk_lens/chunk_offsets and the BF16 k_cache/v_cache are replaced by - UINT8 quant_k/quant_v caches, FP32 per-row scales, the BF16 rotation - matrices, and the FP32 Lloyd-Max codebook. - """ - kv_hidden = num_kv_heads * head_dim - total_tokens = batch * max_seq - runtime_cache_blocks = (max_seq + page_size - 1) // page_size - cache_rows = batch * runtime_cache_blocks * num_layers * num_kv_heads * page_size - rot_rows = num_layers * head_dim + """Compile the greedy sampling HOST wrapper.""" dummy_args = [ - # args 0-1 - torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16), # hidden_states - torch.empty((batch,), dtype=torch.int32), # seq_lens - # args 2-7: weights (NO chunk_lens/chunk_offsets in TQ) - torch.empty((num_layers, hidden_size), dtype=torch.float32), # input_rms_weight - torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wq - torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wk - torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wv - torch.empty((num_layers, head_dim), dtype=torch.float32), # q_norm_weight - torch.empty((num_layers, head_dim), dtype=torch.float32), # k_norm_weight - # args 8-11: rope + addressing - torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_cos - torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_sin - torch.empty((batch * block_table_stride,), dtype=torch.int32), # block_table - torch.empty((total_tokens,), dtype=torch.int32), # slot_mapping - # args 12-17: TQ quantized caches + scales + rotation + codebook - torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_k_cache - torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_v_cache - torch.empty((cache_rows, 1), dtype=torch.float32), # quant_k_scales - torch.empty((cache_rows, 1), dtype=torch.float32), # quant_v_scales - torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), # rot_matrices - torch.empty((1, 16), dtype=torch.float32), # tq_codebook - # args 18-25: standard weights - torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wo - torch.empty((num_layers, hidden_size), dtype=torch.float32), # post_rms_weight - torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_gate - torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_up - torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), # w_down - torch.empty((1, hidden_size), dtype=torch.float32), # final_norm_weight - torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), # lm_head_weight - torch.empty((batch, vocab_size), dtype=torch.float32), # out + torch.empty((batch, vocab_size), dtype=torch.float32), + torch.empty((batch, sampled_ids_width), dtype=torch.int32), ] - return self._compile_jit_fwd_callable("prefill_fwd_tq", jit_fn, dummy_args) + return self._compile_jit_fwd_callable("greedy_sample_fwd", jit_fn, dummy_args) - def _compile_decode_fwd_tq_callable( + def _compile_token_embed_callable( self, jit_fn: object, *, batch: int, - max_seq: int, - block_table_stride: int, hidden_size: int, - intermediate_size: int, - num_heads: int, - num_kv_heads: int, - head_dim: int, - num_layers: int, + sampled_ids_width: int, vocab_size: int, - page_size: int, ) -> _L3Callable: - """Compile the TurboQuant ``decode_fwd_tq`` HOST wrapper into an L3 callable. - - TQ decode replaces the BF16 k_cache/v_cache with the same six TQ tensors - (quant caches + scales + rotation + codebook). Weight order is - wo, post_rms, w_gate, w_up, w_down. Total: 26 args (vs 22 for non-TQ). - """ - kv_hidden = num_kv_heads * head_dim - runtime_cache_blocks = (max_seq + page_size - 1) // page_size - cache_rows = num_layers * batch * runtime_cache_blocks * num_kv_heads * page_size - rot_rows = num_layers * head_dim + """Compile the embedding lookup HOST wrapper.""" dummy_args = [ - # args 0-11: same as non-TQ decode - torch.empty((batch, hidden_size), dtype=torch.bfloat16), # hidden_states - torch.empty((num_layers, hidden_size), dtype=torch.float32), # input_rms_weight - torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wq - torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wk - torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wv - torch.empty((num_layers, head_dim), dtype=torch.float32), # q_norm_weight - torch.empty((num_layers, head_dim), dtype=torch.float32), # k_norm_weight - torch.empty((batch,), dtype=torch.int32), # seq_lens - torch.empty((batch * block_table_stride,), dtype=torch.int32), # block_table - torch.empty((batch,), dtype=torch.int32), # slot_mapping - torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_cos - torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_sin - # args 12-17: TQ quantized caches + scales + rotation + codebook - torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_k_cache - torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_v_cache - torch.empty((cache_rows, 1), dtype=torch.float32), # quant_k_scales - torch.empty((cache_rows, 1), dtype=torch.float32), # quant_v_scales - torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), # rot_matrices - torch.empty((1, 16), dtype=torch.float32), # tq_codebook - # args 18-25: weights (TQ order: wo, post_rms, w_gate, w_up, w_down) - torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wo - torch.empty((num_layers, hidden_size), dtype=torch.float32), # post_rms_weight - torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_gate - torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_up - torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), # w_down - torch.empty((1, hidden_size), dtype=torch.float32), # final_norm_weight - torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), # lm_head_weight - torch.empty((batch, vocab_size), dtype=torch.float32), # out + torch.empty((batch, sampled_ids_width), dtype=torch.int32), + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), + torch.empty((batch, hidden_size), dtype=torch.bfloat16), ] - return self._compile_jit_fwd_callable("decode_fwd_tq", jit_fn, dummy_args) + return self._compile_jit_fwd_callable("token_embed_fwd", jit_fn, dummy_args) def _compile_jit_fwd_callable( self, @@ -620,7 +645,7 @@ def _compile_jit_fwd_callable( config = self._run_config(codegen_only=True) distributed_config = DistributedConfig( - device_ids=[self._device_id], + device_ids=list(self._device_ids), num_sub_workers=0, block_dim=_QWEN14B_BLOCK_DIM, aicpu_thread_num=4, @@ -691,14 +716,13 @@ def cat(attr: str) -> torch.Tensor: @classmethod def _validate_total_kv_pages(cls, model: RuntimeModel, kernel_batch: int) -> None: - """Validate that runtime KV page count matches compiled capacity.""" + """Validate that runtime KV page count covers the batch capacity.""" if model.runtime.total_kv_pages is None: return - expected_pages = kernel_batch * (model.runtime.max_seq_len + model.runtime.page_size - 1) // model.runtime.page_size - if model.runtime.total_kv_pages != expected_pages: + if model.runtime.total_kv_pages < kernel_batch: raise ValueError( - "PyPTO Qwen3-14B kernels require total_kv_pages to match the runtime batch capacity: " - f"{model.runtime.total_kv_pages} provided, {expected_pages} required." + f"total_kv_pages must be at least kernel_batch ({kernel_batch}), " + f"got {model.runtime.total_kv_pages}" ) @staticmethod @@ -774,3 +798,104 @@ def _validate_supported_shape(model: RuntimeModel) -> None: "PyPTO Qwen3-14B kernels require runtime page_size " f"{_QWEN14B_PAGE_SIZE}, got {model.runtime.page_size}." ) + + def _compile_prefill_fwd_tq_callable( + self, + jit_fn: object, + *, + batch: int, + max_seq: int, + block_table_stride: int, + hidden_size: int, + intermediate_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + num_layers: int, + vocab_size: int, + page_size: int, + ) -> _L3Callable: + """Compile the TurboQuant ``prefill_fwd_tq`` HOST wrapper into an L3 callable.""" + kv_hidden = num_kv_heads * head_dim + total_tokens = batch * max_seq + runtime_cache_blocks = (max_seq + page_size - 1) // page_size + cache_rows = batch * runtime_cache_blocks * num_layers * num_kv_heads * page_size + rot_rows = num_layers * head_dim + dummy_args = [ + torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16), + torch.empty((batch,), dtype=torch.int32), + torch.empty((batch,), dtype=torch.int32), + torch.empty((batch,), dtype=torch.int32), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), + torch.empty((num_layers, head_dim), dtype=torch.float32), + torch.empty((num_layers, head_dim), dtype=torch.float32), + torch.empty((max_seq, head_dim), dtype=torch.float32), + torch.empty((max_seq, head_dim), dtype=torch.float32), + torch.empty((batch * block_table_stride,), dtype=torch.int32), + torch.empty((total_tokens,), dtype=torch.int32), + torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), + torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), + torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), + torch.empty((1, hidden_size), dtype=torch.float32), + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), + torch.empty((batch, vocab_size), dtype=torch.float32), + torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), + torch.empty((1, 16), dtype=torch.float32), + ] + return self._compile_jit_fwd_callable("prefill_fwd_tq", jit_fn, dummy_args) + + def _compile_decode_fwd_tq_callable( + self, + jit_fn: object, + *, + batch: int, + max_seq: int, + block_table_stride: int, + hidden_size: int, + intermediate_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + num_layers: int, + vocab_size: int, + page_size: int, + ) -> _L3Callable: + """Compile the TurboQuant ``decode_fwd_tq`` HOST wrapper into an L3 callable.""" + kv_hidden = num_kv_heads * head_dim + runtime_cache_blocks = (max_seq + page_size - 1) // page_size + cache_rows = num_layers * batch * runtime_cache_blocks * num_kv_heads * page_size + rot_rows = num_layers * head_dim + dummy_args = [ + torch.empty((batch, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), + torch.empty((num_layers, head_dim), dtype=torch.float32), + torch.empty((num_layers, head_dim), dtype=torch.float32), + torch.empty((batch,), dtype=torch.int32), + torch.empty((batch * block_table_stride,), dtype=torch.int32), + torch.empty((batch,), dtype=torch.int32), + torch.empty((max_seq, head_dim), dtype=torch.float32), + torch.empty((max_seq, head_dim), dtype=torch.float32), + torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), + torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), + torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), + torch.empty((num_layers, hidden_size), dtype=torch.float32), + torch.empty((1, hidden_size), dtype=torch.float32), + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), + torch.empty((batch, vocab_size), dtype=torch.float32), + torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), + torch.empty((1, 16), dtype=torch.float32), + ] + return self._compile_jit_fwd_callable("decode_fwd_tq", jit_fn, dummy_args) diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index 3c801a2..5ab7344 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -9,6 +9,9 @@ from __future__ import annotations +import logging +import os +import time from dataclasses import dataclass from typing import Any @@ -28,6 +31,9 @@ from python.profile import profile_span +logger = logging.getLogger(__name__) + + def _kernel_trace_name(kernel_name: str) -> str: if "prefill" in kernel_name: return "kernel.prefill_fwd" @@ -75,11 +81,14 @@ class _CompiledKernels: prefill: _L3Callable decode: _L3Callable + greedy_sample: _L3Callable + token_embed: _L3Callable final_norm_weight: torch.Tensor rope_cos: torch.Tensor rope_sin: torch.Tensor padded_vocab: int padded_lm_head_weight: torch.Tensor + padded_embed_weight: torch.Tensor decode_weights: dict[str, torch.Tensor] prefill_hidden_buffer: torch.Tensor prefill_seq_lens_buffer: torch.Tensor @@ -88,15 +97,20 @@ class _CompiledKernels: prefill_block_table_buffer: torch.Tensor prefill_slot_mapping_buffer: torch.Tensor prefill_logits_buffer: torch.Tensor + prefill_sampled_ids_buffer: torch.Tensor + prefill_next_hidden_buffer: torch.Tensor decode_hidden_buffer: torch.Tensor decode_seq_lens_buffer: torch.Tensor decode_block_table_buffer: torch.Tensor decode_slot_mapping_buffer: torch.Tensor decode_logits_buffer: torch.Tensor + 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 # [num_layers*head_dim, head_dim] BF16 - tq_codebook: torch.Tensor | None = None # [1, 16] FP32 Lloyd-Max centroids + rot_matrices: torch.Tensor | None = None + tq_codebook: torch.Tensor | None = None @dataclass @@ -117,6 +131,7 @@ class _DecodeInputs: """Active user rows prepared for decode.""" actual_batch: int + token_ids: torch.Tensor hidden: torch.Tensor seq_lens: torch.Tensor block_table: torch.Tensor @@ -128,6 +143,7 @@ class _DecodeKernelInputs: """Fixed-batch tensors passed to the fused decode kernel.""" actual_batch: int + token_ids: torch.Tensor hidden: torch.Tensor seq_lens: torch.Tensor block_table: torch.Tensor @@ -150,6 +166,7 @@ class _StaticKernelArgs: rope_cos: _StaticDeviceTensor rope_sin: _StaticDeviceTensor 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 @@ -163,10 +180,12 @@ def __init__( self, *, 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] = {} @@ -176,17 +195,295 @@ def __init__( self._share_static_kernel_tensors() self._static_args = self._build_static_kernel_args() - def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConfig) -> None: - """Create the L3 worker-resident cache before the first request.""" + #: Scratch KV pages for the profile pass — slot=-1 means only page 0 + #: is ever touched (reads via block_table=0, writes via slot clamp to 0). + _PROFILE_PAGES = 1 + + def init_kv_cache(self, model_id: str, config: ModelConfig, runtime: RuntimeConfig) -> int: + """Create the L3 worker-resident cache before the first request. + + Order (vLLM-style): run a profile warmup FIRST so the simpler arena + is allocated before the KV cache competes for HBM. The profile uses + ``slot_mapping=-1`` / ``block_table=0`` so only a single dummy page + is needed. The KV cache size is then computed by the estimation + formula and allocated into the remaining space; if allocation fails + the page count is halved and retried. + """ if model_id in self._kv_caches: - return + num_pages = self._kv_caches[model_id].key_pages.shape[0] // ( + config.num_hidden_layers * config.num_key_value_heads * runtime.page_size + ) + return num_pages self._pending_kv_cache_specs[model_id] = (config, runtime) + + logger.info("[init_kv_cache] creating L3 worker …") with profile_span("Qwen314BModelRunner.prepare_l3_worker", cat="executor"): self._shared_l3_worker() + + logger.info("[init_kv_cache] uploading static tensors …") with profile_span("Qwen314BModelRunner.upload_static_tensors", cat="executor"): self._materialize_static_tensors() - with profile_span("Qwen314BModelRunner.init_kv_cache", cat="executor"): - ModelRunner.init_kv_cache(self, model_id, config, runtime) + + # -- phase 1: profile warmup → arena allocated ---------------------- + # Uses slot_mapping=-1 so no real KV cache pages are needed; the + # 1-page scratch is the dummy target for all reads/writes. + logger.info(f"[init_kv_cache] profile warmup (scratch {self._PROFILE_PAGES} page) …") + ModelRunner.init_kv_cache(self, model_id, config, runtime, num_pages=self._PROFILE_PAGES) + try: + self._warmup_dispatch(runtime) + finally: + self.close_kv_cache() + self._kv_caches.pop(model_id, None) + + # -- phase 2: real KV cache, halve-and-retry on OOM ----------------- + logger.info("[init_kv_cache] computing KV cache pages …") + num_pages = self._compute_kv_cache_pages(config, runtime, self._device_id) + num_pages = self._alloc_kv_cache_with_retry(model_id, config, runtime, num_pages) + self._print_memory_breakdown("after KV cache alloc", config, runtime, num_pages, self._device_id) + logger.info("[init_kv_cache] done") + return num_pages + + def _alloc_kv_cache_with_retry( + self, model_id: str, config: ModelConfig, runtime: RuntimeConfig, num_pages: int, + ) -> int: + """Allocate the KV cache, halving the page count on OOM.""" + floor = max(runtime.max_batch_size, 1) + requested = num_pages + num_pages = max(num_pages, floor) # always try at least the floor + while num_pages >= floor: + 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 + ) + logger.info( + f"[init_kv_cache] allocated {num_pages} pages " + f"(requested {requested}, downgraded after OOM): " + f"{num_pages * bytes_per_page / 1e9:.2f} GB KV cache, " + f"{num_pages * runtime.page_size} context tokens", + + ) + return num_pages + except (RuntimeError, MemoryError) as e: + prev = num_pages + num_pages //= 2 + if num_pages < floor and prev > floor: + num_pages = floor + logger.info( + f"[init_kv_cache] alloc failed ({e}); retrying {prev} -> {num_pages}", + ) + raise RuntimeError( + f"KV cache allocation failed even at floor {floor} pages" + ) + + @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. + + Called AFTER the profile warm-up, so weights, the simpler ring-heap + arena, compiled buffers and any persistent scratch are already + allocated — ``peak_non_kv = total − free`` captures all of it. The KV + budget is ``total x utilization − peak_non_kv``, leaving + ``total x (1 − utilization)`` as a fixed absolute headroom (more robust + 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 + ) + 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, " + "peak_non_kv=%.2f GB, kv_budget=%.2f GB, requested_pages=%d (%.1f MB/page)", + total_bytes / 1e9, utilization, peak_non_kv / 1e9, + kv_budget / 1e9, num_pages, bytes_per_page / 1e6, + ) + return num_pages + + @staticmethod + 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``. + + ``torch.npu.mem_get_info`` only reports a single total, so each part + is reconstructed rather than queried: weights (estimated from the + model config), KV cache (exact = num_pages x bytes_per_page), simpler + ring-heap arena (from the ``PTO2_RING_HEAP`` env x 4), and the + residual (compiled buffers + transient activation scratch + overhead). + """ + free_bytes, total_bytes = torch.npu.mem_get_info(f"npu:{device_id}") + used_bytes = total_bytes - free_bytes + dtype_bytes = getattr(torch, runtime.kv_dtype).itemsize + + # Weights — GQA: Q/O are hiddenxhidden, K/V are hiddenxkv_hidden. + hidden = config.hidden_size + kv_hidden = config.num_key_value_heads * config.head_dim + wt_params = ( + config.num_hidden_layers * ( + hidden * hidden * 2 + + hidden * kv_hidden * 2 + + hidden * config.intermediate_size * 3 + + hidden * 4 + ) + + config.vocab_size * hidden + ) + 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 + ) + kv_bytes = num_pages * bytes_per_page + + # Simpler ring-heap arena — from env (matches _compute_kv_cache_pages). + ring_heap = int(os.environ.get("PTO2_RING_HEAP", 256 * 1024 * 1024)) + arena_bytes = ring_heap * 4 + 128 * 1024 * 1024 + + residual = used_bytes - weight_bytes - kv_bytes - arena_bytes + + logger.info(f"[mem-breakdown] {label}:") + logger.info( + f" total used (measured): {used_bytes / 1e9:7.2f} GB " + f"/ {total_bytes / 1e9:.2f} GB (free {free_bytes / 1e9:.2f} GB)", + + ) + logger.info(f" ├─ weights (estimated): {weight_bytes / 1e9:7.2f} GB") + kv_tokens = num_pages * runtime.page_size + max_seq_len = runtime.max_seq_len + worst_case_demand = runtime.max_batch_size * max_seq_len + max_len_reqs = kv_tokens // max(max_seq_len, 1) + logger.info( + f" ├─ KV cache ({num_pages} pages): {kv_bytes / 1e9:7.2f} GB " + f"({bytes_per_page / 1e6:.1f} MB/page)", + + ) + logger.info( + f" │ capacity = {kv_tokens} tokens " + f"≈ {max_len_reqs} x full-len({max_seq_len}) reqs; " + f"worst-case need {runtime.max_batch_size}x{max_seq_len}=" + f"{worst_case_demand} tokens" + + (" [OK]" if kv_tokens >= worst_case_demand else " [TIGHT]"), + + ) + logger.info(f" ├─ simpler arena (env x 4): {arena_bytes / 1e9:7.2f} GB") + logger.info( + f" └─ residual (buffers/scratch): {residual / 1e9:6.2f} GB " + f"(compiled buffers + transient activation scratch + overhead)", + + ) + logger.info( + " note: weights/arena are estimates, KV is exact; total is from " + "mem_get_info (may under-count simpler's rtMalloc pool).", + + ) + + def warmup(self, model: RuntimeModel) -> None: + """Dispatch a dummy prefill + decode through the L3 worker.""" + self._warmup_dispatch(model.runtime) + + def _warmup_dispatch(self, runtime: RuntimeConfig) -> None: + """Production-scale prefill + decode warm-up with slot_mapping=-1. + + Sizes the prefill to one serving scheduling step — total tokens = + ``max_num_batched_tokens`` spread across ``max_batch`` requests. + This deliberately exercises the kernel at the configured capacity so + that a too-large ``max_num_batched_tokens`` (which would hit the + single-die attention heap ceiling around seq≈415 in the 40-layer + fused prefill) fails at startup rather than on the first real + request. + """ + batch = runtime.max_batch_size + max_seq = runtime.max_seq_len + mnb = getattr(runtime, "max_num_batched_tokens", 4096) + step_tokens = min(mnb, batch * max_seq) + per_req = max(step_tokens // batch, 1) + total_tokens = per_req * batch + + logger.info( + f"[warmup] starting (batch={batch}, max_num_batched_tokens={mnb}, " + f"max_seq={max_seq}, per_req={per_req}, total_tokens={total_tokens}, slot=-1)", + + ) + compiled = self._compiled + kv_cache = list(self._kv_caches.values())[0] + + # -- prefill --------------------------------------------------------- + compiled.prefill_hidden_buffer[:total_tokens].zero_() + compiled.prefill_seq_lens_buffer.zero_() + compiled.prefill_chunk_lens_buffer.zero_() + compiled.prefill_chunk_offsets_buffer.zero_() + compiled.prefill_block_table_buffer.fill_(0) # all reads from page 0 + compiled.prefill_slot_mapping_buffer.fill_(-1) # all writes to page 0 + + token_offset = 0 + for b in range(batch): + compiled.prefill_seq_lens_buffer[b] = per_req + compiled.prefill_chunk_lens_buffer[b] = per_req + compiled.prefill_chunk_offsets_buffer[b] = token_offset + token_offset += per_req + + prefill_inputs = _PrefillInputs( + actual_batch=batch, + hidden=compiled.prefill_hidden_buffer[:total_tokens], + seq_lens=compiled.prefill_seq_lens_buffer, + chunk_lens=compiled.prefill_chunk_lens_buffer, + chunk_offsets=compiled.prefill_chunk_offsets_buffer, + block_table=compiled.prefill_block_table_buffer, + slot_mapping=compiled.prefill_slot_mapping_buffer, + ) + + logger.info(f"[warmup] prefill dispatch … (batch={batch}, tokens={total_tokens})") + t0 = time.perf_counter() + self._run_distributed_program( + compiled.prefill, + *self._prefill_kernel_args( + prefill_inputs, kv_cache.key_pages, kv_cache.value_pages, + compiled.prefill_logits_buffer, + ), + ) + logger.info(f"[warmup] prefill done ({time.perf_counter() - t0:.2f} s)") + + # -- decode (full fixed batch, minimal seq) ------------------------- + compiled.decode_hidden_buffer.zero_() + compiled.decode_token_ids_buffer.zero_() + compiled.decode_seq_lens_buffer.zero_() + compiled.decode_block_table_buffer.fill_(0) # all reads from page 0 + compiled.decode_slot_mapping_buffer.fill_(-1) # all writes to page 0 + + for b in range(batch): + compiled.decode_seq_lens_buffer[b] = min(per_req + 1, max_seq) + + decode_kernel_inputs = _DecodeKernelInputs( + actual_batch=batch, + token_ids=compiled.decode_token_ids_buffer, + hidden=compiled.decode_hidden_buffer, + seq_lens=compiled.decode_seq_lens_buffer, + block_table=compiled.decode_block_table_buffer, + slot_mapping=compiled.decode_slot_mapping_buffer, + logits=compiled.decode_logits_buffer, + ) + + logger.info(f"[warmup] decode dispatch … (batch={batch}, seq_len={per_req + 1})") + 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), + ) + logger.info(f"[warmup] decode done ({time.perf_counter() - t0:.2f} s)") + + logger.info("[warmup] complete") def _alloc_kv_cache_tensor(self, shape: tuple[int, ...], dtype: torch.dtype) -> DeviceTensor: """Allocate one worker-resident KV cache tensor shared by prefill/decode.""" @@ -243,11 +540,12 @@ def _share_static_kernel_tensors(self) -> None: def _iter_static_host_tensors(self) -> tuple[torch.Tensor, ...]: """Return host tensors that must be shared before the worker forks.""" compiled = self._compiled - tensors: tuple[torch.Tensor, ...] = ( + return ( compiled.final_norm_weight, compiled.rope_cos, compiled.rope_sin, compiled.padded_lm_head_weight, + compiled.padded_embed_weight, *compiled.decode_weights.values(), compiled.prefill_hidden_buffer, compiled.prefill_seq_lens_buffer, @@ -256,18 +554,17 @@ def _iter_static_host_tensors(self) -> tuple[torch.Tensor, ...]: compiled.prefill_block_table_buffer, compiled.prefill_slot_mapping_buffer, compiled.prefill_logits_buffer, + compiled.prefill_sampled_ids_buffer, + compiled.prefill_next_hidden_buffer, compiled.decode_hidden_buffer, compiled.decode_seq_lens_buffer, compiled.decode_block_table_buffer, compiled.decode_slot_mapping_buffer, compiled.decode_logits_buffer, + compiled.decode_token_ids_buffer, + compiled.decode_sampled_ids_buffer, + compiled.decode_next_hidden_buffer, ) - # TurboQuant rotation matrices + codebook are also static kernel inputs. - if compiled.rot_matrices is not None: - tensors += (compiled.rot_matrices,) - if compiled.tq_codebook is not None: - tensors += (compiled.tq_codebook,) - return tensors def _build_static_kernel_args(self) -> _StaticKernelArgs: """Create static device-upload markers once per runner.""" @@ -277,20 +574,11 @@ def _build_static_kernel_args(self) -> _StaticKernelArgs: rope_cos=self._static_device_tensor(compiled.rope_cos), rope_sin=self._static_device_tensor(compiled.rope_sin), padded_lm_head_weight=self._static_device_tensor(compiled.padded_lm_head_weight), + padded_embed_weight=self._static_device_tensor(compiled.padded_embed_weight), decode_weights={ 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: @@ -307,46 +595,37 @@ def run_prefill(self, model: RuntimeModel, batch: PrefillBatch) -> PrefillResult logits_padded = compiled.prefill_logits_buffer kv_cache = self._materialize_kv_cache(model) - if self._tq_mode: - quant_k = kv_cache.quant_k_pages - quant_v = kv_cache.quant_v_pages - k_scales = kv_cache.k_scales_pages - v_scales = kv_cache.v_scales_pages - if quant_k is None or quant_v is None or k_scales is None or v_scales is None: - raise RuntimeError("TurboQuant KV caches are not initialized") - self._validate_kv_cache_bounds( - model, prefill_inputs.block_table, prefill_inputs.slot_mapping, quant_k - ) - self._run_distributed_program( - compiled.prefill, - *self._prefill_tq_kernel_args( - prefill_inputs, quant_k, quant_v, k_scales, v_scales, logits_padded - ), - ) - else: - 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 - ) - self._run_distributed_program( - compiled.prefill, - *self._prefill_kernel_args(prefill_inputs, k_cache, v_cache, logits_padded), - ) + 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) + + self._run_distributed_program( + compiled.prefill, + *self._prefill_kernel_args(prefill_inputs, k_cache, v_cache, logits_padded), + ) for batch_idx, alloc in enumerate(batch.kv_allocations): seq_len = int(batch.seq_lens[batch_idx].item()) alloc.tokens_used = max(alloc.tokens_used, seq_len) + sampled_ids, next_hidden = self._maybe_run_sample_embed( + logits_padded, + compiled.prefill_sampled_ids_buffer, + compiled.prefill_next_hidden_buffer, + prefill_inputs.actual_batch, + allow=batch.allow_device_greedy_sampling, + ) return PrefillResult( last_hidden=None, logits=logits_padded[: prefill_inputs.actual_batch, : model.config.vocab_size], + sampled_token_ids=sampled_ids, + next_hidden_states=next_hidden, ) def run_decode(self, model: RuntimeModel, batch: DecodeBatch) -> DecodeResult: """Run the fused all-layer PAGED ``decode_layer.decode_fwd`` and return logits. ``decode_fwd`` runs all NUM_LAYERS + the LM head in one dispatch over the - PAGED KV pool, addressing KV via ``block_table`` + ``slot_mapping`` — the + PAGED KV pool, addressing KV via ``block_table`` + ``slot_mapping``, the SAME device-resident KV pool prefill writes (``self._kv_caches``), so prompt KV is already in place with no bridge. KV is keyed by block_table page id, not by kernel row, so a request may occupy any row each step (no stable-slot shim). @@ -365,44 +644,82 @@ 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 kernel_inputs = self._pad_decode_inputs(model, decode_inputs) - if self._tq_mode: - quant_k = kv_cache.quant_k_pages - quant_v = kv_cache.quant_v_pages - k_scales = kv_cache.k_scales_pages - v_scales = kv_cache.v_scales_pages - if quant_k is None or quant_v is None or k_scales is None or v_scales is None: - raise RuntimeError("TurboQuant KV caches are not initialized") - # 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, quant_k - ) - self._run_distributed_program( - compiled.decode, - *self._decode_tq_kernel_args(kernel_inputs, quant_k, quant_v, k_scales, v_scales), - ) - else: - k_cache = kv_cache.key_pages - v_cache = kv_cache.value_pages - # 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._run_distributed_program( - compiled.decode, - *self._decode_kernel_args(kernel_inputs, k_cache, v_cache), - ) + # 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._run_distributed_program( + compiled.decode, + *self._decode_kernel_args(kernel_inputs, k_cache, v_cache), + ) for batch_idx, alloc in enumerate(batch.kv_allocations): alloc.tokens_used = max(alloc.tokens_used, int(batch.seq_lens[batch_idx].item())) + sampled_ids, next_hidden = self._integrated_sample_result( + compiled.decode_sampled_ids_buffer, + # decode_fwd's next_hidden output is the embedding for sampled_ids_in + # used by this decode step. The newly sampled token is embedded at the + # start of the following decode_fwd call, so there is no next-step + # hidden row to return here. + None, + kernel_inputs.actual_batch, + allow=batch.allow_device_greedy_sampling, + ) return DecodeResult( hidden_states=decode_inputs.hidden.float(), logits=kernel_inputs.logits[: kernel_inputs.actual_batch, : model.config.vocab_size].to( decode_inputs.hidden.device ), + sampled_token_ids=sampled_ids, + next_hidden_states=next_hidden, + ) + + @staticmethod + def _integrated_sample_result( + sampled_ids_buffer: torch.Tensor, + next_hidden_buffer: torch.Tensor | None, + actual_batch: int, + *, + allow: bool, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Read device sampling output and optional precomputed next hidden rows.""" + if not allow: + return None, None + next_hidden = ( + next_hidden_buffer[:actual_batch].clone() + if next_hidden_buffer is not None + else None + ) + return ( + sampled_ids_buffer[:actual_batch, :1].clone(), + next_hidden, + ) + + def _maybe_run_sample_embed( + self, + logits: torch.Tensor, + sampled_ids_buffer: torch.Tensor, + next_hidden_buffer: torch.Tensor, + actual_batch: int, + *, + allow: bool, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Run device greedy sampling when the request is greedy.""" + if not allow: + return None, None + compiled = self._compiled + self._run_distributed_program( + compiled.greedy_sample, + logits, + sampled_ids_buffer, + ) + return ( + sampled_ids_buffer[:actual_batch, :1].clone(), + None, ) def _prefill_kernel_args( @@ -433,10 +750,10 @@ def _prefill_kernel_args( k_cache, v_cache, weights["decode_wo"], - weights["decode_post_rms_weight"], weights["decode_w_gate"], weights["decode_w_up"], weights["decode_w_down"], + weights["decode_post_rms_weight"], static.final_norm_weight, static.padded_lm_head_weight, logits, @@ -474,99 +791,10 @@ def _decode_kernel_args( static.final_norm_weight, static.padded_lm_head_weight, inputs.logits, - ) - - def _prefill_tq_kernel_args( - self, - inputs: _PrefillInputs, - quant_k: DeviceTensor, - quant_v: DeviceTensor, - k_scales: DeviceTensor, - v_scales: DeviceTensor, - logits: torch.Tensor, - ) -> tuple[Any, ...]: - """Return arguments in ``qwen3_prefill_tq_host`` signature order. - - TQ prefill drops chunk_lens/chunk_offsets and replaces the BF16 k/v cache - with UINT8 quantized caches + FP32 scales + rotation matrices + codebook. - """ - static = self._require_static_args() - if static.rot_matrices is None or static.tq_codebook is None: - raise RuntimeError("TurboQuant rotation matrices / codebook are not compiled") - weights = static.decode_weights - 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, - quant_k, - quant_v, - k_scales, - v_scales, - 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, - ) - - def _decode_tq_kernel_args( - self, - inputs: _DecodeKernelInputs, - quant_k: DeviceTensor, - quant_v: DeviceTensor, - k_scales: DeviceTensor, - v_scales: DeviceTensor, - ) -> tuple[Any, ...]: - """Return arguments in ``qwen3_decode_tq_host`` signature order. - - TQ decode replaces the BF16 k/v cache with the six TQ tensors. Weight - order is wo, post_rms, w_gate, w_up, w_down (differs from non-TQ decode). - """ - static = self._require_static_args() - if static.rot_matrices is None or static.tq_codebook is None: - raise RuntimeError("TurboQuant rotation matrices / codebook are not compiled") - weights = static.decode_weights - 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, - quant_k, - quant_v, - k_scales, - v_scales, - 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, + static.padded_embed_weight, + inputs.token_ids, + self._compiled.decode_sampled_ids_buffer, + self._compiled.decode_next_hidden_buffer, ) def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _DecodeKernelInputs: @@ -592,8 +820,24 @@ def _pad_decode_inputs(self, model: RuntimeModel, inputs: _DecodeInputs) -> _Dec if actual_batch < kernel_batch: hidden[actual_batch:].copy_(inputs.hidden[0:1].expand(kernel_batch - actual_batch, -1)) + token_ids = compiled.decode_token_ids_buffer + token_ids.zero_() + active_token_ids = inputs.token_ids.reshape(actual_batch, -1) + if active_token_ids.shape[1] != 1: + raise ValueError( + "decode token_ids must contain exactly one token per row, " + f"got shape {tuple(inputs.token_ids.shape)}" + ) + width = 1 + token_ids[:actual_batch, :width].copy_(active_token_ids[:, :width]) + if actual_batch < kernel_batch: + token_ids[actual_batch:, :width].copy_( + active_token_ids[0:1, :width].expand(kernel_batch - actual_batch, width) + ) + return _DecodeKernelInputs( actual_batch=actual_batch, + token_ids=token_ids, hidden=hidden, seq_lens=self._copy_replicated_rows( compiled.decode_seq_lens_buffer, @@ -652,7 +896,12 @@ def _shared_l3_worker(self) -> Any: if worker is None: from pypto.runtime import DistributedWorker # noqa: PLC0415 - worker = DistributedWorker([self._compiled.prefill.compiled, self._compiled.decode.compiled]) + worker = DistributedWorker([ + self._compiled.prefill.compiled, + self._compiled.decode.compiled, + self._compiled.greedy_sample.compiled, + self._compiled.token_embed.compiled, + ]) self._l3_worker = worker return worker @@ -678,14 +927,10 @@ def _materialize_static_tensors(self) -> None: static.rope_cos, static.rope_sin, static.padded_lm_head_weight, + static.padded_embed_weight, *static.decode_weights.values(), ): self._coerce_l3_arg(worker, arg) - # TurboQuant rotation matrices + codebook are also uploaded once. - if static.rot_matrices is not None: - self._coerce_l3_arg(worker, static.rot_matrices) - if static.tq_codebook is not None: - self._coerce_l3_arg(worker, static.tq_codebook) @staticmethod def _copy_replicated_rows( @@ -878,6 +1123,7 @@ def _prepare_decode_inputs( return _DecodeInputs( actual_batch=actual_batch, + token_ids=batch.token_ids.to(torch.int32).cpu(), hidden=hidden, seq_lens=seq_lens, block_table=block_table, diff --git a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py index ef03fe0..7a58e5e 100644 --- a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py +++ b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py @@ -187,6 +187,123 @@ def qwen3_prefill_tq_host( ) +@pl.jit.host +def qwen3_decode_tq_host( + hidden_states: pl.Tensor, + input_rms_weight: pl.Tensor, + wq: pl.Tensor, + wk: pl.Tensor, + wv: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + seq_lens: pl.Tensor, + block_table: pl.Tensor, + slot_mapping: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + quant_k_cache: pl.Tensor, + quant_v_cache: pl.Tensor, + quant_k_scales: pl.Tensor, + quant_v_scales: pl.Tensor, + rot_matrices: pl.Tensor, + tq_codebook: pl.Tensor, + wo: pl.Tensor, + post_rms_weight: pl.Tensor, + w_gate: pl.Tensor, + w_up: pl.Tensor, + w_down: pl.Tensor, + final_norm_weight: pl.Tensor, + lm_head_weight: pl.Tensor, + out: pl.Out[pl.Tensor], +) -> pl.Tensor: + return decode_fwd_tq( + hidden_states, + input_rms_weight, + wq, + wk, + wv, + q_norm_weight, + k_norm_weight, + seq_lens, + block_table, + slot_mapping, + rope_cos, + rope_sin, + quant_k_cache, + quant_v_cache, + quant_k_scales, + quant_v_scales, + rot_matrices, + tq_codebook, + wo, + post_rms_weight, + w_gate, + w_up, + w_down, + final_norm_weight, + lm_head_weight, + out, + ) +prefill_fwd_tq = None +decode_fwd_tq = None +def qwen3_prefill_tq_host( + hidden_states: pl.Tensor, + seq_lens: pl.Tensor, + input_rms_weight: pl.Tensor, + wq: pl.Tensor, + wk: pl.Tensor, + wv: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + block_table: pl.Tensor, + slot_mapping: pl.Tensor, + quant_k_cache: pl.Tensor, + quant_v_cache: pl.Tensor, + quant_k_scales: pl.Tensor, + quant_v_scales: pl.Tensor, + rot_matrices: pl.Tensor, + tq_codebook: pl.Tensor, + wo: pl.Tensor, + post_rms_weight: pl.Tensor, + w_gate: pl.Tensor, + w_up: pl.Tensor, + w_down: pl.Tensor, + final_norm_weight: pl.Tensor, + lm_head_weight: pl.Tensor, + out: pl.Out[pl.Tensor], +) -> pl.Tensor: + return prefill_fwd_tq( + hidden_states, + seq_lens, + input_rms_weight, + wq, + wk, + wv, + q_norm_weight, + k_norm_weight, + rope_cos, + rope_sin, + block_table, + slot_mapping, + quant_k_cache, + quant_v_cache, + quant_k_scales, + quant_v_scales, + rot_matrices, + tq_codebook, + wo, + post_rms_weight, + w_gate, + w_up, + w_down, + final_norm_weight, + lm_head_weight, + out, + ) + + @pl.jit.host def qwen3_decode_tq_host( hidden_states: pl.Tensor, From affc6b932b91cad51448f528b6b1971bc45c020d Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 12:17:35 +0800 Subject: [PATCH 04/17] kv_cache: remove TQ hardcoded num_pages, use same dynamic sizing as non-TQ --- python/core/kv_cache.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/python/core/kv_cache.py b/python/core/kv_cache.py index 49cfb6f..d44c62f 100644 --- a/python/core/kv_cache.py +++ b/python/core/kv_cache.py @@ -198,19 +198,6 @@ def register_model(self, model_id: str, config: ModelConfig, runtime: RuntimeCon # total_kv_pages explicitly to cover simultaneous workspace demand. kv_quant_config = runtime.kv_quant_config num_pages = runtime.total_kv_pages - if num_pages is None and kv_quant_config is not None and kv_quant_config.enabled: - # NOTE: The fused NPU decode kernel has the per-layer cache stride - # (num_pages * num_kv_heads * page_size) baked in at compile time. - # Reducing the page count would cause incorrect layer offsets and a - # device hang. Keep the full page count; TurboQuant saves memory by - # compressing old tokens in-place (freeing bf16 pages back to the - # pool for reuse by *other* requests on the same model). - num_pages = runtime.max_batch_size * max_blocks_per_seq - print( - f"[KV-Quant] Using full bf16 pages ({num_pages}) to match " - f"compiled kernel stride. Compression frees pages for reuse.", - flush=True, - ) if num_pages is None: num_pages = runtime.max_batch_size * max_blocks_per_seq self._init_blocks(num_pages, runtime.page_size) From 998503804f61fac87ef95d1b80b79c1fa0421974 Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 12:19:22 +0800 Subject: [PATCH 05/17] fix: resolve duplicate kv_dtype/weight_dtype in npu_generate.py --- examples/model/qwen3_14b/npu_generate.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/examples/model/qwen3_14b/npu_generate.py b/examples/model/qwen3_14b/npu_generate.py index abd699c..cf99603 100644 --- a/examples/model/qwen3_14b/npu_generate.py +++ b/examples/model/qwen3_14b/npu_generate.py @@ -473,11 +473,9 @@ def main() -> None: max_seq_len=args.max_seq_len, max_new_tokens=args.max_new_tokens, device="cpu", - kv_dtype="bfloat16", - weight_dtype="bfloat16", - kv_quant_config=KvQuantConfig(enabled=True) if args.tq_mode else None, 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 From 62f74c2c80a84d72e5e666c0f380dfeec24d7029 Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 14:07:02 +0800 Subject: [PATCH 06/17] fix: remove hardcoded total_kv_pages=200 from npu_generate.py --- examples/model/qwen3_14b/npu_generate.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/examples/model/qwen3_14b/npu_generate.py b/examples/model/qwen3_14b/npu_generate.py index cf99603..ec3ddd0 100644 --- a/examples/model/qwen3_14b/npu_generate.py +++ b/examples/model/qwen3_14b/npu_generate.py @@ -478,10 +478,6 @@ def main() -> None: 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: From adb3e2a89f3b0223d7239d27690c14e38de3fcad Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 14:33:50 +0800 Subject: [PATCH 07/17] kv_cache: restore TQ full-page-count override in register_model Commit affc6b9 removed the TurboQuant-specific num_pages override so TQ would use the same dynamic sizing as non-TQ. This caused a device hang in TQ mode: the fused NPU decode kernel has the per-layer cache stride (num_pages * num_kv_heads * page_size) baked in at compile time, so reducing the page count yields incorrect layer offsets and hangs. Restore the override (num_pages = max_batch_size * max_blocks_per_seq when kv_quant_config.enabled and total_kv_pages is unset) with the explanatory comment so it isn't removed again. Compression still saves memory by freeing bf16 pages in-place for reuse. Co-Authored-By: Claude --- python/core/kv_cache.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/core/kv_cache.py b/python/core/kv_cache.py index d44c62f..6c0ab6e 100644 --- a/python/core/kv_cache.py +++ b/python/core/kv_cache.py @@ -198,6 +198,14 @@ def register_model(self, model_id: str, config: ModelConfig, runtime: RuntimeCon # total_kv_pages explicitly to cover simultaneous workspace demand. kv_quant_config = runtime.kv_quant_config num_pages = runtime.total_kv_pages + if num_pages is None and kv_quant_config is not None and kv_quant_config.enabled: + # NOTE: The fused NPU decode kernel has the per-layer cache stride + # (num_pages * num_kv_heads * page_size) baked in at compile time. + # Reducing the page count would cause incorrect layer offsets and a + # device hang. Keep the full page count; TurboQuant saves memory by + # compressing old tokens in-place (freeing bf16 pages back to the + # pool for reuse by *other* requests on the same model). + num_pages = runtime.max_batch_size * max_blocks_per_seq if num_pages is None: num_pages = runtime.max_batch_size * max_blocks_per_seq self._init_blocks(num_pages, runtime.page_size) From 6ecc16dd8c490bba1f0f7ccdd732bd0ce78da4e0 Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 15:02:49 +0800 Subject: [PATCH 08/17] turboquant: fix KVCompressor construction hang (CPU thread contention) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TQ mode hung at "[KV-Quant] Creating KVCompressor ..." because the merged compressor (MSECompressor + LloydMaxCodebook) was pathologically slow on this 320-CPU-core host — two independent causes: 1. solve_lloyd_max used pure-Python loops: a list comprehension calling pdf() over 2048 samples, nested in 2^bits levels x 200 iterations, per codebook, x80 compressors. Vectorized the numerical integration into one (n_levels, n_samples) tensor op and memoized codebooks by (d, bits) since 80 compressors share only ~1-2 unique bit-widths. 2. Rotation-matrix QR (_generate_rotation_matrix) ran one 128x128 QR per compressor. With torch defaulting to all 320 CPU threads, thread- contention overhead on tiny matrices dominated: ~140ms/matrix (11s for 80). Batched all per-layer rotations into a single batched QR in KVCompressor and capped torch CPU threads to 8 for setup-only ops via a limited_cpu_threads() context manager (restored after). Inference runs on the NPU, so CPU thread count does not affect serving throughput. Net: KVCompressor(40 layers) from >26min (effectively hung) to ~0.17s. Rotation matrices verified bit-identical to the per-seed single-matrix path; codebook math unchanged. Co-Authored-By: Claude --- python/core/turboquant/compressor.py | 71 +++++++++++++-- python/core/turboquant/lloyd_max.py | 125 ++++++++++++++++++--------- 2 files changed, 147 insertions(+), 49 deletions(-) diff --git a/python/core/turboquant/compressor.py b/python/core/turboquant/compressor.py index 9ba6c24..e0fe69d 100644 --- a/python/core/turboquant/compressor.py +++ b/python/core/turboquant/compressor.py @@ -6,13 +6,12 @@ from __future__ import annotations -import math from dataclasses import dataclass import torch import torch.nn.functional as F -from .lloyd_max import LloydMaxCodebook +from .lloyd_max import get_codebook, limited_cpu_threads def _generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.Tensor: @@ -27,6 +26,29 @@ def _generate_rotation_matrix(d: int, seed: int, device: str = "cpu") -> torch.T return Q.to(device) +def _generate_rotation_matrices( + seeds: list[int], d: int, device: str = "cpu" +) -> torch.Tensor: + """Generate orthogonal rotation matrices for many seeds in one batched QR. + + Batching is essential on hosts with many CPU cores: torch dispatches small + per-matrix QR calls across all cores, and the thread-contention overhead on + 128x128 matrices dominates wall time (observed ~140 ms/matrix at 320 threads + vs <2 ms batched). Returns a (len(seeds), d, d) tensor. + """ + mats = [] + for seed in seeds: + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + mats.append(torch.randn(d, d, generator=gen)) + G = torch.stack(mats) # (N, d, d) + with limited_cpu_threads(): + Q, R = torch.linalg.qr(G) # batched QR + diag_sign = torch.sign(torch.diagonal(R, dim1=1, dim2=2)) # (N, d) + diag_sign[diag_sign == 0] = 1.0 + return (Q * diag_sign.unsqueeze(1)).to(device) # (N, d, d) + + class MSECompressor: """Single-stage MSE-optimal compressor for one side (keys or values). @@ -34,13 +56,26 @@ class MSECompressor: and stores bit-packed indices + norms. """ - def __init__(self, head_dim: int, bits: int, seed: int, device: str = "cpu"): + def __init__( + self, + head_dim: int, + bits: int, + seed: int, + device: str = "cpu", + rotation: torch.Tensor | None = None, + ): self.head_dim = head_dim self.bits = bits self.device = device - self.Pi = _generate_rotation_matrix(head_dim, seed=seed, device=device) - self.centroids = LloydMaxCodebook(head_dim, bits).centroids.to(device) + # Allow callers that batch many rotation matrices (e.g. KVCompressor) to + # pass a precomputed matrix and skip the per-instance QR. + self.Pi = ( + rotation + if rotation is not None + else _generate_rotation_matrix(head_dim, seed=seed, device=device) + ) + self.centroids = get_codebook(head_dim, bits).centroids.to(device) @torch.no_grad() def compress(self, states: torch.Tensor) -> dict: @@ -144,6 +179,15 @@ def __init__( self.key_compressors: list[MSECompressor] = [] self.val_compressors: list[MSECompressor] = [] + # Precompute every layer's key/value rotation matrix in a single batched + # QR (see _generate_rotation_matrices). seeds are interleaved key, val. + rotation_seeds: list[int] = [] + for layer_idx in range(num_layers): + seed_base = seed + layer_idx * 1000 + rotation_seeds.append(seed_base) # key + rotation_seeds.append(seed_base + 500) # val + rotations = _generate_rotation_matrices(rotation_seeds, head_dim, device=device) + for layer_idx in range(num_layers): is_protected = ( layer_idx < config.protected_layers @@ -154,12 +198,23 @@ def __init__( effective_key_bits = min(effective_key_bits, 8) effective_val_bits = min(effective_val_bits, 8) - seed_base = seed + layer_idx * 1000 self.key_compressors.append( - MSECompressor(head_dim, effective_key_bits, seed=seed_base, device=device) + MSECompressor( + head_dim, + effective_key_bits, + seed=rotation_seeds[2 * layer_idx], + device=device, + rotation=rotations[2 * layer_idx], + ) ) self.val_compressors.append( - MSECompressor(head_dim, effective_val_bits, seed=seed_base + 500, device=device) + MSECompressor( + head_dim, + effective_val_bits, + seed=rotation_seeds[2 * layer_idx + 1], + device=device, + rotation=rotations[2 * layer_idx + 1], + ) ) @torch.no_grad() diff --git a/python/core/turboquant/lloyd_max.py b/python/core/turboquant/lloyd_max.py index 36db0b2..467b34d 100644 --- a/python/core/turboquant/lloyd_max.py +++ b/python/core/turboquant/lloyd_max.py @@ -5,10 +5,37 @@ We solve the Lloyd-Max conditions to find optimal centroids. """ +import contextlib import math import torch +# Lloyd-Max codebooks are a deterministic function of (d, bits); many layers +# share the same bit-width, so cache the solved codebook instead of recomputing +# it (the per-layer compressors create one codebook each — 80+ for Qwen3-14B). +_CODEBOOK_CACHE: dict[tuple[int, int], "LloydMaxCodebook"] = {} + + +@contextlib.contextmanager +def limited_cpu_threads(n: int = 8): + """Temporarily cap torch's CPU thread count. + + On hosts with many cores torch defaults to using all of them, which causes + severe thread-contention overhead on the small CPU tensors used during + TurboQuant setup (QR decomposition, Lloyd-Max numerical integration) — + observed >100x slowdown at 320 threads vs 8. The inference workload itself + runs on the NPU and does not depend on torch CPU threading, so capping + threads during these setup-only computations is safe. + """ + prev = torch.get_num_threads() + if prev > n: + torch.set_num_threads(n) + try: + yield + finally: + if prev != torch.get_num_threads(): + torch.set_num_threads(prev) + def solve_lloyd_max(d: int, bits: int, max_iter: int = 200, tol: float = 1e-10): """Solve Lloyd-Max optimal quantizer for N(0, 1/d). @@ -19,50 +46,56 @@ def solve_lloyd_max(d: int, bits: int, max_iter: int = 200, tol: float = 1e-10): """ n_levels = 2 ** bits sigma = 1.0 / math.sqrt(d) + var = sigma * sigma + # Precompute the Gaussian N(0, sigma^2) normalization once. + norm = 1.0 / math.sqrt(2.0 * math.pi * var) - # Initialize centroids uniformly in [-3.5*sigma, 3.5*sigma] + # Initialize centroids uniformly in [-3.5*sigma, 3.5*sigma]. lo, hi = -3.5 * sigma, 3.5 * sigma - centroids = [lo + (hi - lo) * (i + 0.5) / n_levels for i in range(n_levels)] - - # Gaussian PDF: N(0, sigma^2) - def pdf(x): - return (1.0 / math.sqrt(2 * math.pi * sigma * sigma)) * math.exp( - -x * x / (2 * sigma * sigma) - ) - - for _ in range(max_iter): - # Step 1: Compute boundaries (midpoints) - boundaries = [ - (centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1) - ] - - # Step 2: Update centroids as conditional expectations - edges = [lo * 3] + boundaries + [hi * 3] - new_centroids = [] - for i in range(n_levels): - a, b = edges[i], edges[i + 1] - # Numerical integration via sampling (avoids scipy dependency) - n_samples = 2048 - xs = torch.linspace(a, b, n_samples) - pdf_vals = torch.tensor([pdf(x) for x in xs]) - weighted = xs * pdf_vals - if pdf_vals.sum() > 1e-15: - new_centroids.append((weighted.sum() / pdf_vals.sum()).item()) - else: - new_centroids.append(centroids[i]) - - # Check convergence - max_shift = max( - abs(new_centroids[i] - centroids[i]) for i in range(n_levels) - ) - centroids = new_centroids - if max_shift < tol: - break - - boundaries = [(centroids[i] + centroids[i + 1]) / 2.0 for i in range(n_levels - 1)] + idx = torch.arange(n_levels, dtype=torch.float64) + centroids = lo + (hi - lo) * (idx + 0.5) / n_levels + + # Integration grid shared across all levels each iteration. Using float64 + # keeps the conditional-expectation integration accurate; the result is cast + # back to float32 at the end. + n_samples = 2048 + grid = torch.linspace(0.0, 1.0, n_samples, dtype=torch.float64) + + with limited_cpu_threads(): + for _ in range(max_iter): + # Step 1: boundaries are midpoints between adjacent centroids. + boundaries = (centroids[:-1] + centroids[1:]) / 2.0 + edges = torch.cat( + [ + torch.tensor([lo * 3.0], dtype=torch.float64), + boundaries, + torch.tensor([hi * 3.0], dtype=torch.float64), + ] + ) + a = edges[:-1] # (n_levels,) + b = edges[1:] # (n_levels,) + + # Step 2: update each centroid as the conditional expectation over + # its bucket [a, b]. Build the (n_levels, n_samples) sample grid in + # one shot and evaluate the Gaussian density vectorized (the + # pure-Python loop over samples/levels here previously dominated + # KVCompressor construction). + xs = a.unsqueeze(1) + (b - a).unsqueeze(1) * grid # (n_levels, n_samples) + pdf_vals = norm * torch.exp(-xs * xs / (2.0 * var)) # (n_levels, n_samples) + weight = pdf_vals.sum(dim=1) + new_centroids = (xs * pdf_vals).sum(dim=1) / weight.clamp_min(1e-15) + # Buckets with negligible probability keep their previous centroid. + new_centroids = torch.where(weight > 1e-15, new_centroids, centroids) + + max_shift = (new_centroids - centroids).abs().max().item() + centroids = new_centroids + if max_shift < tol: + break + + boundaries = (centroids[:-1] + centroids[1:]) / 2.0 return ( - torch.tensor(centroids, dtype=torch.float32), - torch.tensor(boundaries, dtype=torch.float32), + centroids.to(torch.float32), + boundaries.to(torch.float32), ) @@ -77,3 +110,13 @@ def __init__(self, d: int, bits: int): def __repr__(self): return f"LloydMaxCodebook(d={self.d}, bits={self.bits}, levels={self.n_levels})" + + +def get_codebook(d: int, bits: int) -> LloydMaxCodebook: + """Return a cached LloydMaxCodebook for (d, bits), solving it once.""" + key = (d, bits) + cb = _CODEBOOK_CACHE.get(key) + if cb is None: + cb = LloydMaxCodebook(d, bits) + _CODEBOOK_CACHE[key] = cb + return cb From e1b34249e97026d0edc3e0c724ba9fece85b3b5e Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 17:02:44 +0800 Subject: [PATCH 09/17] npu_executor: gate REAL_VOCAB/sampling validation under non-TQ only TQ mode crashed in _compile_model with "module '_pypto_lib_qwen3_14b_decode_tq' has no attribute 'REAL_VOCAB'" because the REAL_VOCAB / greedy_sample / token_embed validation block ran unconditionally. In TQ mode qwen3_decode_layer is the decode_tq module (which has VOCAB but not REAL_VOCAB), and qwen3_greedy_sample / qwen3_token_embed are never loaded at all. Gate the whole block under `if not self._tq_mode:` (the intended structure per the device-sampling merge); TQ falls through to the existing `else: sampled_ids_width = 1`. Co-Authored-By: Claude --- .../model/qwen3_14b/runner/npu_executor.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/examples/model/qwen3_14b/runner/npu_executor.py b/examples/model/qwen3_14b/runner/npu_executor.py index c615ad1..d6f3453 100644 --- a/examples/model/qwen3_14b/runner/npu_executor.py +++ b/examples/model/qwen3_14b/runner/npu_executor.py @@ -212,12 +212,18 @@ def _mark(label: str) -> None: f"{padded_vocab} (round_up({model.config.vocab_size}, {_VOCAB_PAD_MULTIPLE})); " "they must match for the decode logits buffer / lm_head weight to line up." ) - if model.config.vocab_size != int(qwen3_decode_layer.REAL_VOCAB): - raise ValueError( - "decode_layer.decode_fwd hard-codes REAL_VOCAB for padded-token masking, " - f"but the runtime model vocab_size is {model.config.vocab_size}; expected " - f"{int(qwen3_decode_layer.REAL_VOCAB)}." - ) + if not self._tq_mode: + # REAL_VOCAB / greedy_sample / token_embed are constants of the + # non-TQ fused decode + device-side sampling kernels. The TQ decode + # kernel (decode_tq) has none of these (sampling/embedding differ), + # so the whole validation block is non-TQ only; TQ uses + # sampled_ids_width=1 below. + if model.config.vocab_size != int(qwen3_decode_layer.REAL_VOCAB): + raise ValueError( + "decode_layer.decode_fwd hard-codes REAL_VOCAB for padded-token masking, " + f"but the runtime model vocab_size is {model.config.vocab_size}; expected " + f"{int(qwen3_decode_layer.REAL_VOCAB)}." + ) if int(qwen3_greedy_sample.BATCH) != kernel_batch: raise ValueError( "greedy_sample_fwd is compiled for a fixed kernel BATCH of " From 7042928dd4658049cfc4cbfdeeed71c4a20f3b3e Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 17:10:12 +0800 Subject: [PATCH 10/17] qwen3_l3_dispatch: remove duplicate undecorated TQ host wrappers The merge appended a second copy of the TQ host wrappers (prefill_fwd_tq/ decode_fwd_tq = None reset + qwen3_prefill_tq_host + qwen3_decode_tq_host) after the decorated originals. The duplicate qwen3_prefill_tq_host was a plain `def` (no @pl.jit.host), so it shadowed the decorated version and qwen3_prefill_tq_host.compile(...) failed with "'function' object has no attribute 'compile'" during TQ kernel compile. Deleted the entire duplicate block (117 lines); the single decorated originals (lines 131/190) are kept, so both TQ wrappers now carry .compile. Co-Authored-By: Claude --- .../qwen3_14b/runner/qwen3_l3_dispatch.py | 117 ------------------ 1 file changed, 117 deletions(-) diff --git a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py index 7a58e5e..ef03fe0 100644 --- a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py +++ b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py @@ -187,123 +187,6 @@ def qwen3_prefill_tq_host( ) -@pl.jit.host -def qwen3_decode_tq_host( - hidden_states: pl.Tensor, - input_rms_weight: pl.Tensor, - wq: pl.Tensor, - wk: pl.Tensor, - wv: pl.Tensor, - q_norm_weight: pl.Tensor, - k_norm_weight: pl.Tensor, - seq_lens: pl.Tensor, - block_table: pl.Tensor, - slot_mapping: pl.Tensor, - rope_cos: pl.Tensor, - rope_sin: pl.Tensor, - quant_k_cache: pl.Tensor, - quant_v_cache: pl.Tensor, - quant_k_scales: pl.Tensor, - quant_v_scales: pl.Tensor, - rot_matrices: pl.Tensor, - tq_codebook: pl.Tensor, - wo: pl.Tensor, - post_rms_weight: pl.Tensor, - w_gate: pl.Tensor, - w_up: pl.Tensor, - w_down: pl.Tensor, - final_norm_weight: pl.Tensor, - lm_head_weight: pl.Tensor, - out: pl.Out[pl.Tensor], -) -> pl.Tensor: - return decode_fwd_tq( - hidden_states, - input_rms_weight, - wq, - wk, - wv, - q_norm_weight, - k_norm_weight, - seq_lens, - block_table, - slot_mapping, - rope_cos, - rope_sin, - quant_k_cache, - quant_v_cache, - quant_k_scales, - quant_v_scales, - rot_matrices, - tq_codebook, - wo, - post_rms_weight, - w_gate, - w_up, - w_down, - final_norm_weight, - lm_head_weight, - out, - ) -prefill_fwd_tq = None -decode_fwd_tq = None -def qwen3_prefill_tq_host( - hidden_states: pl.Tensor, - seq_lens: pl.Tensor, - input_rms_weight: pl.Tensor, - wq: pl.Tensor, - wk: pl.Tensor, - wv: pl.Tensor, - q_norm_weight: pl.Tensor, - k_norm_weight: pl.Tensor, - rope_cos: pl.Tensor, - rope_sin: pl.Tensor, - block_table: pl.Tensor, - slot_mapping: pl.Tensor, - quant_k_cache: pl.Tensor, - quant_v_cache: pl.Tensor, - quant_k_scales: pl.Tensor, - quant_v_scales: pl.Tensor, - rot_matrices: pl.Tensor, - tq_codebook: pl.Tensor, - wo: pl.Tensor, - post_rms_weight: pl.Tensor, - w_gate: pl.Tensor, - w_up: pl.Tensor, - w_down: pl.Tensor, - final_norm_weight: pl.Tensor, - lm_head_weight: pl.Tensor, - out: pl.Out[pl.Tensor], -) -> pl.Tensor: - return prefill_fwd_tq( - hidden_states, - seq_lens, - input_rms_weight, - wq, - wk, - wv, - q_norm_weight, - k_norm_weight, - rope_cos, - rope_sin, - block_table, - slot_mapping, - quant_k_cache, - quant_v_cache, - quant_k_scales, - quant_v_scales, - rot_matrices, - tq_codebook, - wo, - post_rms_weight, - w_gate, - w_up, - w_down, - final_norm_weight, - lm_head_weight, - out, - ) - - @pl.jit.host def qwen3_decode_tq_host( hidden_states: pl.Tensor, From e555a4524d713374790fe3e0574f760c771b0402 Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 17:21:32 +0800 Subject: [PATCH 11/17] npu_executor: fix TQ prefill/decode dummy_args to match host wrapper sig MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TQ kernel compile failed with "missing inferred tensor metadata for parameter 'quant_indices'". The device kernel slices [CMP_CHUNK, head_dim//2] uint8 indices out of quant_k_cache, but the dummy_args passed quant_k_cache/quant_v_cache as BF16, so the sliced quant_indices shape could not be inferred. Root cause: both _compile_prefill_fwd_tq_callable and _compile_decode_fwd_tq_callable had mangled dummy_args copied from the non-TQ layout — duplicate seq_lens / chunk args, BF16 quant caches, no scales, and rot_matrices/codebook appended at the end instead of in signature order. Counts were wrong (prefill 27, decode 24 vs 26 params). Rebuilt both dummy_args to match the qwen3_prefill_tq_host / qwen3_decode_tq_host signatures exactly (26 args each): UINT8 quantized K/V caches, (cache_rows, 1) FP32 scales, (num_layers*head_dim, head_dim) BF16 rot_matrices, (1, 16) FP32 codebook, in wrapper param order. Co-Authored-By: Claude --- .../model/qwen3_14b/runner/npu_executor.py | 111 ++++++++++-------- 1 file changed, 61 insertions(+), 50 deletions(-) diff --git a/examples/model/qwen3_14b/runner/npu_executor.py b/examples/model/qwen3_14b/runner/npu_executor.py index d6f3453..ca0be72 100644 --- a/examples/model/qwen3_14b/runner/npu_executor.py +++ b/examples/model/qwen3_14b/runner/npu_executor.py @@ -827,33 +827,38 @@ def _compile_prefill_fwd_tq_callable( runtime_cache_blocks = (max_seq + page_size - 1) // page_size cache_rows = batch * runtime_cache_blocks * num_layers * num_kv_heads * page_size rot_rows = num_layers * head_dim + # Must match qwen3_prefill_tq_host param order exactly (26 params: + # 25 in + 1 out). TQ quantized caches are UINT8 with per-row FP32 + # scales; the device kernel slices [CMP_CHUNK, head_dim//2] uint8 + # indices out of quant_k_cache, so the dtype/shape here must be uint8 + # or the specializer cannot infer the sliced quant_indices tensor. dummy_args = [ - torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16), - torch.empty((batch,), dtype=torch.int32), - torch.empty((batch,), dtype=torch.int32), - torch.empty((batch,), dtype=torch.int32), - torch.empty((num_layers, hidden_size), dtype=torch.float32), - torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), - torch.empty((num_layers, head_dim), dtype=torch.float32), - torch.empty((num_layers, head_dim), dtype=torch.float32), - torch.empty((max_seq, head_dim), dtype=torch.float32), - torch.empty((max_seq, head_dim), dtype=torch.float32), - torch.empty((batch * block_table_stride,), dtype=torch.int32), - torch.empty((total_tokens,), dtype=torch.int32), - torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), - torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), - torch.empty((num_layers, hidden_size), dtype=torch.float32), - torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), - torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), - torch.empty((1, hidden_size), dtype=torch.float32), - torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), - torch.empty((batch, vocab_size), dtype=torch.float32), - torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), - torch.empty((1, 16), dtype=torch.float32), + torch.empty((total_tokens, hidden_size), dtype=torch.bfloat16), # hidden_states + torch.empty((batch,), dtype=torch.int32), # seq_lens + torch.empty((num_layers, hidden_size), dtype=torch.float32), # input_rms_weight + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wq + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wk + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wv + torch.empty((num_layers, head_dim), dtype=torch.float32), # q_norm_weight + torch.empty((num_layers, head_dim), dtype=torch.float32), # k_norm_weight + torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_cos + torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_sin + torch.empty((batch * block_table_stride,), dtype=torch.int32), # block_table + torch.empty((total_tokens,), dtype=torch.int32), # slot_mapping + torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_k_cache + torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_v_cache + torch.empty((cache_rows, 1), dtype=torch.float32), # quant_k_scales + torch.empty((cache_rows, 1), dtype=torch.float32), # quant_v_scales + torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), # rot_matrices + torch.empty((1, 16), dtype=torch.float32), # tq_codebook + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wo + torch.empty((num_layers, hidden_size), dtype=torch.float32), # post_rms_weight + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_gate + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_up + torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), # w_down + torch.empty((1, hidden_size), dtype=torch.float32), # final_norm_weight + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), # lm_head_weight + torch.empty((batch, vocab_size), dtype=torch.float32), # out ] return self._compile_jit_fwd_callable("prefill_fwd_tq", jit_fn, dummy_args) @@ -878,30 +883,36 @@ def _compile_decode_fwd_tq_callable( runtime_cache_blocks = (max_seq + page_size - 1) // page_size cache_rows = num_layers * batch * runtime_cache_blocks * num_kv_heads * page_size rot_rows = num_layers * head_dim + # Must match qwen3_decode_tq_host param order exactly (26 params: + # 25 in + 1 out). Same TQ dtype/shape rules as prefill: uint8 quantized + # caches + (cache_rows, 1) FP32 scales, so the specializer can infer the + # sliced quant_indices tensor. dummy_args = [ - torch.empty((batch, hidden_size), dtype=torch.bfloat16), - torch.empty((num_layers, hidden_size), dtype=torch.float32), - torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), - torch.empty((num_layers, head_dim), dtype=torch.float32), - torch.empty((num_layers, head_dim), dtype=torch.float32), - torch.empty((batch,), dtype=torch.int32), - torch.empty((batch * block_table_stride,), dtype=torch.int32), - torch.empty((batch,), dtype=torch.int32), - torch.empty((max_seq, head_dim), dtype=torch.float32), - torch.empty((max_seq, head_dim), dtype=torch.float32), - torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), - torch.empty((cache_rows, head_dim), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), - torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), - torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), - torch.empty((num_layers, hidden_size), dtype=torch.float32), - torch.empty((1, hidden_size), dtype=torch.float32), - torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), - torch.empty((batch, vocab_size), dtype=torch.float32), - torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), - torch.empty((1, 16), dtype=torch.float32), + torch.empty((batch, hidden_size), dtype=torch.bfloat16), # hidden_states + torch.empty((num_layers, hidden_size), dtype=torch.float32), # input_rms_weight + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wq + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wk + torch.empty((num_layers * hidden_size, kv_hidden), dtype=torch.bfloat16), # wv + torch.empty((num_layers, head_dim), dtype=torch.float32), # q_norm_weight + torch.empty((num_layers, head_dim), dtype=torch.float32), # k_norm_weight + torch.empty((batch,), dtype=torch.int32), # seq_lens + torch.empty((batch * block_table_stride,), dtype=torch.int32), # block_table + torch.empty((batch,), dtype=torch.int32), # slot_mapping + torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_cos + torch.empty((max_seq, head_dim), dtype=torch.float32), # rope_sin + torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_k_cache + torch.empty((cache_rows, head_dim), dtype=torch.uint8), # quant_v_cache + torch.empty((cache_rows, 1), dtype=torch.float32), # quant_k_scales + torch.empty((cache_rows, 1), dtype=torch.float32), # quant_v_scales + torch.empty((rot_rows, head_dim), dtype=torch.bfloat16), # rot_matrices + torch.empty((1, 16), dtype=torch.float32), # tq_codebook + torch.empty((num_layers * hidden_size, hidden_size), dtype=torch.bfloat16), # wo + torch.empty((num_layers, hidden_size), dtype=torch.float32), # post_rms_weight + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_gate + torch.empty((num_layers * hidden_size, intermediate_size), dtype=torch.bfloat16), # w_up + torch.empty((num_layers * intermediate_size, hidden_size), dtype=torch.bfloat16), # w_down + torch.empty((1, hidden_size), dtype=torch.float32), # final_norm_weight + torch.empty((vocab_size, hidden_size), dtype=torch.bfloat16), # lm_head_weight + torch.empty((batch, vocab_size), dtype=torch.float32), # out ] return self._compile_jit_fwd_callable("decode_fwd_tq", jit_fn, dummy_args) From d0879009b05f6d1142835314cc0efc9599b2492c Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 17:28:21 +0800 Subject: [PATCH 12/17] npu_runner: skip device sampling kernels in TQ worker; gate supports_device_* MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit TQ mode crashed building the L3 worker: AttributeError: 'NoneType' object has no attribute 'compiled' at self._compiled.greedy_sample.compiled — TQ sets greedy_sample/token_embed to None (it samples on CPU, not via the device greedy_sample/token_embed kernels). _shared_l3_worker now builds the program list conditionally, including greedy_sample/token_embed only when present. worker.run dispatches by compiled-program reference, so omitting them does not shift prefill/decode indexing. Also gate Qwen314BPyptoExecutor.supports_device_sampling / supports_device_embedding on `not self._tq_mode`. They unconditionally returned True, which would set allow_device_greedy_sampling=True in TQ mode and crash _maybe_run_sample_embed (compiled.greedy_sample is None). TQ must report no device sampling/embedding so the engine falls back to CPU sampling. Co-Authored-By: Claude --- examples/model/qwen3_14b/runner/npu_executor.py | 15 +++++++++++---- examples/model/qwen3_14b/runner/npu_runner.py | 14 ++++++++++---- 2 files changed, 21 insertions(+), 8 deletions(-) diff --git a/examples/model/qwen3_14b/runner/npu_executor.py b/examples/model/qwen3_14b/runner/npu_executor.py index ca0be72..c3180b4 100644 --- a/examples/model/qwen3_14b/runner/npu_executor.py +++ b/examples/model/qwen3_14b/runner/npu_executor.py @@ -126,13 +126,20 @@ def profile_verbose(self) -> bool: @property def supports_device_sampling(self) -> bool: - """Qwen3 NPU runner can return greedy sampled token ids.""" - return True + """Qwen3 NPU runner can return greedy sampled token ids. + + Non-TQ only: TurboQuant does not compile the device sampling/embedding + kernels (greedy_sample/token_embed), it samples on the CPU instead. + """ + return not self._tq_mode @property def supports_device_embedding(self) -> bool: - """Qwen3 NPU decode embeds greedy token ids inside the device kernel.""" - return True + """Qwen3 NPU decode embeds greedy token ids inside the device kernel. + + Non-TQ only (see supports_device_sampling). + """ + return not self._tq_mode def _create_runner(self, model_id: str, compiled: object) -> ModelRunner: """Create the Qwen3-14B runtime runner for compiled kernels.""" diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index 5ab7344..d807c35 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -896,12 +896,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 From 07fc464844fd668d046756952c99b06249b7951b Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 18:26:46 +0800 Subject: [PATCH 13/17] npu_runner: route TQ kernel args through quant pool; add rot/codebook statics Warmup prefill crashed with "DistributedWorker expects 26 arguments, got 24": _prefill_kernel_args/_decode_kernel_args built the non-TQ 24/26-arg layout (k_cache/v_cache, chunk_lens/offsets, in-kernel embed/sample buffers), but the TQ wrappers expect the TQ layout (uint8 quant caches + scales + rot_matrices + codebook, no chunk/embed/sample buffers). - _prefill_kernel_args / _decode_kernel_args now take the kv_cache pool and branch on tq_mode, emitting the 26-arg TQ layout (quant_k/v_pages, k/v_scales_pages, static.rot_matrices, static.tq_codebook; weight order wo, post_rms, w_gate, w_up, w_down) or the original non-TQ layout. - _build_static_kernel_args now wraps compiled.rot_matrices/tq_codebook as _StaticDeviceTensor so they upload via _coerce_l3_arg. - _validate_kv_cache_bounds uses quant_k_pages for TQ (key_pages is None in the TQ pool). - Updated all 4 call sites (run_prefill, run_decode, warmup prefill/decode) to pass the pool. Co-Authored-By: Claude --- examples/model/qwen3_14b/runner/npu_runner.py | 110 ++++++++++++++---- 1 file changed, 90 insertions(+), 20 deletions(-) diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index d807c35..c374b13 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -449,7 +449,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, ), ) @@ -479,7 +479,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)") @@ -579,6 +579,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: @@ -595,13 +605,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): @@ -644,18 +653,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())) @@ -725,13 +733,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, @@ -747,8 +786,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"], @@ -762,12 +801,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"], @@ -781,8 +851,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"], From b66dcb9301435d40d3db95f5af2496ea5b9a7c3e Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Tue, 7 Jul 2026 18:49:04 +0800 Subject: [PATCH 14/17] npu_runner: TQ-aware KV cache bytes_per_page (uint8 + fp32 scales) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _compute_kv_cache_pages / _print_memory_breakdown / _alloc_kv_cache_with_retry all sized the KV cache assuming runtime.kv_dtype (bf16, 2 bytes/elem), but TurboQuant allocates uint8 quantized K/V (1 byte/elem) + one fp32 scale per K and per V row. The sizing therefore over-counted TQ's footprint ~2x, requesting far fewer pages than the budget allows. Add _kv_bytes_per_page(config, runtime) returning the real per-page HBM: TQ = layers * kv_heads * page_size * (2*head_dim + 2*4) # uint8 + scales bf16= layers * kv_heads * page_size * 2 * head_dim * dtype_bytes For Qwen3-14B this is 10.8 MB/page (TQ) vs 21.0 MB/page (bf16) — TQ now correctly sizes ~1.94x more pages for the same budget. Co-Authored-By: Claude --- examples/model/qwen3_14b/runner/npu_runner.py | 39 +++++++++++-------- 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/examples/model/qwen3_14b/runner/npu_runner.py b/examples/model/qwen3_14b/runner/npu_runner.py index c374b13..b26469a 100644 --- a/examples/model/qwen3_14b/runner/npu_runner.py +++ b/examples/model/qwen3_14b/runner/npu_runner.py @@ -254,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): " @@ -279,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. @@ -291,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 @@ -341,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). From 6fcea159aa34928774bd77c9705709cc2613eecd Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Wed, 8 Jul 2026 01:28:14 -0700 Subject: [PATCH 15/17] qwen3_l3_dispatch: fix prefill_host body post_rms_weight arg order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The non-TQ qwen3_prefill_host body passed args to prefill_fwd in the wrong weight order: wo, w_gate, w_up, w_down, post_rms_weight — but the prefill_fwd kernel signature is wo, post_rms_weight, w_gate, w_up, w_down. The misaligned args made the kernel receive w_gate (BF16) in the post_rms_weight (FP32) position, so the kernel's `gamma = pl.slice(post_rms_weight)` came out BF16 while line 174's `gamma = pl.slice(input_rms_weight)` was FP32 — pypto's SSA checker then raised "Cannot reassign 'gamma' with a different type", blocking all non-TQ prefill compilation. Reorder the body to wo, post_rms_weight, w_gate, w_up, w_down, matching the prefill_fwd kernel (and the upstream main repo, which runs non-TQ fine). The wrapper signature and dummy_args are left as-is (they are internally consistent and the body reorders for the kernel); decode_host was already correct (decode kernel has post_rms after w_down). Co-Authored-By: Claude --- examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py index ef03fe0..fb25a93 100644 --- a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py +++ b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py @@ -67,10 +67,10 @@ def qwen3_prefill_host( k_cache, v_cache, wo, + post_rms_weight, w_gate, w_up, w_down, - post_rms_weight, final_norm_weight, lm_head_weight, out, From 2856ac0a7ac78ffe92b1416b87948f3c05a46270 Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Wed, 8 Jul 2026 01:43:23 -0700 Subject: [PATCH 16/17] qwen3_l3_dispatch: bring qwen3_decode_host to 26-param device-sampling sig Non-TQ decode compilation failed with "too many positional arguments": the decode dummy_args (26) and decode_fwd kernel (26, PR#47 device sampling) had embed_weight/sampled_ids_in/sampled_ids/next_hidden, but qwen3_decode_host was still the old 22-param wrapper. Update it to the 26-param device-sampling signature matching the upstream main repo (and the kernel + dummy + the non-TQ _decode_kernel_args dispatch, which already passed 26 args). Co-Authored-By: Claude --- .../model/qwen3_14b/runner/qwen3_l3_dispatch.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py index fb25a93..4a4f98e 100644 --- a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py +++ b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py @@ -101,8 +101,12 @@ def qwen3_decode_host( final_norm_weight: pl.Tensor, lm_head_weight: pl.Tensor, out: pl.Out[pl.Tensor], -) -> pl.Tensor: - return decode_fwd( + embed_weight: pl.Tensor, + sampled_ids_in: pl.Tensor, + sampled_ids: pl.Out[pl.Tensor], + next_hidden: pl.Out[pl.Tensor], +) -> tuple[pl.Tensor, pl.Tensor, pl.Tensor]: + logits, sampled_ids, next_hidden = decode_fwd( hidden_states, input_rms_weight, wq, @@ -125,7 +129,12 @@ def qwen3_decode_host( final_norm_weight, lm_head_weight, out, + embed_weight, + sampled_ids_in, + sampled_ids, + next_hidden, ) + return logits, sampled_ids, next_hidden @pl.jit.host From 52125a7d7604a3d961dc50fe85489ddfcdbbfd26 Mon Sep 17 00:00:00 2001 From: sunghajung6688 <2901782080@qq.com> Date: Wed, 8 Jul 2026 01:49:00 -0700 Subject: [PATCH 17/17] qwen3_l3_dispatch: add greedy_sample/token_embed host wrappers + bindings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The non-TQ compile also builds greedy_sample/token_embed (for the prefill first-token device sampling via _maybe_run_sample_embed), referencing qwen3_greedy_sample_host / qwen3_token_embed_host — but those wrappers and the greedy_sample_fwd / token_embed_fwd module bindings were absent, so it would AttributeError after decode. Add them, matching the upstream main repo. Co-Authored-By: Claude --- .../qwen3_14b/runner/qwen3_l3_dispatch.py | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py index 4a4f98e..b840927 100644 --- a/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py +++ b/examples/model/qwen3_14b/runner/qwen3_l3_dispatch.py @@ -16,6 +16,10 @@ prefill_fwd = None decode_fwd = None +# Device-side sampling + embedding kernels (assigned by the executor before +# compilation). Used by _maybe_run_sample_embed for the prefill first token. +greedy_sample_fwd = None +token_embed_fwd = None # TurboQuant (TQ) device kernels; assigned by the executor before compilation # when tq_mode is enabled. The TQ wrappers below forward to these. prefill_fwd_tq = None @@ -137,6 +141,30 @@ def qwen3_decode_host( return logits, sampled_ids, next_hidden +@pl.jit.host +def qwen3_greedy_sample_host( + logits: pl.Tensor, + sampled_ids: pl.Out[pl.Tensor], +) -> pl.Tensor: + return greedy_sample_fwd( + logits, + sampled_ids, + ) + + +@pl.jit.host +def qwen3_token_embed_host( + sampled_ids: pl.Tensor, + embed_weight: pl.Tensor, + next_hidden: pl.Out[pl.Tensor], +) -> pl.Tensor: + return token_embed_fwd( + sampled_ids, + embed_weight, + next_hidden, + ) + + @pl.jit.host def qwen3_prefill_tq_host( hidden_states: pl.Tensor,