From 5a0027c52505466d9c91702406e639965b104a77 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Sun, 23 Aug 2026 17:19:31 -0400 Subject: [PATCH] feat(kvcache): 8-bit DSV4 window/compressed KV behind --kv-cache-dtype DeepSeek-V4's attention KV is one MLA latent row per token, so unlike a separate-K-and-V pool there is a single buffer to quantize and a single scale array beside it. `--kv-cache-dtype fp8_e4m3` stores the window and compressed pools as fp8 plus one fp16 scale per 32 elements along head_dim: 1.0625 bytes per element against bf16's 2. The precision cost is small because the model already rounds this KV onto the e4m3 grid before it reaches the pool -- act_quant_fp8_inplace(kv[..., :-rd], 64) runs immediately before every store_window and before the compressor's scatter. Re-blocking 64 -> 32 moves values within the grid they already live on. The rope tail is the exception: those dims are genuine bf16 and this rounds them. Scope, and what is deliberately left alone: - window_pool and cmp_pool are quantized together. The gather selects a pool BASE pointer per column and issues one tl.load, so the two must share a row layout -- they quantize together or not at all. - idx_pool stays bf16. The indexer's K decides top-k *selection*, so an error there perturbs routing rather than a value, and it is 3.7% of the pool. - The compress-state rings stay fp32. They are accumulators, not storage, and at 53% of the pool they are the larger target -- separately. Both sparse-attention kernels dequantize the gathered tile before the dot (the scale varies along head_dim, the reduction axis, so it cannot be folded in after). QUANT=0 compiles the existing bf16 path unchanged. Measured on an RTX 6000 Ada (sm_89), DeepSeek-V4-Flash, offload MoE, TP=1, --memory-ratio 0.96, with --moe-cache-auto sizing the split: num_pages expert slots single 8-concurrent bf16 560 2551 18.81 34.88 fp8_e4m3 670 2571 18.31 30.18 +19.6% +20 -2.7% -13.5% So it buys 19.6% more KV capacity for the same VRAM and costs throughput: the dequant is on the critical path and there is little to amortise it against at these batch sizes. Whether that trade is worth taking depends on whether a deployment is context-bound or throughput-bound; it is off by default. A code planted ~9000 tokens back was recalled under both dtypes. The cost model has to agree with the pool: --moe-cache-auto divides the VRAM budget by bytes-per-KV-element to pick the expert-cache size, and it resolves before any pool exists, so kv_quant is stamped on dsv4_args at config resolution rather than at pool construction. Getting that ordering wrong is silent -- the pool shrinks and the saving never reaches the budget. Gated at config time to a DeepSeek-V4 checkpoint, head_dim a multiple of 32, and native fp8 (sm_89+, since a float8 pointer is illegal in triton below it). 17 tests: store kernel against a torch reference, quantized attention against the bf16 reference on both the prefill and split-k decode kernels, masked columns, cmp_counts bounds, saturation, and the already-on-grid regime. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun --- python/freetoken/attention/dsv4_compress.py | 13 ++ python/freetoken/attention/dsv4_sparse.py | 9 + python/freetoken/engine/config.py | 2 + python/freetoken/engine/engine.py | 40 ++++ .../freetoken/kernel/triton/dsv4/kv_quant.py | 77 ++++++ .../kernel/triton/dsv4/sparse_attn.py | 67 +++++- python/freetoken/kvcache/__init__.py | 4 + python/freetoken/kvcache/dsv4_cost_model.py | 14 +- python/freetoken/kvcache/dsv4_kv_quant.py | 77 ++++++ python/freetoken/kvcache/dsv4_paged_pool.py | 56 ++++- python/freetoken/server/args.py | 13 ++ tests/kvcache/test_dsv4_kv_quant.py | 219 ++++++++++++++++++ 12 files changed, 577 insertions(+), 14 deletions(-) create mode 100644 python/freetoken/kernel/triton/dsv4/kv_quant.py create mode 100644 python/freetoken/kvcache/dsv4_kv_quant.py create mode 100644 tests/kvcache/test_dsv4_kv_quant.py diff --git a/python/freetoken/attention/dsv4_compress.py b/python/freetoken/attention/dsv4_compress.py index c034b1b9..a04ee475 100644 --- a/python/freetoken/attention/dsv4_compress.py +++ b/python/freetoken/attention/dsv4_compress.py @@ -36,6 +36,12 @@ class CompressorBackendMixin: def compress_pool(self, layer_id: int, tier: str) -> torch.Tensor: return getattr(self.pool, self._TIER_ATTRS[tier][0])[layer_id] + def compress_scale(self, layer_id: int, tier: str): + """Parallel scale array for an 8-bit compressed pool, else None.""" + if tier != "attn": + return None + return getattr(self.pool, "cmp_scale", [None] * (layer_id + 1))[layer_id] + def compress_state_ring(self, layer_id: int, tier: str): return getattr(self.pool, self._TIER_ATTRS[tier][1])[layer_id] @@ -68,6 +74,13 @@ def scatter_compressed( self, layer_id: int, tier: str, rows: torch.Tensor, kv: torch.Tensor ) -> None: pool = self.compress_pool(layer_id, tier) + # The attention tier's cmp_pool may be 8-bit; the indexer tier never is. + scales = self.compress_scale(layer_id, tier) + if scales is not None: + from freetoken.kernel.triton.dsv4.kv_quant import store_kv_quant + + store_kv_quant(pool, scales, rows, kv.reshape(-1, pool.shape[1])) + return pool.index_copy_(0, rows, kv.to(pool.dtype)) # ----- compress-state ring --------------------------------------------------------- diff --git a/python/freetoken/attention/dsv4_sparse.py b/python/freetoken/attention/dsv4_sparse.py index 2c740e2a..985f4342 100644 --- a/python/freetoken/attention/dsv4_sparse.py +++ b/python/freetoken/attention/dsv4_sparse.py @@ -261,9 +261,18 @@ def attend( # ratio-0 layers have no compressed pool; the kernel never reads it there (n_window == # topk), so alias the window pool to keep the two-pool stride assert happy. cmp = pool.cmp_pool[layer_id] if has_compression else pool.window_pool[layer_id] + # 8-bit pools carry a parallel scale array; alias it the same way the pool is aliased + # on a ratio-0 layer so the kernel's stride-sharing assert still holds. + win_s = getattr(pool, "window_scale", None) + if win_s is not None and win_s[layer_id] is not None: + cmp_s = pool.cmp_scale[layer_id] if has_compression else win_s[layer_id] + else: + cmp_s = None return sparse_attn_paged( q, pool.window_pool[layer_id], cmp, attn_sink, topk_idxs.int(), n_window, softmax_scale, cmp_counts=cmp_counts, + window_scale=win_s[layer_id] if win_s is not None else None, + cmp_scale=cmp_s, ) # ----- internals -------------------------------------------------------------------- diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f3..e44c8ceb 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -32,6 +32,8 @@ class EngineConfig: moe_cache_rate: float | None = None moe_cache_auto: bool = False kv_reserve_tokens: int = 8192 # KV floor for --moe-cache-auto; small by design (MoE-priority) + # 8-bit storage for the DSV4 window/compressed KV pools: "auto" (bf16) or "fp8_e4m3". + kv_cache_dtype: str = "auto" moe_cache_policy: str = "lru" moe_prefill_overlap: bool = True # Prefill hit/miss split: serve cache-resident experts D2D during prefill diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ff3c985f..38f2609e 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -154,6 +154,45 @@ def _resolve_auto_attention_backend( ) +def _validate_kv_cache_dtype(config, model_config) -> None: + """Gate --kv-cache-dtype against what the 8-bit path actually implements. + + The storage lives in the DSV4 window/compressed pools and the two sparse-attention + kernels that read them. Every other model's KV goes through different pools and + kernels, so reject those here, at config time, rather than letting a wrong-dtype + tensor reach a kernel. + """ + name = getattr(config, "kv_cache_dtype", "auto") + dsv4_args = getattr(model_config, "dsv4_args", None) + # Stamp the args here, at config resolution, not when the pool is built: the cost model + # reads this to size the KV budget, and --moe-cache-auto divides that budget to pick the + # expert-cache size well before any pool exists. Setting it later leaves the auto-sizer + # working from bf16 bytes and the 8-bit saving never reaches the expert cache. + if dsv4_args is not None: + dsv4_args.kv_quant = name != "auto" + if name == "auto": + return + + if dsv4_args is None: + raise ValueError( + f"--kv-cache-dtype {name} is implemented for DeepSeek-V4 KV pools only; this " + "checkpoint uses a different pool family. Drop the flag." + ) + + from freetoken.kvcache.dsv4_kv_quant import BLOCK + + if dsv4_args.head_dim % BLOCK: + raise ValueError( + f"--kv-cache-dtype {name} needs head_dim to be a multiple of {BLOCK}, the " + f"quantization block; this checkpoint has {dsv4_args.head_dim}." + ) + if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9): + raise ValueError( + f"--kv-cache-dtype {name} needs native fp8 (sm_89 or newer): below that a " + "float8 pointer is illegal inside a triton kernel. Use --kv-cache-dtype auto." + ) + + 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 +1229,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/dsv4/kv_quant.py b/python/freetoken/kernel/triton/dsv4/kv_quant.py new file mode 100644 index 00000000..5d7e7600 --- /dev/null +++ b/python/freetoken/kernel/triton/dsv4/kv_quant.py @@ -0,0 +1,77 @@ +"""Quantizing store for the DSV4 window / compressed KV pools. + +One program per row: read the bf16 row, take a max-abs per :data:`BLOCK`, and write the +fp8 bytes plus the fp16 scales into the pool at the row's slot. Replaces the plain +``index_copy_`` the unquantized pool does. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.e4m3_compat import round_e4m3 +from freetoken.kvcache.dsv4_kv_quant import BLOCK, MAX_MAG, scale_width + + +@triton.jit +def _store_kv_quant_kernel( + src_ptr, dst_ptr, sc_ptr, slot_ptr, + stride_sm, stride_dm, stride_cm, + D: tl.constexpr, NB: tl.constexpr, QBLOCK: tl.constexpr, MAXMAG: tl.constexpr, +): + row = tl.program_id(0) + slot = tl.load(slot_ptr + row).to(tl.int64) + + # Address the row as [NB, QBLOCK] directly rather than loading [D] and reshaping: + # the block axis is the reduction axis for the max-abs, and triton reshapes from 1-D + # are not free to express. + offs2 = tl.arange(0, NB)[:, None] * QBLOCK + tl.arange(0, QBLOCK)[None, :] + xb = tl.load(src_ptr + row * stride_sm + offs2).to(tl.float32) # [NB, QBLOCK] + + amax = tl.max(tl.abs(xb), 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 / MAXMAG, 1.0) + # Round the scale to its stored precision before dividing, so the value written here + # and the value the attention kernel reads back are scaled by the identical number. + scale = scale.to(tl.float16).to(tl.float32) + + # div_rn, not `/`: the plain operator may lower to a reciprocal multiply, which + # disagrees with the torch reference on values sitting exactly between two steps. + q = tl.math.div_rn(xb, scale[:, None]) + q = tl.minimum(tl.maximum(q, -MAXMAG), MAXMAG) + # round_e4m3 before the cast: triton lowers fp32 -> float8e4nv as a double-round + # (fp32 -> fp16 RTZ -> e4m3), which is one-sided toward zero. Putting the value on + # the grid in fp32 first makes the cast exact. + q = round_e4m3(q) + + tl.store(dst_ptr + slot * stride_dm + offs2, q.to(tl.float8e4nv)) + tl.store(sc_ptr + slot * stride_cm + tl.arange(0, NB), scale.to(tl.float16)) + + +def store_kv_quant( + pool: torch.Tensor, # [slots, D] float8_e4m3fn + scales: torch.Tensor, # [slots, D//BLOCK] fp16 + slots: torch.Tensor, # [rows] int + kv: torch.Tensor, # [rows, D] compute dtype +) -> None: + """Quantize ``kv`` into ``pool``/``scales`` at ``slots``.""" + rows, d = kv.shape + assert pool.shape[1] == d and pool.dtype == torch.float8_e4m3fn, (pool.shape, pool.dtype) + nb = scale_width(d) + assert scales.shape[1] == nb, (scales.shape, nb) + if rows == 0: + return + kv = kv.contiguous() + slots = slots.to(torch.int64).contiguous() + _store_kv_quant_kernel[(rows,)]( + kv, pool, scales, slots, + kv.stride(0), pool.stride(0), scales.stride(0), + D=d, NB=nb, QBLOCK=BLOCK, MAXMAG=MAX_MAG, + num_warps=4, + ) + + +__all__ = ["store_kv_quant"] diff --git a/python/freetoken/kernel/triton/dsv4/sparse_attn.py b/python/freetoken/kernel/triton/dsv4/sparse_attn.py index c22f891e..ae847b1e 100644 --- a/python/freetoken/kernel/triton/dsv4/sparse_attn.py +++ b/python/freetoken/kernel/triton/dsv4/sparse_attn.py @@ -50,12 +50,13 @@ @triton.jit def _sparse_attn_paged_kernel( - q_ptr, win_ptr, cmp_ptr, o_ptr, sink_ptr, idx_ptr, cnt_ptr, + q_ptr, win_ptr, cmp_ptr, win_s_ptr, cmp_s_ptr, o_ptr, sink_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, N_WINDOW, stride_qb, stride_qm, stride_qh, stride_qd, stride_wn, stride_wd, stride_cn, stride_cd, + stride_sn, stride_sd, stride_ob, stride_om, stride_oh, stride_od, stride_ib, stride_im, stride_it, stride_nb, stride_nm, @@ -63,6 +64,8 @@ def _sparse_attn_paged_kernel( BLOCK_H: tl.constexpr, BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, + QUANT: tl.constexpr, + QBLOCK: tl.constexpr, ): pid_m = tl.program_id(0) pid_b = tl.program_id(1) @@ -103,6 +106,17 @@ def _sparse_attn_paged_kernel( kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32) # [BLOCK_T, D] + if QUANT: + # 8-bit pool: one fp16 scale per QBLOCK dims along head_dim. The scale varies + # along the reduction axis of q @ kv, so it cannot be folded in after the dot -- + # the tile is dequantized here, before it is used. Indexing the scale per-d + # (rather than loading [BLOCK_T, D//QBLOCK] and broadcasting) keeps this to one + # extra load site: the D/QBLOCK distinct values per row stay hot in L1. + s_base = tl.where(is_win, win_s_ptr, cmp_s_ptr) + s_ptrs = (s_base[:, None] + idxs[:, None] * stride_sn + + (offs_d[None, :] // QBLOCK) * stride_sd) + kv = kv * tl.load(s_ptrs, mask=valid[:, None], other=1.0).to(tl.float32) + scores = tl.dot(q, tl.trans(kv)) * scale # [BLOCK_H, BLOCK_T] scores = tl.where(valid[None, :], scores, -float("inf")) @@ -126,12 +140,13 @@ def _sparse_attn_paged_kernel( @triton.jit def _sparse_attn_paged_splitk_kernel( - q_ptr, win_ptr, cmp_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, + q_ptr, win_ptr, cmp_ptr, win_s_ptr, cmp_s_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, N_WINDOW, stride_qb, stride_qm, stride_qh, stride_qd, stride_wn, stride_wd, stride_cn, stride_cd, + stride_sn, stride_sd, stride_mb, stride_mm, stride_mh, stride_ms, stride_md, stride_lb, stride_lm, stride_lh, stride_ls, stride_ib, stride_im, stride_it, @@ -140,6 +155,8 @@ def _sparse_attn_paged_splitk_kernel( BLOCK_H: tl.constexpr, BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, + QUANT: tl.constexpr, + QBLOCK: tl.constexpr, NUM_SPLITS: tl.constexpr, ): """Stage 1: each program reduces one BLOCK_T-aligned slice of the candidate list and writes @@ -186,6 +203,17 @@ def _sparse_attn_paged_splitk_kernel( kv_ptrs = base[:, None] + idxs[:, None] * stride_wn + offs_d[None, :] * stride_wd kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0).to(tl.float32) + if QUANT: + # 8-bit pool: one fp16 scale per QBLOCK dims along head_dim. The scale varies + # along the reduction axis of q @ kv, so it cannot be folded in after the dot -- + # the tile is dequantized here, before it is used. Indexing the scale per-d + # (rather than loading [BLOCK_T, D//QBLOCK] and broadcasting) keeps this to one + # extra load site: the D/QBLOCK distinct values per row stay hot in L1. + s_base = tl.where(is_win, win_s_ptr, cmp_s_ptr) + s_ptrs = (s_base[:, None] + idxs[:, None] * stride_sn + + (offs_d[None, :] // QBLOCK) * stride_sd) + kv = kv * tl.load(s_ptrs, mask=valid[:, None], other=1.0).to(tl.float32) + scores = tl.dot(q, tl.trans(kv)) * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -289,6 +317,8 @@ def sparse_attn_paged( n_window: int, # # of window entries (the first n_window cols of topk) softmax_scale: float, cmp_counts: torch.Tensor | None = None, # [b, m] int32, live compressed columns per query + window_scale: torch.Tensor | None = None, # [n_win_slots, d//QBLOCK] fp16, 8-bit pools only + cmp_scale: torch.Tensor | None = None, # [n_cmp, d//QBLOCK] fp16, 8-bit pools only ) -> torch.Tensor: """Paged sparse MLA attention: gather KV from the two global pools inside the kernel. @@ -315,6 +345,22 @@ def sparse_attn_paged( # pools must share strides. Guaranteed here by .contiguous() on their [*, D] shape. assert window_pool.stride() == cmp_pool.stride(), (window_pool.stride(), cmp_pool.stride()) + # 8-bit pools carry a parallel scale array. Same base-pointer selection as the pools, so + # the same stride-sharing rule applies. + quant = window_scale is not None + assert quant == (cmp_scale is not None), "both pools are quantized or neither" + if quant: + window_scale = window_scale.contiguous() + cmp_scale = cmp_scale.contiguous() + assert window_scale.stride() == cmp_scale.stride(), ( + window_scale.stride(), cmp_scale.stride()) + assert d % window_scale.shape[1] == 0, (d, window_scale.shape) + qblock = d // window_scale.shape[1] + stride_sn, stride_sd = window_scale.stride() + else: + window_scale = cmp_scale = window_pool # unused; keeps the arg a valid pointer + qblock, stride_sn, stride_sd = 1, 0, 0 + has_counts = cmp_counts is not None if has_counts: cnt = cmp_counts.contiguous().to(torch.int32).view(b, m) @@ -325,19 +371,20 @@ def sparse_attn_paged( n_splits = split_count(b, m, h, topk, q.device) if n_splits: return _sparse_attn_paged_splitk( - q, window_pool, cmp_pool, sink, idx, cnt, o, + q, window_pool, cmp_pool, window_scale, cmp_scale, sink, idx, cnt, o, b, m, h, d, topk, n_window, softmax_scale, has_counts, stride_nb, stride_nm, - n_splits, + n_splits, quant, qblock, stride_sn, stride_sd, ) grid = (m, b, triton.cdiv(h, BLOCK_H)) _sparse_attn_paged_kernel[grid]( - q, window_pool, cmp_pool, o, sink, idx, cnt, + q, window_pool, cmp_pool, window_scale, cmp_scale, o, sink, idx, cnt, float(softmax_scale), h, topk, int(n_window), q.stride(0), q.stride(1), q.stride(2), q.stride(3), window_pool.stride(0), window_pool.stride(1), cmp_pool.stride(0), cmp_pool.stride(1), + stride_sn, stride_sd, o.stride(0), o.stride(1), o.stride(2), o.stride(3), idx.stride(0), idx.stride(1), idx.stride(2), stride_nb, stride_nm, @@ -345,6 +392,8 @@ def sparse_attn_paged( BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, + QUANT=quant, + QBLOCK=qblock, num_warps=8, num_stages=2, ) @@ -352,20 +401,22 @@ def sparse_attn_paged( def _sparse_attn_paged_splitk( - q, window_pool, cmp_pool, sink, idx, cnt, o, + q, window_pool, cmp_pool, window_scale, cmp_scale, sink, idx, cnt, o, b, m, h, d, topk, n_window, softmax_scale, has_counts, stride_nb, stride_nm, n_splits, + quant, qblock, stride_sn, stride_sd, ): head_blocks = triton.cdiv(h, BLOCK_H) mid_o = torch.empty((b, m, h, n_splits, d), dtype=torch.float32, device=q.device) mid_lse = torch.empty((b, m, h, n_splits), dtype=torch.float32, device=q.device) _sparse_attn_paged_splitk_kernel[(m * n_splits, b, head_blocks)]( - q, window_pool, cmp_pool, mid_o, mid_lse, idx, cnt, + q, window_pool, cmp_pool, window_scale, cmp_scale, mid_o, mid_lse, idx, cnt, float(softmax_scale), h, topk, int(n_window), q.stride(0), q.stride(1), q.stride(2), q.stride(3), window_pool.stride(0), window_pool.stride(1), cmp_pool.stride(0), cmp_pool.stride(1), + stride_sn, stride_sd, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4), mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3), idx.stride(0), idx.stride(1), idx.stride(2), @@ -374,6 +425,8 @@ def _sparse_attn_paged_splitk( BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, HAS_COUNTS=has_counts, + QUANT=quant, + QBLOCK=qblock, NUM_SPLITS=n_splits, num_warps=8, num_stages=2, diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 1bb352b4..50294bd9 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -79,6 +79,9 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt # DSV4 is driven by the generic CacheManager over the shared page table; the pool is # the only DSV4-specific piece (the swa_pool plug-in: window tier + cmp/idx/state # shadows). Sizing reads dsv4_args, never the group spec. + # Stamped on dsv4_args at config resolution (see _validate_kv_cache_dtype), because + # --moe-cache-auto sizes the expert cache off the KV budget before this runs. + kv_quant = bool(getattr(model_config.dsv4_args, "kv_quant", False)) pool = DSV4PagedKVCache( sizes=_dsv4_pool_sizes(config, num_pages + 1), # +1 for dummy page args=model_config.dsv4_args, @@ -86,6 +89,7 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt dtype=dtype, P=model_config.dsv4_args.window_size, n_scratch=config.max_running_req + 1, + kv_quant=kv_quant, ) pool._init_paged_state(config.max_running_req, config.cache_type != "naive") return pool diff --git a/python/freetoken/kvcache/dsv4_cost_model.py b/python/freetoken/kvcache/dsv4_cost_model.py index 2c084aff..389d9195 100644 --- a/python/freetoken/kvcache/dsv4_cost_model.py +++ b/python/freetoken/kvcache/dsv4_cost_model.py @@ -19,10 +19,12 @@ from __future__ import annotations +import math import os - from dataclasses import dataclass, field +from .dsv4_kv_quant import BYTES_PER_ELEMENT as KV_QUANT_BYTES_PER_ELEMENT + _BF16_BYTES = 2 _FP32_BYTES = 4 @@ -48,6 +50,16 @@ def ring_size_for_ratio(ratio: int) -> int: def _kv_bytes(args) -> int: + """Bytes per element in the window / compressed KV tiers. + + Under ``kv_quant`` those two pools store fp8 plus one fp16 scale per 32 elements + (1.0625 B/elem); the indexer tier and the state rings are unaffected. This has to + agree with what the pool actually allocates: ``--moe-cache-auto`` divides the VRAM + budget by these numbers, so a stale 2 here would over-size the expert cache and then + fail at allocation. + """ + if getattr(args, "kv_quant", False): + return math.ceil(args.head_dim * KV_QUANT_BYTES_PER_ELEMENT) return args.head_dim * _BF16_BYTES diff --git a/python/freetoken/kvcache/dsv4_kv_quant.py b/python/freetoken/kvcache/dsv4_kv_quant.py new file mode 100644 index 00000000..f989b484 --- /dev/null +++ b/python/freetoken/kvcache/dsv4_kv_quant.py @@ -0,0 +1,77 @@ +"""8-bit storage for the DeepSeek-V4 window / compressed KV pools. + +DSV4's attention KV is one MLA latent row per token -- K and V are the same slab -- +so unlike a separate-K-and-V pool there is a single buffer to quantize and a single +scale array beside it. + +Storage is fp8-e4m3 held as ``uint8``, with one fp16 scale per :data:`BLOCK` elements +along ``head_dim``:: + + 1 byte + 2/32 = 1.0625 bytes per element, against bf16's 2 + +The scale varies along ``head_dim``, which is the reduction dimension of ``q @ kv``, so +the attention kernel cannot fold it in after the dot: it dequantizes the gathered tile +before the dot. Storage is what this buys, not tensor-core throughput. + +Why the precision cost is small here. The model already rounds this KV onto the e4m3 +grid before it reaches the pool -- ``act_quant_fp8_inplace(kv[..., :-rope_head_dim], 64)`` +runs immediately before every ``store_window`` and before the compressor's scatter -- so +for those dims this is closer to a storage change than a precision change. Re-blocking +from 64 to 32 does move values (a 32-block's own max-abs sets a different scale than the +64-block's did), but only within the grid the values already live on. The rope tail is +the exception: those dims are genuine bf16 and this rounds them for the first time. + +Needs native fp8 (sm_89+). Below that a float8 pointer is illegal inside a triton +kernel, and DSV4 checkpoints are fp8 to begin with, so the flag is rejected at config +time rather than carrying an emulated decode nobody would run. +""" + +from __future__ import annotations + +import torch + +# Elements per scale, along head_dim. +BLOCK = 32 +SCALE_DTYPE = torch.float16 +STORAGE_DTYPE = torch.float8_e4m3fn +# Max-abs of a block maps to this magnitude, e4m3's largest finite value. +MAX_MAG = 448.0 + +BYTES_PER_ELEMENT = 1.0 + SCALE_DTYPE.itemsize / BLOCK # 1.0625 + + +def scale_width(head_dim: int) -> int: + """Scales per row. ``head_dim`` must be a multiple of :data:`BLOCK`.""" + if head_dim % BLOCK: + raise ValueError(f"head_dim {head_dim} is not a multiple of the quant block {BLOCK}") + return head_dim // BLOCK + + +def quantize_ref(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Torch reference for the store kernel: ``[rows, D]`` -> ``(fp8, fp16 scales)``. + + The scale is rounded to its stored precision *before* dividing, so the value written + here and the value the attention kernel reads back are scaled by the identical number. + """ + rows, d = x.shape + nb = scale_width(d) + xb = x.float().reshape(rows, nb, BLOCK) + amax = xb.abs().amax(dim=-1) + scale = torch.where(amax > 0, amax / MAX_MAG, torch.ones_like(amax)) + scale = scale.to(SCALE_DTYPE).float() + q = (xb / scale[..., None]).clamp(-MAX_MAG, MAX_MAG) + return q.to(STORAGE_DTYPE).reshape(rows, d), scale.to(SCALE_DTYPE) + + +def dequantize_ref(q: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Inverse of :func:`quantize_ref`, in fp32.""" + rows, d = q.shape + nb = scale_width(d) + x = q.float().reshape(rows, nb, BLOCK) + return (x * scale.float()[..., None]).reshape(rows, d) + + +__all__ = [ + "BLOCK", "SCALE_DTYPE", "STORAGE_DTYPE", "MAX_MAG", "BYTES_PER_ELEMENT", + "scale_width", "quantize_ref", "dequantize_ref", +] diff --git a/python/freetoken/kvcache/dsv4_paged_pool.py b/python/freetoken/kvcache/dsv4_paged_pool.py index 3ce23eea..681dbc5a 100644 --- a/python/freetoken/kvcache/dsv4_paged_pool.py +++ b/python/freetoken/kvcache/dsv4_paged_pool.py @@ -8,9 +8,11 @@ * ``state_ring[L]`` -- ratio>0 layers; the per-window-page compress-state ring (fp32, ``kv|score`` split), with index ``-1`` a permanent scratch slot. -The KV/compressed/indexer pools are bf16 (the fp8/fp4 quant is an in-place round-trip already -baked into the bf16 value, so ``index_select`` staging is byte-exact); only the compress-state -ring is fp32. +The indexer pool is bf16 and the compress-state ring is fp32. The window and compressed +pools are bf16 by default -- the fp8/fp4 quant is an in-place round-trip already baked into +the bf16 value, so ``index_select`` staging is byte-exact -- or fp8 with a parallel fp16 +scale array under ``kv_quant`` (see :mod:`freetoken.kvcache.dsv4_kv_quant`), which stores +those already-rounded values in 8 bits instead of 16. ``state_loc`` is DERIVED from a window slot, never stored: state_loc = where(ws < 0, -1, (ws // P) * ring_size + ws % ring_size) @@ -30,6 +32,11 @@ dsv4_window_unit_bytes, ring_size_for_ratio, ) +from .dsv4_kv_quant import ( + SCALE_DTYPE as KV_SCALE_DTYPE, + STORAGE_DTYPE as KV_STORAGE_DTYPE, + scale_width, +) logger = init_logger(__name__) @@ -151,6 +158,13 @@ def set_blocks(self, page_base: torch.Tensor, blocks: torch.Tensor) -> None: self._clear_scratch() +def _store_kv_quant(pool, scales, slots, kv) -> None: + """Quantizing counterpart of ``index_copy_`` for an 8-bit pool.""" + from freetoken.kernel.triton.dsv4.kv_quant import store_kv_quant + + store_kv_quant(pool, scales, slots, kv.reshape(-1, pool.shape[1])) + + class DSV4PagedKVCache(BaseKVCachePool): def __init__( self, @@ -160,8 +174,15 @@ def __init__( dtype: torch.dtype = torch.bfloat16, P: int = 128, n_scratch: int = 1, + kv_quant: bool = False, ) -> None: - assert dtype == torch.bfloat16, "KV pools are bf16 (fp4/fp8 is an in-place round-trip)" + assert dtype == torch.bfloat16, "the compute dtype is bf16" + # 8-bit storage for the window / compressed attention KV (see dsv4_kv_quant). The + # indexer pool and the compress-state rings keep their dtype: the indexer's K decides + # top-k *selection*, and the rings are fp32 accumulators, neither of which is a + # storage-bandwidth problem worth a routing or accumulation risk. + self.kv_quant = bool(kv_quant) + self._kv_storage_dtype = KV_STORAGE_DTYPE if self.kv_quant else dtype self.args = args self.sizes = sizes self._device = device @@ -208,10 +229,20 @@ def _alloc_buffers(self) -> None: self.full_loc_map: torch.Tensor | None = None # Window KV: every layer. + kvdt = self._kv_storage_dtype self.window_pool: list[torch.Tensor] = [ - torch.zeros(sizes.n_win_slots, self.head_dim, device=device, dtype=dtype) + torch.zeros(sizes.n_win_slots, self.head_dim, device=device, dtype=kvdt) for _ in range(self._n_layers) ] + # Parallel scale arrays, one fp16 per quant block along head_dim. Empty when the + # pools are bf16, so the unquantized path allocates nothing extra. + nb = scale_width(self.head_dim) if self.kv_quant else 0 + self.window_scale: list[torch.Tensor | None] = [ + torch.zeros(sizes.n_win_slots, nb, device=device, dtype=KV_SCALE_DTYPE) + if self.kv_quant else None + for _ in range(self._n_layers) + ] + self.cmp_scale: list[torch.Tensor | None] = [] # Compressed KV / Indexer KV / compress-state ring: per ratio-class. # ``state_ring`` is the ATTENTION compressor's ring (head_dim). Ratio-4 @@ -225,6 +256,7 @@ def _alloc_buffers(self) -> None: ratio = self.compress_ratios[L] if ratio == 0: self.cmp_pool.append(None) + self.cmp_scale.append(None) self.idx_pool.append(None) self.state_ring.append(None) self.indexer_state_ring.append(None) @@ -235,9 +267,14 @@ def _alloc_buffers(self) -> None: self.cmp_scratch_base.append(sizes.cmp_blocks[L]) self.cmp_pool.append( torch.zeros( - sizes.cmp_blocks[L] + self.n_scratch, self.head_dim, device=device, dtype=dtype + sizes.cmp_blocks[L] + self.n_scratch, self.head_dim, device=device, dtype=kvdt ) ) + self.cmp_scale.append( + torch.zeros(sizes.cmp_blocks[L] + self.n_scratch, nb, + device=device, dtype=KV_SCALE_DTYPE) + if self.kv_quant else None + ) if ratio == 4: self.idx_scratch_base.append(sizes.idx_blocks[L]) self.idx_pool.append( @@ -557,11 +594,18 @@ def set_state(self, layer_id: int, state_loc: torch.Tensor, kv_score: torch.Tens # ----- specialized writes ----- def store_window(self, k: torch.Tensor, layer_id: int, window_slot: torch.Tensor) -> None: + if self.kv_quant: + _store_kv_quant( + self.window_pool[layer_id], self.window_scale[layer_id], window_slot, k) + return self.window_pool[layer_id].index_copy_(0, window_slot, k.to(self._dtype)) def store_compressed(self, kv: torch.Tensor, layer_id: int, cmp_slot: torch.Tensor) -> None: pool = self.cmp_pool[layer_id] assert pool is not None, f"layer {layer_id} (ratio 0) has no compressed pool" + if self.kv_quant: + _store_kv_quant(pool, self.cmp_scale[layer_id], cmp_slot, kv) + return pool.index_copy_(0, cmp_slot, kv.to(self._dtype)) def store_indexer(self, k: torch.Tensor, layer_id: int, idx_slot: torch.Tensor) -> None: diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b75..8f0ca9c3 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -507,6 +507,19 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="KV-cache token floor reserved before --moe-cache-auto fills experts.", ) + parser.add_argument( + "--kv-cache-dtype", + dest="kv_cache_dtype", + choices=("auto", "fp8_e4m3"), + default=ServerArgs.kv_cache_dtype, + help=( + "Storage dtype for the DeepSeek-V4 window/compressed KV pools. 'auto' keeps " + "bf16; 'fp8_e4m3' stores fp8 plus one fp16 scale per 32 elements (1.0625 " + "bytes/element against 2), which frees VRAM for a larger KV pool or a larger " + "expert cache. Needs a DeepSeek-V4 checkpoint and native fp8 (sm_89+)." + ), + ) + parser.add_argument( "--moe-cache-policy", default=ServerArgs.moe_cache_policy, diff --git a/tests/kvcache/test_dsv4_kv_quant.py b/tests/kvcache/test_dsv4_kv_quant.py new file mode 100644 index 00000000..f4bb9d41 --- /dev/null +++ b/tests/kvcache/test_dsv4_kv_quant.py @@ -0,0 +1,219 @@ +"""8-bit storage for the DSV4 window / compressed KV pools.""" +import pytest +import torch + +from freetoken.kvcache.dsv4_kv_quant import ( + BLOCK, + BYTES_PER_ELEMENT, + dequantize_ref, + quantize_ref, + scale_width, +) + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel") + +D = 512 # DSV4-Flash head_dim + + +def _native_fp8() -> bool: + return torch.cuda.get_device_capability() >= (8, 9) + + +def test_bytes_per_element_is_what_the_cost_model_assumes(): + assert BYTES_PER_ELEMENT == 1.0 + 2 / BLOCK == 1.0625 + assert scale_width(D) == D // BLOCK + + +def test_scale_width_rejects_a_ragged_head_dim(): + with pytest.raises(ValueError, match="multiple of the quant block"): + scale_width(D + 1) + + +@pytest.mark.parametrize("rows", [1, 7, 512]) +def test_round_trip_error_is_one_e4m3_step(rows): + """e4m3 is a floating format, so the error is relative, not a uniform step of the + block scale: 3 mantissa bits give at most 2**-4 relative on a normal value. + + Below the smallest normal (``scale * 2**-6``) the values are e4m3 subnormals, whose + step is fixed at ``scale * 2**-9``, so relative error there is unbounded and only the + absolute half-step binds. The bound is the sum of the two. + """ + torch.manual_seed(0) + x = torch.randn(rows, D, device="cuda", dtype=torch.bfloat16) + q, s = quantize_ref(x) + back = dequantize_ref(q, s) + xf = x.float() + step = s.float().repeat_interleave(BLOCK, dim=-1) # the block's own scale + bound = 2.0 ** -4 * xf.abs() + step * (2.0 ** -9) / 2 + assert ((back - xf).abs() <= bound * 1.01).all() + + +def test_an_all_zero_row_survives(): + x = torch.zeros(4, D, device="cuda", dtype=torch.bfloat16) + q, s = quantize_ref(x) + assert torch.equal(dequantize_ref(q, s), torch.zeros(4, D, device="cuda")) + + +def test_values_already_on_the_e4m3_grid_are_reproduced_closely(): + """The model rounds KV onto the e4m3 grid before it reaches the pool, so this is the + regime that matters: re-blocking 64 -> 32 must not move values much.""" + torch.manual_seed(0) + x = torch.randn(64, D, device="cuda", dtype=torch.float32) + on_grid = x.to(torch.float8_e4m3fn).float() # what act_quant_fp8_inplace leaves behind + q, s = quantize_ref(on_grid) + err = (dequantize_ref(q, s) - on_grid).abs() / on_grid.abs().clamp(min=1e-6) + assert err.max() < 0.13, err.max() # within one e4m3 step (2^-3) + + +@pytest.mark.skipif(not _native_fp8(), reason="needs native fp8 (sm_89+)") +@pytest.mark.parametrize("rows", [1, 5, 128]) +def test_store_kernel_matches_the_reference(rows): + from freetoken.kernel.triton.dsv4.kv_quant import store_kv_quant + + torch.manual_seed(0) + slots_total = 256 + kv = torch.randn(rows, D, device="cuda", dtype=torch.bfloat16) + slots = torch.randperm(slots_total, device="cuda")[:rows].to(torch.int32) + + pool = torch.zeros(slots_total, D, device="cuda", dtype=torch.float8_e4m3fn) + scales = torch.zeros(slots_total, scale_width(D), device="cuda", dtype=torch.float16) + store_kv_quant(pool, scales, slots, kv) + + ref_q, ref_s = quantize_ref(kv) + sl = slots.long() + assert torch.equal(pool[sl].view(torch.uint8), ref_q.view(torch.uint8)) + assert torch.equal(scales[sl], ref_s) + + +@pytest.mark.skipif(not _native_fp8(), reason="needs native fp8 (sm_89+)") +def test_store_kernel_leaves_other_slots_untouched(): + from freetoken.kernel.triton.dsv4.kv_quant import store_kv_quant + + pool = torch.zeros(32, D, device="cuda", dtype=torch.float8_e4m3fn) + scales = torch.zeros(32, scale_width(D), device="cuda", dtype=torch.float16) + kv = torch.randn(4, D, device="cuda", dtype=torch.bfloat16) + slots = torch.tensor([1, 3, 5, 7], device="cuda", dtype=torch.int32) + store_kv_quant(pool, scales, slots, kv) + + untouched = [i for i in range(32) if i not in (1, 3, 5, 7)] + assert pool[untouched].view(torch.uint8).eq(0).all() + assert scales[untouched].eq(0).all() + + +@pytest.mark.skipif(not _native_fp8(), reason="needs native fp8 (sm_89+)") +def test_store_kernel_handles_saturation(): + """A block whose max-abs is huge still maps that max onto MAX_MAG, not to inf.""" + from freetoken.kernel.triton.dsv4.kv_quant import store_kv_quant + + kv = torch.full((2, D), 3.0e4, device="cuda", dtype=torch.bfloat16) + pool = torch.zeros(8, D, device="cuda", dtype=torch.float8_e4m3fn) + scales = torch.zeros(8, scale_width(D), device="cuda", dtype=torch.float16) + store_kv_quant(pool, scales, torch.tensor([0, 1], device="cuda", dtype=torch.int32), kv) + back = dequantize_ref(pool[:2], scales[:2]) + assert torch.isfinite(back).all() + assert (back - kv.float()).abs().max() / 3.0e4 < 0.02 + + +# ---------------------------------------------------------------- attention read path +@pytest.mark.skipif(not _native_fp8(), reason="needs native fp8 (sm_89+)") +@pytest.mark.parametrize("m", [1, 4]) +def test_quantized_attention_tracks_the_bf16_reference(m): + """The kernel dequantizes the gathered tile before the dot. Against the same pools + stored bf16, the output should differ only by the storage rounding.""" + from freetoken.kernel.triton.dsv4.sparse_attn import sparse_attn_paged + + torch.manual_seed(0) + b, h, topk, n_window = 2, 8, 64, 32 + n_win_slots, n_cmp = 128, 128 + + q = torch.randn(b, m, h, D, device="cuda", dtype=torch.bfloat16) + win = torch.randn(n_win_slots, D, device="cuda", dtype=torch.bfloat16) + cmp = torch.randn(n_cmp, D, device="cuda", dtype=torch.bfloat16) + sink = torch.randn(h, device="cuda", dtype=torch.float32) + idx = torch.randint(0, 128, (b, m, topk), device="cuda", dtype=torch.int32) + + ref = sparse_attn_paged(q, win, cmp, sink, idx, n_window, D ** -0.5) + + win_q, win_s = quantize_ref(win) + cmp_q, cmp_s = quantize_ref(cmp) + got = sparse_attn_paged(q, win_q, cmp_q, sink, idx, n_window, D ** -0.5, + window_scale=win_s, cmp_scale=cmp_s) + + assert got.shape == ref.shape and got.dtype == ref.dtype + # Attention averages over topk columns, so the per-element storage error largely + # cancels; what survives is well inside an e4m3 step of the output's own scale. + err = (got.float() - ref.float()).abs().max() / ref.float().abs().max() + assert err < 0.05, err + + +@pytest.mark.skipif(not _native_fp8(), reason="needs native fp8 (sm_89+)") +def test_quantized_attention_honours_masked_columns(): + """-1 columns must stay masked: a scale load for a masked column must not resurrect it.""" + from freetoken.kernel.triton.dsv4.sparse_attn import sparse_attn_paged + + torch.manual_seed(0) + b, m, h, topk, n_window = 1, 1, 4, 32, 16 + q = torch.randn(b, m, h, D, device="cuda", dtype=torch.bfloat16) + win = torch.randn(64, D, device="cuda", dtype=torch.bfloat16) + cmp = torch.randn(64, D, device="cuda", dtype=torch.bfloat16) + sink = torch.zeros(h, device="cuda", dtype=torch.float32) + idx = torch.full((b, m, topk), -1, device="cuda", dtype=torch.int32) + idx[..., :4] = torch.arange(4, device="cuda", dtype=torch.int32) + + win_q, win_s = quantize_ref(win) + cmp_q, cmp_s = quantize_ref(cmp) + ref = sparse_attn_paged(q, win, cmp, sink, idx, n_window, D ** -0.5) + got = sparse_attn_paged(q, win_q, cmp_q, sink, idx, n_window, D ** -0.5, + window_scale=win_s, cmp_scale=cmp_s) + assert torch.isfinite(got.float()).all() + err = (got.float() - ref.float()).abs().max() / ref.float().abs().max().clamp(min=1e-6) + assert err < 0.05, err + + +@pytest.mark.skipif(not _native_fp8(), reason="needs native fp8 (sm_89+)") +def test_quantized_attention_on_the_splitk_decode_path(): + """The decode path is a different kernel with its own gather, so it needs its own + check: m == 1 and topk > MIN_TILES_PER_SPLIT * BLOCK_T is what selects it.""" + from freetoken.kernel.triton.dsv4.sparse_attn import sparse_attn_paged, split_count + + torch.manual_seed(0) + b, m, h, topk, n_window = 1, 1, 16, 512, 128 # DSV4's index_topk + assert split_count(b, m, h, topk, "cuda") > 1, "shape did not select the split-k kernel" + + q = torch.randn(b, m, h, D, device="cuda", dtype=torch.bfloat16) + win = torch.randn(256, D, device="cuda", dtype=torch.bfloat16) + cmp = torch.randn(256, D, device="cuda", dtype=torch.bfloat16) + sink = torch.randn(h, device="cuda", dtype=torch.float32) + idx = torch.randint(0, 256, (b, m, topk), device="cuda", dtype=torch.int32) + + ref = sparse_attn_paged(q, win, cmp, sink, idx, n_window, D ** -0.5) + win_q, win_s = quantize_ref(win) + cmp_q, cmp_s = quantize_ref(cmp) + got = sparse_attn_paged(q, win_q, cmp_q, sink, idx, n_window, D ** -0.5, + window_scale=win_s, cmp_scale=cmp_s) + err = (got.float() - ref.float()).abs().max() / ref.float().abs().max() + assert err < 0.05, err + + +@pytest.mark.skipif(not _native_fp8(), reason="needs native fp8 (sm_89+)") +def test_quantized_attention_respects_cmp_counts(): + """cmp_counts bounds the compressed half from device memory; the dequant must not + read past it into stale slots.""" + from freetoken.kernel.triton.dsv4.sparse_attn import sparse_attn_paged + + torch.manual_seed(0) + b, m, h, topk, n_window = 1, 1, 16, 512, 128 + q = torch.randn(b, m, h, D, device="cuda", dtype=torch.bfloat16) + win = torch.randn(256, D, device="cuda", dtype=torch.bfloat16) + cmp = torch.randn(256, D, device="cuda", dtype=torch.bfloat16) + sink = torch.randn(h, device="cuda", dtype=torch.float32) + idx = torch.randint(0, 256, (b, m, topk), device="cuda", dtype=torch.int32) + counts = torch.full((b, m), 40, device="cuda", dtype=torch.int32) + + win_q, win_s = quantize_ref(win) + cmp_q, cmp_s = quantize_ref(cmp) + ref = sparse_attn_paged(q, win, cmp, sink, idx, n_window, D ** -0.5, cmp_counts=counts) + got = sparse_attn_paged(q, win_q, cmp_q, sink, idx, n_window, D ** -0.5, + cmp_counts=counts, window_scale=win_s, cmp_scale=cmp_s) + err = (got.float() - ref.float()).abs().max() / ref.float().abs().max() + assert err < 0.05, err