feat(kvcache): 8-bit DSV4 window/compressed KV behind --kv-cache-dtype

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 ---------------------------------------------------------

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 --------------------------------------------------------------------

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

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

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"]

python/freetoken/kernel/triton/dsv4/sparse_attn.py

@@ -50,19 +50,22 @@
 
 @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,
     D: tl.constexpr,
     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,47 +371,52 @@ 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,
         D=d,
         BLOCK_H=BLOCK_H,
         BLOCK_T=BLOCK_T,
         HAS_COUNTS=has_counts,
+        QUANT=quant,
+        QBLOCK=qblock,
         num_warps=8,
         num_stages=2,
     )
     return o
 
 
 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,

python/freetoken/kvcache/__init__.py

@@ -79,13 +79,17 @@ 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,
             device=device,
             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