feat(kvcache): store the KV cache as fp8 e4m3 codes (--kv-cache-dtype…

docs/cli.md

@@ -71,8 +71,32 @@ ft serve --model ... --gpu GPU-9e8d7c6b  # the same card by UUID (a unique prefi
 | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) |
 | `--page-size` | 1 | KV page size; DSV4 forces 128, the TRTLLM backend needs 16/32/64, SWA models require 1 |
 | `--cache-type` | radix | `radix` (prefix reuse; SWA/GDN-aware variants picked automatically) or `naive` |
+| `--kv-cache-dtype` | bf16 | `bf16` or `fp8`: store the KV cache as e4m3 codes plus one fp32 scale per (token, kv head), roughly doubling the tokens that fit in the same VRAM; see [FP8 KV cache](#fp8-kv-cache) |
 | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU |
 
+### FP8 KV cache
+
+`ft serve --kv-cache-dtype fp8` halves the bytes per cached token (8-bit codes instead
+of 16), so a card that held N tokens holds close to 2N. Each `(token, kv head)` row
+keeps its own fp32 scale, which costs ~3% back at `head_dim=128`. Requirements and
+trade-offs:
+
+- Needs the **triton** attention backend; `--attn auto` selects it (and refuses an
+  explicit `fi`/`fa`/`trtllm`, which cannot be shown to apply these scales).
+- Works on the plain paged, hybrid-SWA and QSA sparse (Qwen3.8-Flash-Next) KV pools.
+  On QSA the block-selection index keys stay 16-bit; only the selected K/V rows are
+  read back as codes. MLA/DSA latent KV, DeepSeek-V4's tiered pool and the block-sparse
+  MiniMax-M3 pool stay 16-bit; asking for fp8 there fails at startup rather than
+  silently ignoring the flag.
+- The same bytes on every GPU FreeToken targets: the codes sit in a plain byte buffer
+  and are decoded in software, so the cache holds identical data and produces identical
+  numbers on any card (the fp8 type is deliberately kept out of the kernels, which is
+  also what makes the feature work on the RTX 30 series).
+- Accuracy is checkpoint-dependent. Expect it to matter most on long contexts and on
+  models with outlier key channels; keep `bf16` when a run must be bit-reproducible.
+- `ft ctl stats` / `/v1/cache/status` report the smaller `kv_bytes_per_token`, and
+  `ft ctl cache --kv N` moves the same (now cheaper) pool.
+
 ### MoE offload
 
 See [models.md](models.md#moe-backends) for what each backend does.

docs/models.md

@@ -41,3 +41,9 @@ for them; other checkpoints of the same architectures work too.
   authoritative model args are read from there.
 - Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM.
 - Multimodal checkpoints are served text-only.
+- `--kv-cache-dtype fp8` (see [cli.md](cli.md#fp8-kv-cache)) covers the plain paged,
+  hybrid-SWA and QSA sparse KV pools — gpt-oss, Qwen3/3.5/3.6, GLM-4.x, Gemma-4,
+  MiniMax-M2.5, Muse-Glimmer, Llama/Qwen2/Mistral, Qwen3.8-Flash-Next (on QSA only the
+  selected K/V rows are read back as codes; block selection keeps 16-bit index keys).
+  MLA/DSA (GLM-5.2), DeepSeek-V4's tiered pool and MiniMax-M3's block-sparse pool stay
+  16-bit and reject it.

python/freetoken/attention/__init__.py

@@ -33,6 +33,11 @@ class BackendInfo:
     # Whether forward() honors a per-call AttentionSpec (window/sm_scale/sinks).
     # Non-consumers raise on a non-None spec instead of silently dropping it.
     consumes_attn_spec: bool = False
+    # Whether forward() reads an fp8 KV pool (codes + per-token/per-head scales).
+    # Backends that hand the cache to an external kernel must opt out until that
+    # kernel is proven to apply our scale layout; the engine then refuses (or auto-
+    # avoids) them for --kv-cache-dtype fp8.
+    supports_fp8_kv: bool = False
 
 
 SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend")
@@ -84,6 +89,7 @@ def create_fa_backend(config: ModelConfig):
     BackendInfo(
         supported_types=frozenset({AttnType.FULL, AttnType.SWA}),
         consumes_attn_spec=True,
+        supports_fp8_kv=True,
     ),
 )
 def create_triton_backend(config: ModelConfig):
@@ -137,6 +143,9 @@ def create_m3_sparse_backend(config: ModelConfig):
     "qsa_sparse",
     BackendInfo(
         supported_types=frozenset({AttnType.QSA}),
+        # The attend kernel dequantizes on load (kernel/triton/qsa/attend.py); the
+        # compressed index keys it scores against are a separate, always-16-bit tier.
+        supports_fp8_kv=True,
         # 64-token pages: a 4-token compress group never straddles a page, so the
         # compressed row of a group is page_base // 4 + block-in-page.
         page_sizes=(64,),

python/freetoken/attention/qsa_sparse.py

@@ -108,7 +108,17 @@ def __init__(self, config: ModelConfig) -> None:
             f"qsa_sparse backend needs a QSA pool, got {type(self.kvcache).__name__}"
         )
         self.device = self.kvcache.device
+        # The pool's COMPUTE dtype, never its store dtype (the contract lives in
+        # kvcache/base.py). These buffers feed the indexer -- qsa_index_norm_rope and
+        # qsa_mqa_paged -- whose tl.dot has no fp8 path, so an e4m3 q_index does not
+        # fail here, it fails at CUDA-graph capture with "Unsupported rhs dtype
+        # fp8e4nv". --kv-cache-dtype fp8 quantizes only the KV tiers; the index tiers
+        # stay 16-bit by design (kvcache/qsa_pool.py).
         self.dtype = self.kvcache.dtype
+        assert self.dtype.itemsize == 2, (
+            f"QSA block selection needs a 16-bit compute dtype, got {self.dtype} -- "
+            "the KV pool must report its compute dtype, not e4m3 codes"
+        )
         self.index_head_dim = self.kvcache.index_head_dim
         self.ratio = self.kvcache.index_ratio
         self.ring_capacity = self.kvcache.ring_capacity
@@ -282,6 +292,8 @@ def qsa_forward(
 
         self._update_index_cache(index, md, slot)
         indices = self._select(index, md, slot)
+        # Scale tensors only exist on an fp8 pool (k_scale returns None otherwise); the
+        # index tier stays bf16 either way, so _select above is quantization-agnostic.
         return qsa_sparse_paged_attention(
             q,
             self.kvcache.k_cache(layer_id),
@@ -290,6 +302,8 @@ def qsa_forward(
             md.block_table,
             md.token_to_req,
             torch.empty_like(q),
+            k_scale=self.kvcache.k_scale(layer_id),
+            v_scale=self.kvcache.v_scale(layer_id),
         )
 
     def _plan_index_writes(self, md: QSASparseMetadata, batch: Batch) -> None:

python/freetoken/attention/triton.py

@@ -155,6 +155,11 @@ 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 fp8 KV pool hands us its per-(token, head) scales; a 16-bit pool returns
+        # None and every kernel below keeps its original (scale-free) code path.
+        k_scale = self.kvcache.k_scale(layer_id)
+        v_scale = self.kvcache.v_scale(layer_id)
+        assert (k_scale is None) == (v_scale is None), "K and V scales come as a pair"
 
         spec = attn_spec or AttentionSpec()
         indices = metadata.indices
@@ -181,6 +186,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 +208,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 +222,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:

python/freetoken/engine/config.py

@@ -59,6 +59,12 @@ class EngineConfig:
     cuda_graph_bs: List[int] | None = None
     cuda_graph_max_bs: int | None = None
     page_size: int = 1
+    # KV-cache storage quantization: "none" stores the compute dtype, "fp8" stores e4m3
+    # codes plus one fp32 scale per (token, slab, layer, kv head) -- about 2x the tokens
+    # per GiB, at a small accuracy cost. --kv-cache-dtype; resolved from "auto" by
+    # _adjust_config, which also refuses it on a pool family or attention backend that
+    # cannot read the scales.
+    kv_quant: str = "none"
     memory_ratio: float = 0.9
     # Hybrid GDN models default to the HybridRadixCache (cross-request GDN-state prefix reuse);
     # `--cache-type naive` opts out. linear_state_cache_ratio sizes the GDN snapshot cache as

python/freetoken/engine/engine.py

@@ -114,10 +114,38 @@ def _backend_requirements_met(name: str) -> bool:
     return True
 
 
-def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str:
+# --kv-cache-dtype spellings -> the stored EngineConfig.kv_quant value.
+KV_QUANT_ALIASES = {"auto": "none", "bf16": "none", "none": "none", "fp8": "fp8"}
+
+
+def _resolve_kv_quant(value: str | None) -> str:
+    """Normalize a --kv-cache-dtype spelling to EngineConfig.kv_quant."""
+    key = (value or "auto").strip().lower()
+    if key not in KV_QUANT_ALIASES:
+        raise ValueError(
+            f"unknown --kv-cache-dtype {value!r}; expected one of "
+            f"{', '.join(sorted(KV_QUANT_ALIASES))}"
+        )
+    return KV_QUANT_ALIASES[key]
+
+
+def _backend_supports_kv_quant(name: str, kv_quant: str) -> bool:
+    """Whether every comma part of an attention-backend string can read a quantized
+    KV pool (an unquantized pool needs nothing from the backend)."""
+    if kv_quant == "none":
+        return True
+    return all(
+        attention_backend_info(part.strip()).supports_fp8_kv for part in name.split(",")
+    )
+
+
+def _resolve_auto_attention_backend(
+    required: frozenset[AttnType], *, kv_quant: str = "none"
+) -> str:
     """First candidate (in per-type priority order) whose arch condition holds,
-    whose packages are installed, and whose every comma part serves ALL required
-    types. Reproduces the historical hardware tree for FULL-only models:
+    whose packages are installed, whose every comma part serves ALL required
+    types, and which can decode a quantized KV cache when one is configured.
+    Reproduces the historical hardware tree for FULL-only models:
     sm_100 -> trtllm, sm_90+sgl_kernel -> "fa,fi", flashinfer -> fi, else triton."""
     candidates: list[tuple[str, bool]] = []
     if AttnType.DSV4 in required:
@@ -144,10 +172,18 @@ def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str:
             continue
         if not _backend_requirements_met(name):
             continue
+        if not _backend_supports_kv_quant(name, kv_quant):
+            continue
         return name
     raise RuntimeError(
         "No attention backend can serve attention types "
-        f"{sorted(t.value for t in required)} on this machine."
+        f"{sorted(t.value for t in required)} on this machine"
+        + (
+            f" with a {kv_quant} KV cache"
+            if kv_quant != "none"
+            else ""
+        )
+        + "."
     )
 
 
@@ -192,6 +228,23 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att
                 f"SWA models require, got {config.attention_backend!r}."
             )
 
+    # A quantized KV pool is only readable by a backend that applies its per-(token,
+    # head) scales; one that hands the cache to an external kernel would silently
+    # attend to raw e4m3 codes. Rejected here, before any weight is resident.
+    kv_quant = getattr(config, "kv_quant", "none")
+    if not _backend_supports_kv_quant(config.attention_backend, kv_quant):
+        fp8_backends = [
+            name
+            for name in ("trtllm", "fi", "fa", "triton")
+            if required <= attention_backend_info(name).supported_types
+            and attention_backend_info(name).supports_fp8_kv
+        ]
+        raise ValueError(
+            f"--kv-cache-dtype {kv_quant} needs an attention backend that decodes the KV "
+            f"scales; {config.attention_backend!r} does not. Valid for this model: "
+            f"{', '.join(fp8_backends) or 'none'} (or use --kv-cache-dtype bf16)."
+        )
+
     # An explicitly-selected backend may require a package that isn't installed. Auto
     # never resolves to one of these when its package is missing, so this only fires for
     # explicit --attention-backend choices.
@@ -1301,6 +1354,27 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
     # lists, then validate whatever is now selected (explicit or auto) -- every
     # comma part must serve every required type, with packages/arch available.
     required_attn_types = _required_attn_types(model_config)
+    # Resolve KV quantization BEFORE the backend tree: a quantized pool narrows both
+    # which pool families are usable and which backend auto may pick.
+    kv_quant = _resolve_kv_quant(getattr(config, "kv_quant", "none"))
+    override("kv_quant", kv_quant)
+    if kv_quant != "none":
+        # fp8 codes are wired through the pools that hand their rows to a Triton
+        # kernel: the plain paged and hybrid-SWA ones, plus the QSA sparse pool, whose
+        # index tier stays bf16 -- only the selected tokens come back as codes.
+        # Everything else (MLA's absorbed cache, DSA/DSV4/BSA sparse) has kernels that
+        # assert on 16-bit rows, and kvcache/__init__.py rejects fp8 for those families
+        # at pool creation.
+        quant_unsupported = required_attn_types - {
+            AttnType.FULL, AttnType.SWA, AttnType.QSA,
+        }
+        if quant_unsupported:
+            raise ValueError(
+                f"--kv-cache-dtype {kv_quant} is implemented for the plain paged, "
+                "hybrid-SWA and QSA sparse KV pools; this model also needs "
+                f"{', '.join(sorted(t.value for t in quant_unsupported))} attention "
+                "(use --kv-cache-dtype bf16)."
+            )
     _dtype = getattr(config, "dtype", None)  # duck-typed test configs omit it
     if (
         required_attn_types & {AttnType.BSA, AttnType.QSA}
@@ -1328,7 +1402,7 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
     if config.attention_backend == "auto":
         override(
             "attention_backend",
-            _resolve_auto_attention_backend(required_attn_types),
+            _resolve_auto_attention_backend(required_attn_types, kv_quant=kv_quant),
         )
         logger.info_rank0(f"Auto-selected attention backend: {config.attention_backend}")
     _validate_attention_backend_choice(config, override, required_attn_types)

python/freetoken/kernel/aot_models.py

@@ -11,9 +11,11 @@
 (a drifted derivation misses the prebuilt cache by spec name and falls back to
 JIT, which needs nvcc):
 
-- store: ``element_size = num_kv_heads * head_dim * 2`` (bf16 KV row), one per
-  paged-KV attention group (kvcache/mha_pool.py, kvcache/hybrid_swa_pool.py).
-  DSV4 writes its MLA latent via torch scatter and contributes nothing.
+- store: ``element_size = num_kv_heads * head_dim * dtype_bytes``, one per
+  paged-KV attention group (kvcache/mha_pool.py, kvcache/hybrid_swa_pool.py) and
+  per KV width: 2 for the 16-bit cache, 1 for an fp8 one (``--kv-cache-dtype
+  fp8``, kvcache/mha_pool.py). DSV4 writes its MLA latent via torch scatter and
+  contributes nothing.
 - index: ``element_size = hidden_size * 2`` (bf16 embedding row) paired with
   the runtime ``num_splits_for`` rule (layers/embedding.py -> kernel/index.py).
   DSV4 (plain nn.Embedding) and GGUF embeddings (GGUFEmbedding) bypass it.
@@ -33,7 +35,8 @@
 
 from .index import num_splits_for
 
-KV_CACHE_DTYPE_BYTES = 2  # every current model allocates bf16 paged KV
+KV_CACHE_DTYPE_BYTES = 2  # the default paged KV is the 16-bit compute dtype
+FP8_KV_CACHE_DTYPE_BYTES = 1  # --kv-cache-dtype fp8 stores one e4m3 code per element
 EMBED_DTYPE_BYTES = 2  # embedding weights stay bf16 on the indexing() path
 
 
@@ -391,8 +394,9 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int
 )
 
 
-def store_element_sizes(model: AotModel) -> set[int]:
-    return {kv * hd * KV_CACHE_DTYPE_BYTES for kv, hd in model.kv_groups}
+def store_element_sizes(model: AotModel, dtype_bytes: int = KV_CACHE_DTYPE_BYTES) -> set[int]:
+    """Store-kernel row sizes for one model's paged-KV groups at a given bytes/elem."""
+    return {kv * hd * dtype_bytes for kv, hd in model.kv_groups}
 
 
 def index_variants(model: AotModel) -> set[tuple[int, int]]:
@@ -414,9 +418,17 @@ def fast_index_copy_feature_sizes(model: AotModel) -> set[int]:
 
 
 def aggregate_store_element_sizes() -> tuple[int, ...]:
+    """Every store row size the runtime can ask for.
+
+    Both KV widths ship: the 16-bit default and the fp8 (``--kv-cache-dtype fp8``)
+    code buffer, whose rows are exactly half as wide. A missing size is not a
+    correctness bug -- it is a kernel-cache miss that falls back to JIT and fails
+    the ``FREETOKEN_DISABLE_JIT=1`` release gate.
+    """
     sizes: set[int] = set()
     for model in SUPPORTED_MODELS:
-        sizes.update(store_element_sizes(model))
+        for dtype_bytes in (KV_CACHE_DTYPE_BYTES, FP8_KV_CACHE_DTYPE_BYTES):
+            sizes.update(store_element_sizes(model, dtype_bytes))
     return tuple(sorted(sizes))