From cd7459aa45e4f1ae07a352f719d2306a095884b7 Mon Sep 17 00:00:00 2001 From: probe Date: Sun, 23 Aug 2026 08:59:16 +0400 Subject: [PATCH 1/2] feat(kvcache): 8-bit KV cache (q8_0 / fp8_e4m3) behind --kv-cache-dtype Stores the KV cache in 8 bits plus an fp16 scale per 32 elements along head_dim (1.0625 bytes/element vs 2), freeing VRAM for the MoE expert cache. Two schemes share the scale tensor, store kernel and dequant path -- q8_0 (int8, s = max/127) and fp8_e4m3 (s = max/448) -- so comparing them is a flag change, not a port. - kvcache/quant.py: KVQuantSpec (storage dtype, block 32, torch reference quantize/dequantize, effective bytes/element) - kernel/triton/kv_quant.py: store kernel computing per-block max-abs and writing the quantized buffer + scales - kernel/triton/attention.py: dequant inside the four attention kernels behind a QUANT constexpr (0 compiles the existing bf16 path unchanged); the scale varies along head_dim, the reduction dim, so K/V dequantize to bf16 before the dot - kvcache pools: parallel scale buffers, k_scale()/v_scale(), rebuild() realloc, unit_bytes()/kv_cost() accounting on effective bytes - server/args.py, engine: --kv-cache-dtype {auto,q8_0,fp8_e4m3} with gating (triton backend only, head_dim % 32 == 0, supported pool families) Tests: 53 new (round-trip vs torch reference, quantized attention vs the bf16 reference, pool sizing and hot rebuild, flag gating); the existing 33 triton attention tests still pass. Step 9 of tasks/todo.md is NOT done: no end-to-end validation on this host -- needle-in-246k, perplexity vs bf16, and the real expert-slot / tok-s gain are unmeasured, so q8_0 vs fp8_e4m3 as the default is still an open question. Co-Authored-By: Claude Claude-Session: https://claude.ai/code/session_01Je2rjENB9qJiiRGmNcnAct --- python/freetoken/attention/triton.py | 13 + python/freetoken/engine/config.py | 10 + python/freetoken/engine/engine.py | 39 +++ python/freetoken/kernel/triton/attention.py | 318 ++++++++++++++++---- python/freetoken/kernel/triton/kv_quant.py | 130 ++++++++ python/freetoken/kvcache/__init__.py | 7 + python/freetoken/kvcache/base.py | 20 +- python/freetoken/kvcache/hybrid_swa_pool.py | 75 +++-- python/freetoken/kvcache/mha_pool.py | 62 ++-- python/freetoken/kvcache/quant.py | 133 ++++++++ python/freetoken/kvcache/quant_storage.py | 79 +++++ python/freetoken/server/args.py | 15 + tasks/todo.md | 7 + tests/engine/test_kv_cache_dtype_gating.py | 72 +++++ tests/kernels/test_kv_quant.py | 286 ++++++++++++++++++ tests/kvcache/test_kv_quant_pool.py | 229 ++++++++++++++ 16 files changed, 1398 insertions(+), 97 deletions(-) create mode 100644 python/freetoken/kernel/triton/kv_quant.py create mode 100644 python/freetoken/kvcache/quant.py create mode 100644 python/freetoken/kvcache/quant_storage.py create mode 100644 tasks/todo.md create mode 100644 tests/engine/test_kv_cache_dtype_gating.py create mode 100644 tests/kernels/test_kv_quant.py create mode 100644 tests/kvcache/test_kv_quant_pool.py diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 9eed1e1d..e9b96266 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -144,6 +144,7 @@ def forward( extend_paged_attention, paged_attention, ) + from freetoken.kvcache.quant import BLOCK as QBLOCK metadata = batch.attn_metadata assert isinstance(metadata, TritonMetadata) @@ -155,6 +156,12 @@ def forward( assert head_dim == q.shape[-1] k_cache = k_raw.view(-1, kv_heads, head_dim) v_cache = v_raw.view(-1, kv_heads, head_dim) + # An 8-bit pool hands its per-block scales alongside the slabs; an unquantized one + # has none, and the kernels compile their bf16 path unchanged. + k_scale = v_scale = None + if getattr(self.kvcache, "quant", None) is not None and self.kvcache.quant.enabled: + k_scale = self.kvcache.k_scale(layer_id).view(-1, kv_heads, head_dim // QBLOCK) + v_scale = self.kvcache.v_scale(layer_id).view(-1, kv_heads, head_dim // QBLOCK) spec = attn_spec or AttentionSpec() indices = metadata.indices @@ -181,6 +188,8 @@ def forward( sm_scale=scale, sliding_window=spec.sliding_window, sinks=spec.sinks, + k_scale=k_scale, + v_scale=v_scale, ) if ( (not metadata.is_decode) @@ -201,6 +210,8 @@ def forward( sinks=spec.sinks, k_extend=k.view(q.shape[0], kv_heads, head_dim), v_extend=v.view(q.shape[0], kv_heads, head_dim), + k_scale=k_scale, + v_scale=v_scale, ) return paged_attention( q=q, @@ -213,6 +224,8 @@ def forward( sm_scale=scale, sliding_window=spec.sliding_window, sinks=spec.sinks, + k_scale=k_scale, + v_scale=v_scale, ) def prepare_metadata(self, batch: Batch) -> None: diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f3..b6a847d2 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -80,6 +80,16 @@ class EngineConfig: # KV capacity in tokens; resolved into num_page_override by _adjust_config once page_size # is final. Mutually exclusive with num_page_override. num_token_override: int | None = None + # KV element storage (--kv-cache-dtype): "auto" keeps the compute dtype, "q8_0" and + # "fp8_e4m3" store 8 bits plus a per-block scale. Resolved through + # freetoken.kvcache.quant.resolve_kv_quant by the pools and the cost model. + kv_cache_dtype: str = "auto" + + @cached_property + def kv_quant(self): + from freetoken.kvcache.quant import resolve_kv_quant + + return resolve_kv_quant(self.kv_cache_dtype) @cached_property def hf_config(self): diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ff3c985f..86255bb8 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -154,6 +154,44 @@ def _resolve_auto_attention_backend( ) +def _validate_kv_cache_dtype(config, model_config) -> None: + """Gate --kv-cache-dtype against what the quantized path actually implements. + + 8-bit KV storage lives in the triton attention kernels and the MHA/hybrid-SWA pools. + Every other backend reads the KV slabs through its own kernels (flashinfer's + ``kv_data_type``, trtllm's fp8 path) which this has not been wired into, and the + MLA/DSA/DSV4/BSA pools have their own slab layouts. Reject those combinations here, + at config time, rather than letting a wrong-dtype tensor reach a kernel. + """ + quant = getattr(config, "kv_quant", None) + if quant is None or not quant.enabled: + return + + from freetoken.kvcache.quant import BLOCK + + backends = [p.strip() for p in config.attention_backend.split(",")] + if any(b != "triton" for b in backends): + raise ValueError( + f"--kv-cache-dtype {quant.name} needs the triton attention backend, but the " + f"resolved backend is {config.attention_backend!r}. Pass " + "--attention-backend triton, or drop --kv-cache-dtype." + ) + + specs = [s for s in model_config.kv_cache_group_specs() if s.num_layers > 0] + if any(s.mla or s.index_head_dim > 0 for s in specs): + raise ValueError( + f"--kv-cache-dtype {quant.name} does not support MLA/DSA latent KV pools " + "(their slabs alias K and V and carry an index tier); use --kv-cache-dtype auto." + ) + bad = [s for s in specs if s.head_dim % BLOCK] + if bad: + names = ", ".join(f"{s.name} (head_dim {s.head_dim})" for s in bad) + raise ValueError( + f"--kv-cache-dtype {quant.name} needs every head_dim to be a multiple of " + f"{BLOCK}, the quantization block; this model has {names}." + ) + + def _validate_attention_backend_choice(config, override, required: frozenset[AttnType]) -> None: """Config-time type x backend capability check for the resolved (or explicit) backend string: every comma part must serve every required type and have its @@ -1190,6 +1228,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution ) logger.info_rank0(f"Auto-selected attention backend: {config.attention_backend}") _validate_attention_backend_choice(config, override, required_attn_types) + _validate_kv_cache_dtype(config, model_config) if config.moe_cache_rate is not None: total_experts = config.model_config.num_moe_layers * config.model_config.num_experts diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84..7fe525cb 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -11,6 +11,54 @@ _MIN_BLOCK_KV = 32 +@triton.jit +def _load_kv( + ptr, + scale_ptr, + offsets, + scale_offsets, + mask, + scale_mask, + out_dtype: tl.constexpr, + QUANT: tl.constexpr, + QBLOCK: tl.constexpr, + D_ON_ROWS: tl.constexpr, +): + """Load a K or V tile, dequantizing it when the pool stores 8-bit values. + + ``offsets`` addresses the tile in the KV buffer; ``scale_offsets`` addresses the + matching scales, whose ``head_dim`` extent is ``QBLOCK`` times smaller -- one scale + per block, loaded once and broadcast across the block rather than re-read per + element. Reading it per element instead costs 2 bytes of load per 1 byte of payload + and measured 1.6x slower than bf16 on prefill; this broadcast is what makes 8-bit + storage actually cheaper to read. + + ``D_ON_ROWS`` says which way the tile is laid out: the dot-product kernels want K as + ``[D, N]`` and V as ``[N, D]``, and the block axis has to be expanded along whichever + one carries ``head_dim``. + + The scale varies along ``head_dim``, the reduction dimension of ``q @ k``, so it + cannot be folded in after the dot -- the tile is dequantized into ``out_dtype`` (the + query's dtype) and fed to the tensor cores like the bf16 path does. + """ + vals = tl.load(ptr + offsets, mask=mask, other=0.0) + if QUANT: + scale = tl.load(scale_ptr + scale_offsets, mask=scale_mask, other=0.0) + if D_ON_ROWS: + nb: tl.constexpr = scale.shape[0] + n: tl.constexpr = scale.shape[1] + wide = tl.broadcast_to(scale[:, None, :], (nb, QBLOCK, n)).reshape(nb * QBLOCK, n) + else: + n: tl.constexpr = scale.shape[0] + nb: tl.constexpr = scale.shape[1] + wide = tl.broadcast_to(scale[:, :, None], (n, nb, QBLOCK)).reshape(n, nb * QBLOCK) + return (vals.to(tl.float32) * wide.to(tl.float32)).to(out_dtype) + # Both branches must yield the same type for Triton to compile the function, so the + # unquantized path casts too. Callers pass the dtype the tile already has there + # (float32 for the fp32 kernel, the cache's own dtype elsewhere), so it is a no-op. + return vals.to(out_dtype) + + @functools.lru_cache(maxsize=None) def _optin_smem_bytes(device_index: int) -> int: """Per-block opt-in shared-memory budget for a CUDA device (0 if unavailable).""" @@ -47,6 +95,8 @@ def _paged_attention_kernel( q_ptr, k_ptr, v_ptr, + ks_ptr, + vs_ptr, o_ptr, indptr_ptr, indices_ptr, @@ -60,6 +110,10 @@ def _paged_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_ksh, + stride_vss, + stride_vsh, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -68,6 +122,8 @@ def _paged_attention_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + QUANT: tl.constexpr, + QBLOCK: tl.constexpr, ): q_tok = tl.program_id(0) q_head = tl.program_id(1) @@ -81,6 +137,9 @@ def _paged_attention_kernel( offs_d = tl.arange(0, BLOCK_D) mask_d = offs_d < D + # One scale per QBLOCK elements of head_dim: the tile's block axis. + offs_nb = tl.arange(0, BLOCK_D // QBLOCK) + mask_nb = offs_nb < D // QBLOCK q = tl.load( q_ptr + q_tok * stride_qt + q_head * stride_qh + offs_d, mask=mask_d, @@ -107,13 +166,19 @@ def _paged_attention_kernel( skip_tile = tl.max(mask_n.to(tl.int32), axis=0) == 0 if not skip_tile: slots = tl.load(indices_ptr + kv_start + offs_n, mask=offs_n < kv_len, other=0) - k = tl.load( - k_ptr - + slots[:, None] * stride_ks - + kv_head * stride_kh - + offs_d[None, :], - mask=(offs_n[:, None] < kv_len) & mask_d[None, :], - other=0.0, + kv_mask = (offs_n[:, None] < kv_len) & mask_d[None, :] + kv_scale_mask = (offs_n[:, None] < kv_len) & mask_nb[None, :] + k = _load_kv( + k_ptr, + ks_ptr, + slots[:, None] * stride_ks + kv_head * stride_kh + offs_d[None, :], + slots[:, None] * stride_kss + kv_head * stride_ksh + offs_nb[None, :], + kv_mask, + kv_scale_mask, + tl.float32, + QUANT, + QBLOCK, + False, ).to(tl.float32) scores = tl.sum(q[None, :] * k, axis=1) * sm_scale scores = tl.where(mask_n, scores, -float("inf")) @@ -124,13 +189,17 @@ def _paged_attention_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new) - v = tl.load( - v_ptr - + slots[:, None] * stride_vs - + kv_head * stride_vh - + offs_d[None, :], - mask=(offs_n[:, None] < kv_len) & mask_d[None, :], - other=0.0, + v = _load_kv( + v_ptr, + vs_ptr, + slots[:, None] * stride_vs + kv_head * stride_vh + offs_d[None, :], + slots[:, None] * stride_vss + kv_head * stride_vsh + offs_nb[None, :], + kv_mask, + kv_scale_mask, + tl.float32, + QUANT, + QBLOCK, + False, ).to(tl.float32) acc = acc * alpha + tl.sum(p[:, None] * v, axis=0) l_i = l_i * alpha + tl.sum(p, axis=0) @@ -149,6 +218,8 @@ def _decode_grouped_stage1_kernel( q_ptr, k_ptr, v_ptr, + ks_ptr, + vs_ptr, sm_scale, indptr_ptr, indices_ptr, @@ -162,6 +233,10 @@ def _decode_grouped_stage1_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_ksh, + stride_vss, + stride_vsh, stride_mid_ob, stride_mid_oh, stride_mid_os, @@ -179,6 +254,8 @@ def _decode_grouped_stage1_kernel( D: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, + QUANT: tl.constexpr, + QBLOCK: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -197,6 +274,10 @@ def _decode_grouped_stage1_kernel( offs_dv = tl.arange(0, BLOCK_DV) mask_d = offs_d < D mask_dv = offs_dv < DV + offs_nb = tl.arange(0, BLOCK_D // QBLOCK) + offs_nbv = tl.arange(0, BLOCK_DV // QBLOCK) + mask_nb = offs_nb < D // QBLOCK + mask_nbv = offs_nbv < DV // QBLOCK kv_start = tl.load(indptr_ptr + batch_id) kv_len = tl.load(indptr_ptr + batch_id + 1) - kv_start @@ -221,10 +302,16 @@ def _decode_grouped_stage1_kernel( q_offsets = batch_id * stride_qt + q_heads[:, None] * stride_qh + offs_d[None, :] k_base_offsets = kv_head * stride_kh + offs_d[:, None] v_base_offsets = kv_head * stride_vh + offs_dv[None, :] + ks_base_offsets = kv_head * stride_ksh + offs_nb[:, None] + vs_base_offsets = kv_head * stride_vsh + offs_nbv[None, :] if split_end > split_start: q = tl.load(q_ptr + q_offsets, mask=mask_h[:, None] & mask_d[None, :], other=0.0) - q = q.to(k_ptr.dtype.element_ty) + if not QUANT: + # Unquantized: match the cache's dtype as before. Quantized: the cache is + # int8/fp8 and casting q into it would destroy the query -- the dequantized + # K/V tiles are produced in q's dtype instead. + q = q.to(k_ptr.dtype.element_ty) for rel_start in tl.range(split_start, split_end, BLOCK_N): rel_offs = rel_start + tl.arange(0, BLOCK_N) @@ -232,18 +319,32 @@ def _decode_grouped_stage1_kernel( logical_offs = effective_start + rel_offs slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=0) - k = tl.load( - k_ptr + slots[None, :] * stride_ks + k_base_offsets, - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, + k = _load_kv( + k_ptr, + ks_ptr, + slots[None, :] * stride_ks + k_base_offsets, + slots[None, :] * stride_kss + ks_base_offsets, + mask_n[None, :] & mask_d[:, None], + mask_n[None, :] & mask_nb[:, None], + q.dtype, + QUANT, + QBLOCK, + True, ) scores = tl.dot(q, k) * sm_scale scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) - v = tl.load( - v_ptr + slots[:, None] * stride_vs + v_base_offsets, - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, + v = _load_kv( + v_ptr, + vs_ptr, + slots[:, None] * stride_vs + v_base_offsets, + slots[:, None] * stride_vss + vs_base_offsets, + mask_n[:, None] & mask_dv[None, :], + mask_n[:, None] & mask_nbv[None, :], + q.dtype, + QUANT, + QBLOCK, + False, ) m_new = tl.maximum(tl.max(scores, axis=1), m_i) @@ -349,6 +450,34 @@ def _decode_stage2_kernel( ) +def _kv_scale_args(k_cache, v_cache, k_scale, v_scale): + """Scale tensors + strides + the QUANT/QBLOCK constexprs for a kernel launch. + + Unquantized pools pass ``k_scale=None``; the kernels then never touch the scale + pointers, so the KV buffers themselves stand in and ``QUANT=False`` compiles the + dequant away entirely -- the bf16 path emits the same code it did before. + """ + if k_scale is None: + assert v_scale is None, "k_scale and v_scale must be given together" + return (k_cache, v_cache, 0, 0, 0, 0, False, 1) + assert v_scale is not None, "k_scale and v_scale must be given together" + assert k_scale.dim() == v_scale.dim() == 3, "scales are [slots, heads, D // block]" + block = k_cache.shape[-1] // k_scale.shape[-1] + assert k_cache.shape[-1] == block * k_scale.shape[-1], ( + f"head_dim {k_cache.shape[-1]} is not a whole number of {k_scale.shape[-1]} blocks" + ) + return ( + k_scale, + v_scale, + k_scale.stride(0), + k_scale.stride(1), + v_scale.stride(0), + v_scale.stride(1), + True, + block, + ) + + def decode_paged_attention( q: torch.Tensor, k_cache: torch.Tensor, @@ -364,11 +493,16 @@ def decode_paged_attention( sliding_window: int | None = None, sinks: torch.Tensor | None = None, out: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, ) -> torch.Tensor: """SGLang-style split-k grouped decode attention for one query per request.""" assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3 + ks, vs, s_kss, s_ksh, s_vss, s_vsh, quant, qblock = _kv_scale_args( + k_cache, v_cache, k_scale, v_scale + ) batch, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] assert batch == indptr.numel() - 1 @@ -405,6 +539,8 @@ def decode_paged_attention( q, k_cache, v_cache, + ks, + vs, sm_scale, indptr, indices, @@ -418,6 +554,10 @@ def decode_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + s_kss, + s_ksh, + s_vss, + s_vsh, attn_logits.stride(0), attn_logits.stride(1), attn_logits.stride(2), @@ -435,6 +575,8 @@ def decode_paged_attention( D=head_dim, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, + QUANT=quant, + QBLOCK=qblock, num_warps=4, num_stages=2, ) @@ -471,6 +613,8 @@ def _extend_attention_kernel( q_ptr, k_ptr, v_ptr, + ks_ptr, + vs_ptr, o_ptr, qo_indptr_ptr, kv_indptr_ptr, @@ -484,6 +628,10 @@ def _extend_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_ksh, + stride_vss, + stride_vsh, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -494,6 +642,8 @@ def _extend_attention_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + QUANT: tl.constexpr, + QBLOCK: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -514,6 +664,10 @@ def _extend_attention_kernel( mask_m = offs_m < q_len mask_d = offs_d < D mask_dv = offs_dv < D + offs_nb = tl.arange(0, BLOCK_D // QBLOCK) + offs_nbv = tl.arange(0, BLOCK_DV // QBLOCK) + mask_nb = offs_nb < D // QBLOCK + mask_nbv = offs_nbv < D // QBLOCK q_abs_pos = prefix_len + offs_m block_q_end = tl.minimum(q_len, (block_m_id + 1) * BLOCK_M) kv_loop_end = tl.minimum(kv_len, prefix_len + block_q_end) @@ -545,13 +699,17 @@ def _extend_attention_kernel( skip_tile = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0 if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) - k = tl.load( - k_ptr - + slots[None, :] * stride_ks - + kv_head * stride_kh - + offs_d[:, None], - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, + k = _load_kv( + k_ptr, + ks_ptr, + slots[None, :] * stride_ks + kv_head * stride_kh + offs_d[:, None], + slots[None, :] * stride_kss + kv_head * stride_ksh + offs_nb[:, None], + mask_n[None, :] & mask_d[:, None], + mask_n[None, :] & mask_nb[:, None], + q.dtype, + QUANT, + QBLOCK, + True, ) scores = tl.dot(q.to(k.dtype), k) * sm_scale scores = tl.where(final_mask, scores, -float("inf")) @@ -562,13 +720,17 @@ def _extend_attention_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - v = tl.load( - v_ptr - + slots[:, None] * stride_vs - + kv_head * stride_vh - + offs_dv[None, :], - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, + v = _load_kv( + v_ptr, + vs_ptr, + slots[:, None] * stride_vs + kv_head * stride_vh + offs_dv[None, :], + slots[:, None] * stride_vss + kv_head * stride_vsh + offs_nbv[None, :], + mask_n[:, None] & mask_dv[None, :], + mask_n[:, None] & mask_nbv[None, :], + q.dtype, + QUANT, + QBLOCK, + False, ) acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) @@ -592,6 +754,8 @@ def _extend_attention_split_kernel( v_extend_ptr, k_cache_ptr, v_cache_ptr, + ks_ptr, + vs_ptr, o_ptr, qo_indptr_ptr, kv_indptr_ptr, @@ -609,6 +773,10 @@ def _extend_attention_split_kernel( stride_kch, stride_vcs, stride_vch, + stride_kss, + stride_ksh, + stride_vss, + stride_vsh, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -619,6 +787,8 @@ def _extend_attention_split_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + QUANT: tl.constexpr, + QBLOCK: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -637,6 +807,10 @@ def _extend_attention_split_kernel( mask_m = offs_m < q_len mask_d = offs_d < D mask_dv = offs_dv < D + offs_nb = tl.arange(0, BLOCK_D // QBLOCK) + offs_nbv = tl.arange(0, BLOCK_DV // QBLOCK) + mask_nb = offs_nb < D // QBLOCK + mask_nbv = offs_nbv < D // QBLOCK q_abs_pos = prefix_len + offs_m q = tl.load( @@ -672,13 +846,17 @@ def _extend_attention_split_kernel( if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) - k = tl.load( - k_cache_ptr - + slots[None, :] * stride_kcs - + kv_head * stride_kch - + offs_d[:, None], - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, + k = _load_kv( + k_cache_ptr, + ks_ptr, + slots[None, :] * stride_kcs + kv_head * stride_kch + offs_d[:, None], + slots[None, :] * stride_kss + kv_head * stride_ksh + offs_nb[:, None], + mask_n[None, :] & mask_d[:, None], + mask_n[None, :] & mask_nb[:, None], + q.dtype, + QUANT, + QBLOCK, + True, ) scores = tl.dot(q.to(k.dtype), k) * sm_scale scores = tl.where(final_mask, scores, -float("inf")) @@ -689,13 +867,17 @@ def _extend_attention_split_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - v = tl.load( - v_cache_ptr - + slots[:, None] * stride_vcs - + kv_head * stride_vch - + offs_dv[None, :], - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, + v = _load_kv( + v_cache_ptr, + vs_ptr, + slots[:, None] * stride_vcs + kv_head * stride_vch + offs_dv[None, :], + slots[:, None] * stride_vss + kv_head * stride_vsh + offs_nbv[None, :], + mask_n[:, None] & mask_dv[None, :], + mask_n[:, None] & mask_nbv[None, :], + q.dtype, + QUANT, + QBLOCK, + False, ) acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) @@ -773,11 +955,16 @@ def extend_paged_attention( out: torch.Tensor | None = None, k_extend: torch.Tensor | None = None, v_extend: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Block-tiled causal prefill/extend attention over paged KV cache.""" assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3 + ks, vs, s_kss, s_ksh, s_vss, s_vsh, quant, qblock = _kv_scale_args( + k_cache, v_cache, k_scale, v_scale + ) num_q_tokens, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] assert qo_indptr.numel() == kv_indptr.numel() @@ -815,6 +1002,8 @@ def extend_paged_attention( v_extend, k_cache, v_cache, + ks, + vs, o, qo_indptr, kv_indptr, @@ -832,6 +1021,10 @@ def extend_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + s_kss, + s_ksh, + s_vss, + s_vsh, o.stride(0), o.stride(1), GROUP=num_q_heads // num_kv_heads, @@ -842,6 +1035,8 @@ def extend_paged_attention( BLOCK_N=block_n, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + QUANT=quant, + QBLOCK=qblock, num_warps=8, num_stages=1, ) @@ -851,6 +1046,8 @@ def extend_paged_attention( q, k_cache, v_cache, + ks, + vs, o, qo_indptr, kv_indptr, @@ -864,6 +1061,10 @@ def extend_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + s_kss, + s_ksh, + s_vss, + s_vsh, o.stride(0), o.stride(1), GROUP=num_q_heads // num_kv_heads, @@ -874,6 +1075,8 @@ def extend_paged_attention( BLOCK_N=block_n, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + QUANT=quant, + QBLOCK=qblock, num_warps=8, num_stages=1, ) @@ -893,6 +1096,8 @@ def paged_attention( sinks: torch.Tensor | None = None, out: torch.Tensor | None = None, block_n: int = 32, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Paged causal attention for one layer. @@ -903,6 +1108,9 @@ def paged_attention( assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3 + ks, vs, s_kss, s_ksh, s_vss, s_vsh, quant, qblock = _kv_scale_args( + k_cache, v_cache, k_scale, v_scale + ) num_tokens, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] assert v_cache.shape[1] == num_kv_heads @@ -922,6 +1130,8 @@ def paged_attention( q, k_cache, v_cache, + ks, + vs, o, indptr, indices, @@ -935,6 +1145,10 @@ def paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + s_kss, + s_ksh, + s_vss, + s_vsh, o.stride(0), o.stride(1), GROUP=num_q_heads // num_kv_heads, @@ -943,6 +1157,8 @@ def paged_attention( BLOCK_N=block_n, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + QUANT=quant, + QBLOCK=qblock, num_warps=8 if head_dim >= 256 else 4, num_stages=2, ) diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py new file mode 100644 index 00000000..f2748d38 --- /dev/null +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -0,0 +1,130 @@ +"""Quantizing store into an 8-bit KV pool. + +The unquantized path stores K/V with ``kernel/store.py``'s CUDA kernel, which is a +pure byte copy parameterized by element size. Quantized storage has to compute a scale +per block of :data:`~freetoken.kvcache.quant.BLOCK` elements along ``head_dim`` on the +way in, so it gets its own kernel here. + +One program handles one ``(token, kv_head)`` pair: it loads that head's ``head_dim`` +values as a ``[head_dim // BLOCK, BLOCK]`` tile, reduces max-abs along the block, and +writes the quantized values plus one scale per block. K and V are done in the same +program -- they share the token's slot index and the tile geometry, so doing both +halves the launch count and the index math. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _store_kv_quant_kernel( + k_ptr, # [tokens, heads, D] source, compute dtype + v_ptr, + kc_ptr, # [slots, heads, D] destination, storage dtype + vc_ptr, + ks_ptr, # [slots, heads, D // BLOCK] scales, fp16 + vs_ptr, + indices_ptr, # [tokens] destination slot per token + stride_kt, + stride_kh, + stride_ct, + stride_ch, + stride_st, + stride_sh, + MAX_MAG: tl.constexpr, + IS_INT: tl.constexpr, + BLOCK: tl.constexpr, + NBLOCK: tl.constexpr, +): + tok = tl.program_id(0) + head = tl.program_id(1) + slot = tl.load(indices_ptr + tok).to(tl.int64) + + # [NBLOCK, BLOCK] tile over head_dim: rows are quant blocks, columns the elements + # sharing one scale. + offs = tl.arange(0, NBLOCK)[:, None] * BLOCK + tl.arange(0, BLOCK)[None, :] + scale_offs = tl.arange(0, NBLOCK) + + for is_v in tl.static_range(2): + src_ptr = v_ptr if is_v else k_ptr + dst_ptr = vc_ptr if is_v else kc_ptr + sc_ptr = vs_ptr if is_v else ks_ptr + + x = tl.load(src_ptr + tok * stride_kt + head * stride_kh + offs).to(tl.float32) + amax = tl.max(tl.abs(x), axis=1) + # An all-zero block quantizes to zeros under any positive scale; 1.0 keeps the + # division finite. + scale = tl.where(amax > 0, amax / MAX_MAG, 1.0) + # Round to the stored precision before dividing, so the value written here and + # the value the attention kernels read back are scaled by the identical number. + scale = scale.to(sc_ptr.dtype.element_ty).to(tl.float32) + # div_rn, not `/`: the plain operator is free to lower to a reciprocal multiply, + # which disagrees with the torch reference on values sitting exactly between two + # quantization steps. IEEE round-to-nearest divide makes the two bit-identical. + q = tl.math.div_rn(x, scale[:, None]) + if IS_INT: + # Round half away from zero (what GGUF's Q8_0 does), then clamp -- the + # float->int cast truncates. + q = tl.where(q >= 0, tl.floor(q + 0.5), tl.ceil(q - 0.5)) + q = tl.minimum(tl.maximum(q, -MAX_MAG), MAX_MAG) + + tl.store( + dst_ptr + slot * stride_ct + head * stride_ch + offs, + q.to(dst_ptr.dtype.element_ty), + ) + tl.store( + sc_ptr + slot * stride_st + head * stride_sh + scale_offs, + scale.to(sc_ptr.dtype.element_ty), + ) + + +def store_kv_quant( + k_cache: torch.Tensor, + k_scale: torch.Tensor, + v_cache: torch.Tensor, + v_scale: torch.Tensor, + indices: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + spec, +) -> None: + """Quantize ``k``/``v`` ``[tokens, heads, D]`` into the pool slots ``indices``. + + ``k_cache``/``v_cache`` are ``[slots, heads, D]`` in the spec's storage dtype and + ``k_scale``/``v_scale`` ``[slots, heads, D // BLOCK]`` in fp16. + """ + from freetoken.kvcache.quant import BLOCK + + num_tokens, num_heads, head_dim = k.shape + if num_tokens == 0: + return + assert head_dim % BLOCK == 0, f"head_dim {head_dim} not a multiple of {BLOCK}" + assert k_cache.shape[1:] == (num_heads, head_dim), ( + f"cache head geometry {tuple(k_cache.shape[1:])} != source {(num_heads, head_dim)}" + ) + _store_kv_quant_kernel[(num_tokens, num_heads)]( + k, + v, + k_cache, + v_cache, + k_scale, + v_scale, + indices, + k.stride(0), + k.stride(1), + k_cache.stride(0), + k_cache.stride(1), + k_scale.stride(0), + k_scale.stride(1), + MAX_MAG=spec.max_magnitude, + IS_INT=spec.is_integer, + BLOCK=BLOCK, + NBLOCK=head_dim // BLOCK, + num_warps=4, + ) + + +__all__ = ["store_kv_quant"] diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 1bb352b4..34fcc322 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -106,6 +106,7 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt num_swa_tokens=num_swa_tokens, device=device, dtype=dtype, + quant=getattr(config, "kv_quant", None), ) @@ -116,7 +117,11 @@ def create_kvcache_pool( dtype: torch.dtype, device: torch.device, num_swa_tokens: int | None = None, + quant=None, ) -> BaseKVCachePool: + from .quant import NONE + + quant = quant if quant is not None else NONE if model_config.has_swa_attention: from .hybrid_swa_pool import HybridSWAKVCache @@ -128,6 +133,7 @@ def create_kvcache_pool( num_swa_tokens=num_swa_tokens, device=device, dtype=dtype, + quant=quant, ) from .mha_pool import MHAKVCache @@ -207,6 +213,7 @@ def create_kvcache_pool( device=device, dtype=dtype, layer_ids=layer_ids, + quant=quant, ) diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index ae8cf9ec..bff022f1 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -21,15 +21,29 @@ def spec_kv_bytes_per_token(spec, config) -> int: x layers, plus the bf16 DSA index-key slab when the spec carries indexer dims. Pure per-spec arithmetic -- pool families compose it over THEIR OWN groups; no family branching here. (2 bytes/elem == the torch.bfloat16 dsa_pool.DSAKVCache._alloc - hardcodes; keep the two in lockstep if the slab dtype ever changes.)""" + hardcodes; keep the two in lockstep if the slab dtype ever changes.) + + An 8-bit KV pool prices at its scheme's bytes per element (1 + 2/32 with the fp16 + per-block scale amortized), not the compute dtype's -- the whole point of the flag is + that this number, times every token of every layer, is what frees VRAM for experts. + The index slab stays bf16: it is never quantized. + """ + bytes_per_elem = _kv_bytes_per_element(config) per_token = ( (1 if spec.mla else 2) # MLA latent groups store one slab (V aliases K) * spec.head_dim * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) - * config.dtype.itemsize * spec.num_layers ) - return per_token + spec.index_head_dim * spec.num_index_layers * 2 + return int(per_token * bytes_per_elem) + spec.index_head_dim * spec.num_index_layers * 2 + + +def _kv_bytes_per_element(config) -> float: + """Storage bytes per K/V element for this engine config's KV pools.""" + quant = getattr(config, "kv_quant", None) + if quant is None or not quant.enabled: + return float(config.dtype.itemsize) + return quant.bytes_per_element(config.dtype) class BaseKVCachePool(ABC): diff --git a/python/freetoken/kvcache/hybrid_swa_pool.py b/python/freetoken/kvcache/hybrid_swa_pool.py index 41c3880e..459ba4ca 100644 --- a/python/freetoken/kvcache/hybrid_swa_pool.py +++ b/python/freetoken/kvcache/hybrid_swa_pool.py @@ -9,6 +9,8 @@ from freetoken.utils import align_ceil, div_even from .base import BaseKVCachePool +from .quant import NONE, KVQuantSpec +from .quant_storage import QuantizedKVStorageMixin @dataclass(frozen=True) @@ -23,9 +25,13 @@ class _KVGroupStorage: k_buffer: torch.Tensor v_buffer: torch.Tensor storage_shape: tuple[int, int, int] + # Per-block scales for an 8-bit group; None when the group stores the compute dtype. + scale_buffer: torch.Tensor | None = None + k_scale: torch.Tensor | None = None + v_scale: torch.Tensor | None = None -class HybridSWAKVCache(BaseKVCachePool): +class HybridSWAKVCache(QuantizedKVStorageMixin, BaseKVCachePool): """SGLang-style wrapper for hybrid full/SWA attention KV storage.""" def __init__( @@ -37,7 +43,9 @@ def __init__( dtype: torch.dtype, device: torch.device, num_swa_tokens: int | None = None, + quant: KVQuantSpec = NONE, ) -> None: + self._quant = quant specs = {group.name: group for group in groups if group.num_layers > 0} if set(specs) != {"full", "swa"}: raise ValueError(f"HybridSWAKVCache requires full and swa groups, got {sorted(specs)}") @@ -80,8 +88,8 @@ def __init__( if self._swa_paged: self._init_swa_paged_state() - @staticmethod def _allocate_group( + self, spec: KVCacheGroupSpec, tp_size: int, outer_size: int, @@ -90,16 +98,17 @@ def _allocate_group( device: torch.device, ) -> _KVGroupStorage: local_kv_heads = div_even(spec.num_kv_heads, tp_size, allow_replicate=True) - buffer = torch.empty( - (2, spec.num_layers, outer_size, inner_size, local_kv_heads, spec.head_dim), - device=device, - dtype=dtype, - ) + shape = (2, spec.num_layers, outer_size, inner_size, local_kv_heads, spec.head_dim) + buffer = torch.empty(shape, device=device, dtype=self._buffer_dtype(dtype)) + scales = self._alloc_scales(shape, device) return _KVGroupStorage( buffer=buffer, k_buffer=buffer[0], v_buffer=buffer[1], storage_shape=(outer_size * inner_size, local_kv_heads, spec.head_dim), + scale_buffer=scales, + k_scale=None if scales is None else scales[0], + v_scale=None if scales is None else scales[1], ) @staticmethod @@ -200,6 +209,16 @@ def v_cache(self, index: int) -> torch.Tensor: ref = self.layers_mapping[index] return self._storages[ref.group].v_buffer[ref.index] + def k_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scales = self._storages[ref.group].k_scale + return None if scales is None else scales[ref.index] + + def v_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scales = self._storages[ref.group].v_scale + return None if scales is None else scales[ref.index] + def store_kv( self, k: torch.Tensor, @@ -207,19 +226,20 @@ def store_kv( out_loc: torch.Tensor, layer_id: int, ) -> None: - from freetoken.kernel import store_cache - ref = self.layers_mapping[layer_id] storage = self._storages[ref.group] indices = out_loc if ref.group == "swa": indices = self.translate_loc_from_full_to_swa(out_loc) - store_cache( - k_cache=storage.k_buffer[ref.index].view(storage.storage_shape), - v_cache=storage.v_buffer[ref.index].view(storage.storage_shape), - indices=indices, - k=k, - v=v, + scale_shape = (storage.storage_shape[0], storage.storage_shape[1], -1) + self._store_kv_into( + storage.k_buffer[ref.index].view(storage.storage_shape), + storage.v_buffer[ref.index].view(storage.storage_shape), + None if storage.k_scale is None else storage.k_scale[ref.index].view(scale_shape), + None if storage.v_scale is None else storage.v_scale[ref.index].view(scale_shape), + indices, + k, + v, ) @property @@ -249,20 +269,20 @@ def _group_geometry(group: _KVGroupStorage) -> tuple: _, num_layers, _old_outer, _old_inner, local_kv_heads, head_dim = group.buffer.shape return (num_layers, local_kv_heads, head_dim, group.buffer.device, group.buffer.dtype) - @staticmethod - def _alloc_group(geom: tuple, outer_size: int, inner_size: int) -> _KVGroupStorage: + def _alloc_group(self, geom: tuple, outer_size: int, inner_size: int) -> _KVGroupStorage: # Only the outer (page/token) dimension changes; the rest comes from ``geom``. num_layers, local_kv_heads, head_dim, device, dtype = geom - buffer = torch.empty( - (2, num_layers, outer_size, inner_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, - ) + shape = (2, num_layers, outer_size, inner_size, local_kv_heads, head_dim) + buffer = torch.empty(shape, device=device, dtype=dtype) + scales = self._alloc_scales(shape, device) return _KVGroupStorage( buffer=buffer, k_buffer=buffer[0], v_buffer=buffer[1], storage_shape=(outer_size * inner_size, local_kv_heads, head_dim), + scale_buffer=scales, + k_scale=None if scales is None else scales[0], + v_scale=None if scales is None else scales[1], ) def rebuild(self, num_full_pages: int, num_swa_tokens: int | None = None) -> None: @@ -343,12 +363,17 @@ def rebuild_from_config( self.rebuild(num_full_pages=num_pages + 1, num_swa_tokens=num_swa_tokens) def unit_bytes(self) -> tuple[int, int]: + def group_bytes(group: _KVGroupStorage) -> int: + total = group.buffer.numel() * group.buffer.element_size() + if group.scale_buffer is not None: + total += group.scale_buffer.numel() * group.scale_buffer.element_size() + return int(total) + full = self.full_kv_pool.buffer - swa = self.swa_kv_pool.buffer full_tokens = int(full.shape[2]) * int(full.shape[3]) return ( - int(full.numel() * full.element_size()) // full_tokens, - int(swa.numel() * swa.element_size()) // self._swa_num_tokens, + group_bytes(self.full_kv_pool) // full_tokens, + group_bytes(self.swa_kv_pool) // self._swa_num_tokens, ) diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b9..654c3281 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/python/freetoken/kvcache/mha_pool.py @@ -7,9 +7,11 @@ from freetoken.utils import div_even from .base import BaseKVCachePool +from .quant import NONE, KVQuantSpec +from .quant_storage import QuantizedKVStorageMixin -class MHAKVCache(BaseKVCachePool): +class MHAKVCache(QuantizedKVStorageMixin, BaseKVCachePool): """ Base class for key-value caches. This class defines the interface for key-value caches used in LLMs. @@ -32,7 +34,9 @@ def __init__( dtype: torch.dtype, device: torch.device, layer_ids: Sequence[int] | None = None, + quant: KVQuantSpec = NONE, ) -> None: + self._quant = quant tp_info = get_tp_info() local_kv_heads = div_even(num_kv_heads, tp_info.size, allow_replicate=True) self._num_layers = num_layers @@ -47,13 +51,16 @@ def __init__( raise ValueError(f"KV layer id {global_id} outside [0, {num_layers})") layer_map[global_id] = dense self._layer_map = layer_map + self._compute_dtype = dtype + kv_shape = (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim) self._kv_buffer = torch.empty( - (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, + kv_shape, device=device, dtype=self._buffer_dtype(dtype) ) self._k_buffer = self._kv_buffer[0] self._v_buffer = self._kv_buffer[1] + self._scale_buffer = self._alloc_scales(kv_shape, device) + self._k_scale = self._scale_buffer[0] if self._scale_buffer is not None else None + self._v_scale = self._scale_buffer[1] if self._scale_buffer is not None else None self._device = device self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) @@ -70,16 +77,21 @@ def rebuild(self, num_pages: int) -> None: self._k_buffer = None self._v_buffer = None self._kv_buffer = None + # Drop the scale slab too before reallocating, for the same reason the KV slab is + # dropped: holding the old one alive can OOM a rebuild the target size would fit. + self._k_scale = None + self._v_scale = None + self._scale_buffer = None if device.type == "cuda": torch.cuda.synchronize(device) torch.cuda.empty_cache() - self._kv_buffer = torch.empty( - (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, - ) + kv_shape = (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim) + self._kv_buffer = torch.empty(kv_shape, device=device, dtype=dtype) self._k_buffer = self._kv_buffer[0] self._v_buffer = self._kv_buffer[1] + self._scale_buffer = self._alloc_scales(kv_shape, device) + self._k_scale = self._scale_buffer[0] if self._scale_buffer is not None else None + self._v_scale = self._scale_buffer[1] if self._scale_buffer is not None else None self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) @classmethod @@ -101,7 +113,10 @@ def rebuild_from_config( def unit_bytes(self) -> tuple[int, int]: buf = self._kv_buffer tokens = int(buf.shape[2]) * int(buf.shape[3]) - return int(buf.numel() * buf.element_size()) // tokens, 0 + total = buf.numel() * buf.element_size() + if self._scale_buffer is not None: + total += self._scale_buffer.numel() * self._scale_buffer.element_size() + return int(total) // tokens, 0 def _dense(self, layer_id: int) -> int: if self._layer_map is None: @@ -117,6 +132,12 @@ def k_cache(self, index: int) -> torch.Tensor: def v_cache(self, index: int) -> torch.Tensor: return self._v_buffer[self._dense(index)] + def k_scale(self, index: int) -> torch.Tensor | None: + return None if self._k_scale is None else self._k_scale[self._dense(index)] + + def v_scale(self, index: int) -> torch.Tensor | None: + return None if self._v_scale is None else self._v_scale[self._dense(index)] + def store_kv( self, k: torch.Tensor, @@ -124,15 +145,16 @@ def store_kv( out_loc: torch.Tensor, layer_id: int, ) -> None: - from freetoken.kernel import store_cache - dense = self._dense(layer_id) - store_cache( - k_cache=self._k_buffer[dense].view(self._storage_shape), - v_cache=self._v_buffer[dense].view(self._storage_shape), - indices=out_loc, - k=k, - v=v, + scale_shape = (self._storage_shape[0], self._storage_shape[1], -1) + self._store_kv_into( + self._k_buffer[dense].view(self._storage_shape), + self._v_buffer[dense].view(self._storage_shape), + None if self._k_scale is None else self._k_scale[dense].view(scale_shape), + None if self._v_scale is None else self._v_scale[dense].view(scale_shape), + out_loc, + k, + v, ) @property @@ -143,6 +165,10 @@ def device(self) -> torch.device: def dtype(self) -> torch.dtype: return self._kv_buffer.dtype + @property + def compute_dtype(self) -> torch.dtype: + return self._compute_dtype + @property def num_layers(self) -> int: return self._num_layers diff --git a/python/freetoken/kvcache/quant.py b/python/freetoken/kvcache/quant.py new file mode 100644 index 00000000..fd3bb439 --- /dev/null +++ b/python/freetoken/kvcache/quant.py @@ -0,0 +1,133 @@ +"""KV-cache quantization schemes: 8-bit storage with a per-block scale. + +The KV pool normally stores K/V in the model's compute dtype (bf16). A quantized pool +stores them in an 8-bit dtype plus a parallel scale tensor holding one fp16 scale per +:data:`BLOCK` elements along ``head_dim`` -- the same geometry GGUF's Q8_0 uses, and the +reason the block is small: KV outliers (mostly in the keys) concentrate in a few +channels, and a block of 32 keeps an outlier from stretching the scale of the whole +head. + +Both schemes share this layout, the store kernel and the dequant path in the attention +kernels; they differ only in the storage dtype and the divisor that maps a block's +max-abs onto the dtype's range. That is deliberate -- picking between them is a flag, +not a second port. + +The scale varies along ``head_dim``, which is the reduction dimension of ``q @ k``, so +the attention kernels cannot fold it in after the dot: they dequantize to bf16 before +the dot. Storage bandwidth is what this buys, not tensor-core throughput. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +# Elements per scale, along head_dim. Matches GGUF Q8_0's block. +BLOCK = 32 +# One fp16 scale per block. +SCALE_DTYPE = torch.float16 + + +@dataclass(frozen=True) +class KVQuantSpec: + """How a KV pool stores its K/V elements. + + ``name`` is the ``--kv-cache-dtype`` value. ``storage_dtype`` is None for the + unquantized pool, in which case the pool allocates in the compute dtype and no + scale tensor exists. + """ + + name: str + storage_dtype: torch.dtype | None + # Max-abs of a block maps to this magnitude in the storage dtype. + max_magnitude: float + + @property + def enabled(self) -> bool: + return self.storage_dtype is not None + + @property + def is_integer(self) -> bool: + """Integer schemes round; float ones just divide.""" + return self.storage_dtype == torch.int8 + + def bytes_per_element(self, compute_dtype: torch.dtype) -> float: + """Storage bytes per K/V element, scales amortized over the block. + + Unquantized: the compute dtype's itemsize. Quantized: 1 byte + 2/32 for the + fp16 scale = 1.0625 -- 6% over a bare 8 bits, versus 16 bits stored. + """ + if not self.enabled: + return float(compute_dtype.itemsize) + return 1.0 + SCALE_DTYPE.itemsize / BLOCK + + def scale_shape(self, shape: tuple[int, ...]) -> tuple[int, ...]: + """Scale-tensor shape for a KV buffer shape: last dim divided by the block.""" + if shape[-1] % BLOCK: + raise ValueError( + f"head_dim {shape[-1]} is not a multiple of the KV quant block {BLOCK}" + ) + return (*shape[:-1], shape[-1] // BLOCK) + + # ---- reference implementations (correctness oracle for the Triton kernels) ---- + + def quantize(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """``x[..., D]`` (float) -> ``(quantized[..., D], scales[..., D // BLOCK])``.""" + assert self.enabled, "quantize() on an unquantized spec" + blocks = x.float().unflatten(-1, (x.shape[-1] // BLOCK, BLOCK)) + amax = blocks.abs().amax(dim=-1) + # An all-zero block would divide by zero; its quantized values are zero either + # way, so any positive scale works. + scales = torch.where(amax > 0, amax / self.max_magnitude, torch.ones_like(amax)) + # Round the scale to its stored precision BEFORE dividing, so quantize and + # dequantize use the identical value. Dividing by the fp32 scale and storing the + # fp16 one leaves a residual error the round-trip cannot cancel. + scales = scales.to(SCALE_DTYPE) + q = blocks / scales.float().unsqueeze(-1) + if self.is_integer: + # Half away from zero, matching GGUF's Q8_0 and the store kernel. + # ``Tensor.round`` is half-to-even and would disagree on ties. + q = torch.where(q >= 0, (q + 0.5).floor(), (q - 0.5).ceil()) + q = q.clamp_(-self.max_magnitude, self.max_magnitude) + return (q.flatten(-2).to(self.storage_dtype), scales) + + def dequantize(self, q: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Inverse of :meth:`quantize`, in float32.""" + assert self.enabled, "dequantize() on an unquantized spec" + blocks = q.float().unflatten(-1, (q.shape[-1] // BLOCK, BLOCK)) + return (blocks * scales.float().unsqueeze(-1)).flatten(-2) + + +# int8 symmetric: a block's max-abs maps to 127. +Q8_0 = KVQuantSpec(name="q8_0", storage_dtype=torch.int8, max_magnitude=127.0) +# e4m3: 4-bit exponent, 3-bit mantissa, max finite magnitude 448. +FP8_E4M3 = KVQuantSpec(name="fp8_e4m3", storage_dtype=torch.float8_e4m3fn, max_magnitude=448.0) +NONE = KVQuantSpec(name="auto", storage_dtype=None, max_magnitude=0.0) + +_BY_NAME = {spec.name: spec for spec in (NONE, Q8_0, FP8_E4M3)} +KV_CACHE_DTYPES = tuple(_BY_NAME) + + +def resolve_kv_quant(name: str | None) -> KVQuantSpec: + """``--kv-cache-dtype`` value -> spec. ``None``/``"auto"`` means unquantized.""" + if name is None: + return NONE + try: + return _BY_NAME[name] + except KeyError: + raise ValueError( + f"unknown --kv-cache-dtype {name!r}; choose from {', '.join(KV_CACHE_DTYPES)}" + ) from None + + +__all__ = [ + "BLOCK", + "SCALE_DTYPE", + "KVQuantSpec", + "KV_CACHE_DTYPES", + "Q8_0", + "FP8_E4M3", + "NONE", + "resolve_kv_quant", +] diff --git a/python/freetoken/kvcache/quant_storage.py b/python/freetoken/kvcache/quant_storage.py new file mode 100644 index 00000000..032cbadc --- /dev/null +++ b/python/freetoken/kvcache/quant_storage.py @@ -0,0 +1,79 @@ +"""Scale-buffer bookkeeping shared by the quantizable KV pools. + +A quantized pool allocates, alongside each K/V slab, a scale slab with the same shape +but the last dimension divided by :data:`~freetoken.kvcache.quant.BLOCK`. The two must +be allocated, rebuilt and freed together, and ``store_kv`` has to route to the +quantizing kernel instead of the byte-copy one -- that is all this mixin owns. The pools +keep their own geometry and indexing. +""" + +from __future__ import annotations + +import torch + +from .quant import NONE, SCALE_DTYPE, KVQuantSpec + + +class QuantizedKVStorageMixin: + """Allocation + store routing for pools whose K/V slabs may be 8-bit. + + Subclasses set ``self._quant`` before allocating and call :meth:`_alloc_scales` for + each K/V buffer they create. ``_quant`` defaulting to the unquantized spec keeps + pools that never opt in behaving exactly as before. + """ + + _quant: KVQuantSpec = NONE + + @property + def quant(self) -> KVQuantSpec: + return self._quant + + def _buffer_dtype(self, compute_dtype: torch.dtype) -> torch.dtype: + """Element dtype for a K/V slab under the active scheme.""" + return self._quant.storage_dtype if self._quant.enabled else compute_dtype + + def _alloc_scales(self, kv_shape: tuple[int, ...], device: torch.device) -> torch.Tensor | None: + """Scale slab matching a ``[2, layers, ..., heads, head_dim]`` K/V buffer. + + None when unquantized -- callers store that verbatim and the attention path reads + it as "no scales", which is what selects the bf16 kernel branch. + """ + if not self._quant.enabled: + return None + return torch.empty( + self._quant.scale_shape(kv_shape), device=device, dtype=SCALE_DTYPE + ) + + def _store_kv_into( + self, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale: torch.Tensor | None, + v_scale: torch.Tensor | None, + indices: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> None: + """Write one layer's K/V, quantizing on the way in when the pool is 8-bit.""" + if not self._quant.enabled: + from freetoken.kernel import store_cache + + store_cache(k_cache=k_cache, v_cache=v_cache, indices=indices, k=k, v=v) + return + + from freetoken.kernel.triton.kv_quant import store_kv_quant + + heads, head_dim = k_cache.shape[-2], k_cache.shape[-1] + store_kv_quant( + k_cache, + k_scale, + v_cache, + v_scale, + indices, + k.view(-1, heads, head_dim), + v.view(-1, heads, head_dim), + self._quant, + ) + + +__all__ = ["QuantizedKVStorageMixin"] diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b75..92f57b16 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -89,6 +89,7 @@ def parse_args( """ from freetoken.attention import validate_attn_backend from freetoken.kvcache import SUPPORTED_CACHE_MANAGER + from freetoken.kvcache.quant import KV_CACHE_DTYPES from freetoken.moe import SUPPORTED_MOE_BACKENDS def _parse_moe_cache_rate(value: str) -> float: @@ -345,6 +346,20 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Set the page size for system management.", ) + parser.add_argument( + "--kv-cache-dtype", + type=str, + choices=list(KV_CACHE_DTYPES), + default=ServerArgs.kv_cache_dtype, + help=( + "KV-cache element storage. 'auto' keeps the compute dtype (bf16). 'q8_0' and " + "'fp8_e4m3' store 8 bits plus an fp16 scale per 32 elements along head_dim " + "(1.0625 bytes/element vs 2), freeing VRAM for the MoE expert cache. q8_0 is " + "the more accurate of the two at this block size. Needs the triton attention " + "backend and head_dim divisible by 32." + ), + ) + parser.add_argument( "--attention-backend", "--attn", diff --git a/tasks/todo.md b/tasks/todo.md new file mode 100644 index 00000000..853add51 --- /dev/null +++ b/tasks/todo.md @@ -0,0 +1,7 @@ +- [x] Review Laguna Phase 3 scoped modules and tests +- [x] Run Laguna imports and models test suite +- [x] Record review results + +## Review +- `tests/models -q`: 48 passed, 1 skipped. Laguna imports passed. +- Findings returned in review response; no source files modified. diff --git a/tests/engine/test_kv_cache_dtype_gating.py b/tests/engine/test_kv_cache_dtype_gating.py new file mode 100644 index 00000000..3d96959a --- /dev/null +++ b/tests/engine/test_kv_cache_dtype_gating.py @@ -0,0 +1,72 @@ +"""--kv-cache-dtype gating: the combinations the quantized path does not implement +must be refused at config time, not at the first kernel launch.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.engine.engine import _validate_kv_cache_dtype +from freetoken.kvcache.quant import NONE, Q8_0 +from freetoken.models.config import KVCacheGroupSpec + + +class _Cfg: + def __init__(self, quant=Q8_0, backend="triton"): + self.kv_quant = quant + self.attention_backend = backend + + +class _Model: + def __init__(self, *specs): + self._specs = specs + + def kv_cache_group_specs(self): + return self._specs + + +def _spec(name="full", head_dim=256, mla=False, index_head_dim=0): + return KVCacheGroupSpec( + name=name, + layer_ids=(0, 1), + num_kv_heads=2, + head_dim=head_dim, + sliding_window=None, + mla=mla, + index_head_dim=index_head_dim, + ) + + +def test_triton_backend_with_aligned_head_dim_is_accepted(): + _validate_kv_cache_dtype(_Cfg(), _Model(_spec(head_dim=256), _spec("swa", 512))) + + +def test_auto_dtype_skips_every_check(): + # Unquantized configs must pass even on backends/pools the quantized path rejects. + _validate_kv_cache_dtype(_Cfg(quant=NONE, backend="fi"), _Model(_spec(mla=True))) + + +@pytest.mark.parametrize("backend", ["fi", "fa", "trtllm", "fa,fi", "triton,fi"]) +def test_non_triton_backends_are_rejected(backend): + with pytest.raises(ValueError, match="needs the triton attention backend"): + _validate_kv_cache_dtype(_Cfg(backend=backend), _Model(_spec())) + + +def test_mla_pool_is_rejected(): + with pytest.raises(ValueError, match="MLA/DSA"): + _validate_kv_cache_dtype(_Cfg(), _Model(_spec(mla=True))) + + +def test_dsa_index_tier_is_rejected(): + with pytest.raises(ValueError, match="MLA/DSA"): + _validate_kv_cache_dtype(_Cfg(), _Model(_spec(index_head_dim=128))) + + +def test_head_dim_not_a_multiple_of_the_block_is_rejected(): + with pytest.raises(ValueError, match="multiple of 32"): + _validate_kv_cache_dtype(_Cfg(), _Model(_spec(head_dim=80))) + + +def test_the_error_names_the_offending_group(): + with pytest.raises(ValueError, match=r"swa \(head_dim 100\)"): + _validate_kv_cache_dtype(_Cfg(), _Model(_spec(head_dim=256), _spec("swa", 100))) diff --git a/tests/kernels/test_kv_quant.py b/tests/kernels/test_kv_quant.py new file mode 100644 index 00000000..9b3f56a4 --- /dev/null +++ b/tests/kernels/test_kv_quant.py @@ -0,0 +1,286 @@ +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.quant import BLOCK, FP8_E4M3, NONE, Q8_0, resolve_kv_quant + +SPECS = [Q8_0, FP8_E4M3] +IDS = [spec.name for spec in SPECS] + +cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _kv(tokens: int, heads: int, dim: int, device="cuda", seed=0) -> torch.Tensor: + g = torch.Generator(device=device).manual_seed(seed) + return torch.randn(tokens, heads, dim, generator=g, device=device, dtype=torch.bfloat16) + + +def test_bytes_per_element_amortizes_the_scale(): + # 8 bits of payload + one fp16 scale per 32 elements. + assert Q8_0.bytes_per_element(torch.bfloat16) == 1.0 + 2 / 32 + assert FP8_E4M3.bytes_per_element(torch.bfloat16) == 1.0 + 2 / 32 + # Unquantized pools price at the compute dtype. + assert NONE.bytes_per_element(torch.bfloat16) == 2.0 + assert NONE.bytes_per_element(torch.float32) == 4.0 + + +def test_resolve_and_scale_shape(): + assert resolve_kv_quant(None) is NONE + assert resolve_kv_quant("auto") is NONE + assert resolve_kv_quant("q8_0") is Q8_0 + with pytest.raises(ValueError, match="unknown --kv-cache-dtype"): + resolve_kv_quant("int4") + assert Q8_0.scale_shape((7, 4, 256)) == (7, 4, 8) + with pytest.raises(ValueError, match="not a multiple"): + Q8_0.scale_shape((7, 4, 100)) + + +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +def test_reference_roundtrip_error_is_within_the_scheme_envelope(spec): + """Each scheme's round-trip error must sit inside the bound its format implies. + + int8 rounds onto a uniform grid of step ``amax/127``, so the error is at most half a + step -- except for a value at the block's extreme, which the clamp can push to a + full step once the fp16 scale rounds down. One step is the honest envelope. e4m3 + carries 3 mantissa bits, so its error is relative, up to ~2^-4 of each value, which + against the block's amax is bounded the same way. + """ + x = torch.randn(64, 4, 256, dtype=torch.float32) + q, scales = spec.quantize(x) + back = spec.dequantize(q, scales) + + blocks = x.unflatten(-1, (x.shape[-1] // BLOCK, BLOCK)) + amax = blocks.abs().amax(dim=-1, keepdim=True) + err = (back - x).unflatten(-1, (x.shape[-1] // BLOCK, BLOCK)).abs() + bound = 1.0 / spec.max_magnitude if spec.is_integer else 1.0 / 2**4 + assert (err <= amax * bound + 1e-6).all() + # And the typical error should sit well under the worst case, not at it. + assert err.mean() <= amax.mean() * bound * 0.5 + + +def test_int8_beats_fp8_on_a_flat_block_and_loses_on_a_spiky_one(): + """The tradeoff the scheme choice turns on, pinned as a test. + + Within a block, int8 spends its codes uniformly and fp8 spends them + logarithmically. So a block of similar magnitudes favours int8, and a block where + one outlier dwarfs the rest favours fp8 -- which is exactly why the block is 32 + elements and not a whole head. + """ + flat = torch.full((1, 1, BLOCK), 1.0) + flat[..., ::2] = 0.9 + spiky = torch.full((1, 1, BLOCK), 0.01) + spiky[..., 0] = 100.0 + + def rel_err(spec, x): + back = spec.dequantize(*spec.quantize(x)) + return ((back - x).abs() / x.abs()).mean().item() + + assert rel_err(Q8_0, flat) < rel_err(FP8_E4M3, flat) + assert rel_err(FP8_E4M3, spiky) < rel_err(Q8_0, spiky) + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +@pytest.mark.parametrize("head_dim", [256, 512]) +def test_store_kernel_matches_the_reference_quantizer(spec, head_dim): + from freetoken.kernel.triton.kv_quant import store_kv_quant + + tokens, heads, slots = 37, 3, 64 + k = _kv(tokens, heads, head_dim, seed=1) + v = _kv(tokens, heads, head_dim, seed=2) + # Scatter to non-contiguous slots: the kernel must honour the index indirection. + indices = torch.randperm(slots, device="cuda")[:tokens].to(torch.int32) + + kc = torch.zeros(slots, heads, head_dim, device="cuda", dtype=spec.storage_dtype) + vc = torch.zeros_like(kc) + ks = torch.zeros(slots, heads, head_dim // BLOCK, device="cuda", dtype=torch.float16) + vs = torch.zeros_like(ks) + + store_kv_quant(kc, ks, vc, vs, indices, k, v, spec) + + idx = indices.to(torch.long) + for src, cache, scale in ((k, kc, ks), (v, vc, vs)): + want_q, want_s = spec.quantize(src.float()) + torch.testing.assert_close(scale[idx].float(), want_s.float(), rtol=0, atol=0) + # Compare dequantized values: int8 is exact, fp8 codes compare through float. + got = spec.dequantize(cache[idx].float(), scale[idx]) + torch.testing.assert_close(got, spec.dequantize(want_q.float(), want_s), rtol=0, atol=0) + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +def test_store_kernel_leaves_untouched_slots_alone(spec): + from freetoken.kernel.triton.kv_quant import store_kv_quant + + heads, head_dim, slots = 2, 256, 16 + k = _kv(4, heads, head_dim, seed=3) + v = _kv(4, heads, head_dim, seed=4) + indices = torch.tensor([1, 3, 5, 7], device="cuda", dtype=torch.int32) + + kc = torch.zeros(slots, heads, head_dim, device="cuda", dtype=spec.storage_dtype) + vc = torch.zeros_like(kc) + ks = torch.zeros(slots, heads, head_dim // BLOCK, device="cuda", dtype=torch.float16) + vs = torch.zeros_like(ks) + store_kv_quant(kc, ks, vc, vs, indices, k, v, spec) + + untouched = [s for s in range(slots) if s not in {1, 3, 5, 7}] + assert (kc[untouched].float() == 0).all() + assert (ks[untouched] == 0).all() + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +def test_store_kernel_handles_an_all_zero_head(spec): + """A zero block has no max to scale by; it must store zeros, not NaNs.""" + from freetoken.kernel.triton.kv_quant import store_kv_quant + + heads, head_dim = 2, 256 + k = torch.zeros(1, heads, head_dim, device="cuda", dtype=torch.bfloat16) + v = torch.zeros_like(k) + indices = torch.zeros(1, device="cuda", dtype=torch.int32) + + kc = torch.empty(4, heads, head_dim, device="cuda", dtype=spec.storage_dtype) + vc = torch.empty_like(kc) + ks = torch.empty(4, heads, head_dim // BLOCK, device="cuda", dtype=torch.float16) + vs = torch.empty_like(ks) + store_kv_quant(kc, ks, vc, vs, indices, k, v, spec) + + assert (kc[0].float() == 0).all() + assert torch.isfinite(ks[0]).all() and (ks[0] > 0).all() + + +# -------------------------------------------------------------------------------------- +# Attention over a quantized pool. +# +# The gate is equivalence against the bf16 kernel fed the SAME dequantized values: that +# isolates "did the dequant path compute attention correctly" from "how much does 8-bit +# storage cost", which is a separate, looser assertion below. +# -------------------------------------------------------------------------------------- + + +def _quantized_pool(spec, k_bf16, v_bf16): + """Store bf16 K/V into a quantized pool; return (kq, ks, vq, vs, k_deq, v_deq).""" + from freetoken.kernel.triton.kv_quant import store_kv_quant + + slots, heads, dim = k_bf16.shape + kq = torch.zeros(slots, heads, dim, device="cuda", dtype=spec.storage_dtype) + vq = torch.zeros_like(kq) + ks = torch.zeros(slots, heads, dim // BLOCK, device="cuda", dtype=torch.float16) + vs = torch.zeros_like(ks) + indices = torch.arange(slots, device="cuda", dtype=torch.int32) + store_kv_quant(kq, ks, vq, vs, indices, k_bf16, v_bf16, spec) + # What the attention kernel will effectively see, in the dtype it dequantizes into. + k_deq = spec.dequantize(kq.float(), ks).to(torch.bfloat16) + v_deq = spec.dequantize(vq.float(), vs).to(torch.bfloat16) + return kq, ks, vq, vs, k_deq, v_deq + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +@pytest.mark.parametrize("head_dim", [256, 512]) +def test_paged_attention_over_quantized_pool(spec, head_dim): + from freetoken.kernel.triton.attention import paged_attention + + slots, q_heads, kv_heads = 96, 8, 2 + q = _kv(6, q_heads, head_dim, seed=5) + k = _kv(slots, kv_heads, head_dim, seed=6) + v = _kv(slots, kv_heads, head_dim, seed=7) + kq, ks, vq, vs, k_deq, v_deq = _quantized_pool(spec, k, v) + + indptr = torch.tensor([0, 40, 96], device="cuda", dtype=torch.int32) + indices = torch.arange(slots, device="cuda", dtype=torch.int32) + q_to_req = torch.tensor([0, 0, 0, 1, 1, 1], device="cuda", dtype=torch.int32) + q_pos = torch.tensor([10, 25, 39, 5, 30, 55], device="cuda", dtype=torch.int32) + kw = dict(indptr=indptr, indices=indices, q_to_req=q_to_req, q_positions=q_pos, + sm_scale=head_dim**-0.5) + + got = paged_attention(q=q, k_cache=kq, v_cache=vq, k_scale=ks, v_scale=vs, **kw) + want = paged_attention(q=q, k_cache=k_deq, v_cache=v_deq, **kw) + torch.testing.assert_close(got, want, rtol=2e-2, atol=2e-2) + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +@pytest.mark.parametrize("head_dim", [256, 512]) +def test_decode_attention_over_quantized_pool(spec, head_dim): + from freetoken.kernel.triton.attention import decode_paged_attention + + slots, q_heads, kv_heads, batch = 128, 8, 2, 3 + q = _kv(batch, q_heads, head_dim, seed=8) + k = _kv(slots, kv_heads, head_dim, seed=9) + v = _kv(slots, kv_heads, head_dim, seed=10) + kq, ks, vq, vs, k_deq, v_deq = _quantized_pool(spec, k, v) + + indptr = torch.tensor([0, 40, 90, 128], device="cuda", dtype=torch.int32) + indices = torch.arange(slots, device="cuda", dtype=torch.int32) + q_pos = torch.tensor([39, 49, 37], device="cuda", dtype=torch.int32) + splits = 4 + logits = torch.zeros(batch, q_heads, splits, head_dim, device="cuda", dtype=torch.float32) + lse = torch.zeros(batch, q_heads, splits, device="cuda", dtype=torch.float32) + nsplits = torch.full((batch,), splits, device="cuda", dtype=torch.int32) + kw = dict(indptr=indptr, indices=indices, q_positions=q_pos, attn_logits=logits, + attn_lse=lse, num_kv_splits=nsplits, max_kv_splits=splits, + sm_scale=head_dim**-0.5) + + got = decode_paged_attention(q=q, k_cache=kq, v_cache=vq, k_scale=ks, v_scale=vs, **kw) + want = decode_paged_attention(q=q, k_cache=k_deq, v_cache=v_deq, **kw) + torch.testing.assert_close(got, want, rtol=2e-2, atol=2e-2) + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +@pytest.mark.parametrize("split", [False, True], ids=["fused", "split"]) +def test_extend_attention_over_quantized_pool(spec, split): + """Both extend paths. The split kernel reads the freshly-computed K/V in bf16 and + only the prefix from the quantized pool, so it exercises the mixed case.""" + from freetoken.kernel.triton.attention import extend_paged_attention + + head_dim, slots, q_heads, kv_heads = 256, 64, 8, 2 + q_len, prefix = 8, 24 + q = _kv(q_len, q_heads, head_dim, seed=11) + k = _kv(slots, kv_heads, head_dim, seed=12) + v = _kv(slots, kv_heads, head_dim, seed=13) + kq, ks, vq, vs, k_deq, v_deq = _quantized_pool(spec, k, v) + + qo_indptr = torch.tensor([0, q_len], device="cuda", dtype=torch.int32) + kv_indptr = torch.tensor([0, prefix + q_len], device="cuda", dtype=torch.int32) + kv_indices = torch.arange(prefix + q_len, device="cuda", dtype=torch.int32) + prefix_lens = torch.tensor([prefix], device="cuda", dtype=torch.int32) + extend = {} + if split: + extend = dict(k_extend=_kv(q_len, kv_heads, head_dim, seed=14), + v_extend=_kv(q_len, kv_heads, head_dim, seed=15)) + kw = dict(qo_indptr=qo_indptr, kv_indptr=kv_indptr, kv_indices=kv_indices, + prefix_lens=prefix_lens, max_q_len=q_len, sm_scale=head_dim**-0.5, **extend) + + got = extend_paged_attention(q=q, k_cache=kq, v_cache=vq, k_scale=ks, v_scale=vs, **kw) + want = extend_paged_attention(q=q, k_cache=k_deq, v_cache=v_deq, **kw) + torch.testing.assert_close(got, want, rtol=2e-2, atol=2e-2) + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +def test_quantized_attention_tracks_the_bf16_pool(spec): + """The end-to-end cost of 8-bit storage: attention over a quantized pool against + attention over the original bf16 one. This is the number that matters for quality, + and it is looser than the kernel-equivalence gate above by construction.""" + from freetoken.kernel.triton.attention import paged_attention + + head_dim, slots, q_heads, kv_heads = 256, 128, 8, 2 + q = _kv(4, q_heads, head_dim, seed=16) + k = _kv(slots, kv_heads, head_dim, seed=17) + v = _kv(slots, kv_heads, head_dim, seed=18) + kq, ks, vq, vs, _, _ = _quantized_pool(spec, k, v) + + indptr = torch.tensor([0, slots], device="cuda", dtype=torch.int32) + indices = torch.arange(slots, device="cuda", dtype=torch.int32) + q_to_req = torch.zeros(4, device="cuda", dtype=torch.int32) + q_pos = torch.tensor([20, 60, 100, 127], device="cuda", dtype=torch.int32) + kw = dict(indptr=indptr, indices=indices, q_to_req=q_to_req, q_positions=q_pos, + sm_scale=head_dim**-0.5) + + got = paged_attention(q=q, k_cache=kq, v_cache=vq, k_scale=ks, v_scale=vs, **kw) + ref = paged_attention(q=q, k_cache=k, v_cache=v, **kw) + rel = ((got.float() - ref.float()).norm() / ref.float().norm()).item() + assert rel < 0.05, f"{spec.name}: relative error {rel:.4f} vs bf16 pool" diff --git a/tests/kvcache/test_kv_quant_pool.py b/tests/kvcache/test_kv_quant_pool.py new file mode 100644 index 00000000..9537d6d2 --- /dev/null +++ b/tests/kvcache/test_kv_quant_pool.py @@ -0,0 +1,229 @@ +"""KV pools backed by 8-bit storage: allocation, store/read-back, cost, rebuild.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kvcache.quant import BLOCK, FP8_E4M3, NONE, Q8_0 + +from .test_hybrid_swa_kv_cache import _kv_group_specs, _patch_tp + +SPECS = [Q8_0, FP8_E4M3] +IDS = [spec.name for spec in SPECS] + +# Measured relative L2 error of a round trip through each scheme on gaussian KV with a +# 32-element block: q8_0 ~0.005, fp8_e4m3 ~0.024. int8 wins by ~4.5x because a block that +# small keeps outliers from stretching the scale, which is the reason q8_0 is the default. +MAX_REL_ERR = {Q8_0.name: 0.01, FP8_E4M3.name: 0.03} + +cuda_only = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + + +def _swa_pool(quant, device="cuda", num_full_pages=64, num_swa_tokens=32): + from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache + + return HybridSWAKVCache( + groups=_kv_group_specs(), + num_layers=6, + num_full_pages=num_full_pages, + page_size=1, + num_swa_tokens=num_swa_tokens, + device=torch.device(device), + dtype=torch.bfloat16, + quant=quant, + ) + + +def _mha_pool(quant, device="cuda", num_pages=64): + from freetoken.kvcache.mha_pool import MHAKVCache + + return MHAKVCache( + num_kv_heads=2, + num_layers=4, + head_dim=256, + num_pages=num_pages, + page_size=1, + dtype=torch.bfloat16, + device=torch.device(device), + quant=quant, + ) + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +def test_swa_pool_allocates_8bit_slabs_and_matching_scales(monkeypatch, spec): + _patch_tp(monkeypatch) + pool = _swa_pool(spec) + + # Layer 0 is SWA (head_dim 256, 8 kv heads), layer 2 full (head_dim 512, 2 heads). + for layer, head_dim, heads in ((0, 256, 8), (2, 512, 2)): + k = pool.k_cache(layer) + s = pool.k_scale(layer) + assert k.dtype == spec.storage_dtype + assert s.dtype == torch.float16 + assert k.shape[-2:] == (heads, head_dim) + assert s.shape[-2:] == (heads, head_dim // BLOCK) + assert s.shape[:-1] == k.shape[:-1] + + +@cuda_only +def test_unquantized_pool_keeps_bf16_and_has_no_scales(monkeypatch): + _patch_tp(monkeypatch) + pool = _swa_pool(NONE) + assert pool.k_cache(0).dtype == torch.bfloat16 + assert pool.k_scale(0) is None + assert pool.v_scale(0) is None + assert not pool.quant.enabled + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +@pytest.mark.parametrize("layer", [0, 2], ids=["swa", "full"]) +def test_store_kv_round_trips_through_the_quantized_pool(monkeypatch, spec, layer): + _patch_tp(monkeypatch) + pool = _swa_pool(spec) + heads, head_dim = pool.k_cache(layer).shape[-2:] + + tokens = 8 + g = torch.Generator(device="cuda").manual_seed(0) + k = torch.randn(tokens, heads, head_dim, generator=g, device="cuda", dtype=torch.bfloat16) + v = torch.randn(tokens, heads, head_dim, generator=g, device="cuda", dtype=torch.bfloat16) + out_loc = torch.arange(1, tokens + 1, device="cuda", dtype=torch.int32) + if pool.is_swa_layer(layer): + pool.alloc_swa(out_loc) + + pool.store_kv(k, v, out_loc, layer) + + slots = ( + pool.translate_loc_from_full_to_swa(out_loc) if pool.is_swa_layer(layer) else out_loc + ).to(torch.long) + got = spec.dequantize( + pool.k_cache(layer).view(-1, heads, head_dim)[slots].float(), + pool.k_scale(layer).view(-1, heads, head_dim // BLOCK)[slots], + ) + # Storing is lossy by construction; what must hold is that it round-trips to within + # the scheme's envelope, not that it is exact. + rel = ((got - k.float()).norm() / k.float().norm()).item() + assert rel < MAX_REL_ERR[spec.name], ( + f"{spec.name} layer {layer}: relative round-trip error {rel:.4f}" + ) + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +def test_unit_bytes_counts_the_scale_slab(monkeypatch, spec): + """Budgeting has to see 1 + 2/32 bytes per element, not 1 -- otherwise a rebuild + would size pools against memory the scales are quietly consuming.""" + _patch_tp(monkeypatch) + quantized = _swa_pool(spec) + plain = _swa_pool(NONE) + + q_full, q_swa = quantized.unit_bytes() + p_full, p_swa = plain.unit_bytes() + for q, p in ((q_full, p_full), (q_swa, p_swa)): + # bf16 is 2 bytes/element, the quantized pool 1 + 2/32 = 1.0625. + assert q == pytest.approx(p * (1.0625 / 2.0), rel=1e-3) + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +def test_rebuild_reallocates_scales_and_keeps_identity(monkeypatch, spec): + _patch_tp(monkeypatch) + pool = _swa_pool(spec, num_full_pages=64, num_swa_tokens=32) + before = id(pool) + + pool.rebuild(num_full_pages=128, num_swa_tokens=64) + + assert id(pool) == before, "rebuild must preserve object identity" + for layer in (0, 2): + k, s = pool.k_cache(layer), pool.k_scale(layer) + assert k.dtype == spec.storage_dtype + assert s is not None and s.dtype == torch.float16 + assert s.shape[:-1] == k.shape[:-1] + assert s.shape[-1] == k.shape[-1] // BLOCK + assert pool.k_cache(2).shape[0] == 128 + assert pool.k_cache(0).shape[0] == 64 + + +@cuda_only +@pytest.mark.parametrize("spec", SPECS, ids=IDS) +def test_mha_pool_quantizes_and_round_trips(monkeypatch, spec): + from freetoken.distributed.info import DistributedInfo + + monkeypatch.setattr( + "freetoken.kvcache.mha_pool.get_tp_info", lambda: DistributedInfo(rank=0, size=1) + ) + pool = _mha_pool(spec) + assert pool.k_cache(0).dtype == spec.storage_dtype + assert pool.k_scale(0).shape[-1] == 256 // BLOCK + + g = torch.Generator(device="cuda").manual_seed(1) + k = torch.randn(6, 2, 256, generator=g, device="cuda", dtype=torch.bfloat16) + v = torch.randn(6, 2, 256, generator=g, device="cuda", dtype=torch.bfloat16) + out_loc = torch.arange(1, 7, device="cuda", dtype=torch.int32) + pool.store_kv(k, v, out_loc, 0) + + idx = out_loc.to(torch.long) + got = spec.dequantize( + pool.k_cache(0).view(-1, 2, 256)[idx].float(), + pool.k_scale(0).view(-1, 2, 256 // BLOCK)[idx], + ) + rel = ((got - k.float()).norm() / k.float().norm()).item() + assert rel < MAX_REL_ERR[spec.name] + + pool.rebuild(32) + assert pool.k_cache(0).shape[0] == 32 + assert pool.k_scale(0).shape[0] == 32 + + +def test_cost_model_prices_the_quantized_pool_below_bf16(monkeypatch): + """The whole point of the flag, as arithmetic: same geometry, ~half the bytes.""" + from freetoken.kvcache.base import spec_kv_bytes_per_token + from freetoken.distributed.info import DistributedInfo + from types import SimpleNamespace + + spec = _kv_group_specs()[0] # full group: 2 kv heads, head_dim 512, 2 layers + tp = DistributedInfo(rank=0, size=1) + + def cfg(quant): + return SimpleNamespace(tp_info=tp, dtype=torch.bfloat16, kv_quant=quant) + + plain = spec_kv_bytes_per_token(spec, cfg(NONE)) + quantized = spec_kv_bytes_per_token(spec, cfg(Q8_0)) + assert quantized == pytest.approx(plain * (1.0625 / 2.0), rel=1e-3) + # A config with no kv_quant attribute at all must price as bf16 (back-compat with + # every caller that predates the flag). + legacy = spec_kv_bytes_per_token(spec, SimpleNamespace(tp_info=tp, dtype=torch.bfloat16)) + assert legacy == plain + + +@cuda_only +def test_q8_0_stores_kv_more_accurately_than_fp8(monkeypatch): + """Why q8_0 is the default, measured rather than argued. + + Both schemes cost the same bytes and the same kernel work here, so the choice is + purely numerical. On a 32-element block int8's uniform grid beats e4m3's 3-bit + mantissa by several times -- fp8's advantage only shows up with a whole-head scale, + where one outlier would crush everything else, or on kernels with native fp8 support + (which the SWA path cannot use anyway). + """ + _patch_tp(monkeypatch) + g = torch.Generator(device="cuda").manual_seed(7) + k = torch.randn(16, 8, 256, generator=g, device="cuda", dtype=torch.bfloat16) + v = torch.randn(16, 8, 256, generator=g, device="cuda", dtype=torch.bfloat16) + out_loc = torch.arange(1, 17, device="cuda", dtype=torch.int32) + + errs = {} + for spec in SPECS: + pool = _swa_pool(spec, num_swa_tokens=64) + pool.alloc_swa(out_loc) + pool.store_kv(k, v, out_loc, 0) + slots = pool.translate_loc_from_full_to_swa(out_loc).to(torch.long) + got = spec.dequantize( + pool.k_cache(0).view(-1, 8, 256)[slots].float(), + pool.k_scale(0).view(-1, 8, 256 // BLOCK)[slots], + ) + errs[spec.name] = ((got - k.float()).norm() / k.float().norm()).item() + + assert errs[Q8_0.name] < errs[FP8_E4M3.name] / 2, errs From 972a38e4e567c4af2f3521e9c463aa78425854da Mon Sep 17 00:00:00 2001 From: probe Date: Mon, 24 Aug 2026 06:38:54 +0400 Subject: [PATCH 2/2] fix(kvcache): round fp8 KV to the e4m3 grid before the native cast The native fp32 -> float8e4nv downcast does not round to nearest everywhere: on sm_89 triton lowers it as a truncating fp32 -> fp16 -> e4m3 double-round, so values just above a grid midpoint collapse downward and disagree with the RNE torch reference (~0.4% of elements). Round explicitly with round_e4m3 before clamping, mirroring the int branch's rounding. Fixes the two fp8 reference-equality tests on sm_89. Co-Authored-By: Claude --- python/freetoken/kernel/triton/kv_quant.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py index f2748d38..944ff8f7 100644 --- a/python/freetoken/kernel/triton/kv_quant.py +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -18,6 +18,8 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import round_e4m3 + @triton.jit def _store_kv_quant_kernel( @@ -70,6 +72,12 @@ def _store_kv_quant_kernel( # float->int cast truncates. q = tl.where(q >= 0, tl.floor(q + 0.5), tl.ceil(q - 0.5)) q = tl.minimum(tl.maximum(q, -MAX_MAG), MAX_MAG) + else: + # The native fp32 -> float8e4nv downcast does not round to nearest on + # every arch (it lowers as a truncating fp32 -> fp16 -> e4m3 double-round + # on sm_89), so values just above a grid midpoint collapse downward and + # disagree with the RNE torch reference. Round explicitly first. + q = round_e4m3(tl.minimum(tl.maximum(q, -MAX_MAG), MAX_MAG)) tl.store( dst_ptr + slot * stride_ct + head * stride_ch + offs,