feat(kvcache): add nvfp4 kv quantization

benchmarks/README.md

@@ -29,3 +29,12 @@ python benchmarks/bench_offload_cache_copy.py
 
 For host RAM vs PCIe bandwidth and the offload/hybrid backend pick, use `ft bench bw`
 instead — it writes the JSON profile the engine reads.
+
+**`bench_kv_quant.py`** compares BF16, FP8 and NVFP4 KV storage bytes, one-step
+scatter latency and paged decode latency on synthetic inputs. No checkpoint is
+required. Keep the GPU idle and use identical arguments for A/B comparisons;
+this does not measure model quality or end-to-end serving throughput.
+
+```bash
+PYTHONPATH=python:. uv run python benchmarks/bench_kv_quant.py --lengths 1024,8192,32768
+```

benchmarks/bench_kv_quant.py

@@ -0,0 +1,77 @@
+"""Paged KV storage/scatter/decode microbenchmark, independent of model weights.
+
+Run with PYTHONPATH=python:. uv run python benchmarks/bench_kv_quant.py.
+Compare identical arguments on the baseline and candidate; this does not measure
+end-to-end model quality, TTFT, or serving throughput.
+"""
+
+import argparse
+import json
+
+import torch
+import triton.testing
+
+from freetoken.distributed import set_tp_info
+from freetoken.kernel.triton.attention import decode_paged_attention
+from freetoken.kvcache.mha_pool import MHAKVCache
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--formats", default="none,fp8,nvfp4")
+    parser.add_argument("--lengths", default="1024,8192,32768")
+    parser.add_argument("--batch", type=int, default=1)
+    parser.add_argument("--heads", type=int, default=4)
+    parser.add_argument("--group", type=int, default=4)
+    parser.add_argument("--dim", type=int, default=128)
+    args = parser.parse_args()
+    set_tp_info(rank=0, size=1)
+    torch.manual_seed(42)
+    batch, heads, dim = args.batch, args.heads, args.dim
+    qheads = heads * args.group
+    device = torch.device("cuda")
+    results = []
+    for length in map(int, args.lengths.split(",")):
+        slots = batch * length
+        k, v = [torch.randn(slots, heads * dim, device=device, dtype=torch.bfloat16)
+                for _ in range(2)]
+        loc = torch.arange(slots, device=device, dtype=torch.int32)
+        q = torch.randn(batch, qheads, dim, device=device, dtype=torch.bfloat16)
+        indptr = torch.arange(batch + 1, device=device, dtype=torch.int32) * length
+        pos = torch.full((batch,), length - 1, device=device, dtype=torch.int32)
+        scratch = torch.empty(batch, qheads, 8, dim, device=device)
+        lse = torch.empty(batch, qheads, 8, device=device)
+        splits = torch.full((batch,), 8, device=device, dtype=torch.int32)
+        out = torch.empty_like(q)
+        for quant in args.formats.split(","):
+            pool = MHAKVCache(heads, 1, dim, slots, 1, q.dtype, device, kv_quant=quant)
+            pool.store_kv(k, v, loc, 0)
+            extra = {}
+            if quant == "nvfp4":
+                extra = dict(kv_quant=quant, k_block_scale=pool.k_block_scale(0),
+                             v_block_scale=pool.v_block_scale(0))
+            kc, vc = [getattr(pool, name)(0).flatten(0, 1) for name in ("k_cache", "v_cache")]
+
+            def decode():
+                return decode_paged_attention(q, kc, vc, indptr, loc, pos,
+                    scratch, lse, splits, 8, dim ** -.5, out=out,
+                    k_scale=pool.k_scale(0), v_scale=pool.v_scale(0), **extra)
+
+            decode()
+            decode_ms = triton.testing.do_bench(decode, warmup=100, rep=300)
+            store_ms = triton.testing.do_bench(
+                lambda: pool.store_kv(k[-batch:], v[-batch:], loc[-batch:], 0),
+                warmup=100, rep=300)
+            record = dict(format=quant, length=length, batch=batch,
+                          kv_bytes=pool.unit_bytes()[0] * slots,
+                          decode_ms=decode_ms, store_ms=store_ms)
+            results.append(record)
+            print(json.dumps(record), flush=True)
+            del pool, kc, vc
+        del k, v
+    print(json.dumps(dict(gpu=torch.cuda.get_device_name(), torch=torch.__version__,
+                         heads=heads, group=args.group, dim=dim, results=results)))
+
+
+if __name__ == "__main__":
+    main()

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`, `fp8`, or `nvfp4` (see [NVFP4 KV cache](#nvfp4-kv-cache)): FP8 stores 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.
@@ -173,3 +197,28 @@ profile that `ft serve --moe-backend auto` and `--moe-hybrid-max-fetch -1` then
 - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe
   by that factor.
 
+
+### NVFP4 KV cache
+
+`ft serve --model <checkpoint> --kv-cache-dtype nvfp4 --attention-backend triton`
+opts into packed E2M1 KV storage. The initial implementation supports plain paged
+FULL attention (MHA/GQA), hybrid-SWA, and the full-attention portion of hybrid-linear
+models, and QSA. Head dimensions must be divisible by 16. MLA/DSA, DSV4, and BSA
+pools are rejected at startup. `auto` selects Triton or QSA sparse attention for
+supported models.
+
+Each K or V row stores `head_dim / 2` packed bytes, `head_dim / 16` E4M3 block-scale
+bytes, and one FP32 row scale. At head_dim 128 this is 76 bytes, versus 256 for
+BF16 and 132 for the existing FP8 format. Pool management, recurrent states,
+attention workspace and model weights consume additional memory.
+
+The second-level scale is dynamic per token/head, so appending a token never
+rescales an existing prefix. This is a FreeToken KV layout, not an external
+NVFP4 checkpoint or attention-library ABI. K/V are restored inside attention;
+Q and attention arithmetic retain their compute precision. The MoE weight option
+`--nvfp4-backend` is independent. Prefill uses fresh compute-dtype K/V while
+cached prefixes are restored, as in the FP8 path.
+
+NVFP4 is opt-in: assess quality on your checkpoint and workload before using it
+for long-context inference. Capacity savings do not guarantee faster decode;
+packing, reconstruction, and the selected attention backend affect throughput.

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,12 @@ 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
+    supports_nvfp4_kv: bool = False
 
 
 SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend")
@@ -84,6 +90,8 @@ def create_fa_backend(config: ModelConfig):
     BackendInfo(
         supported_types=frozenset({AttnType.FULL, AttnType.SWA}),
         consumes_attn_spec=True,
+        supports_fp8_kv=True,
+        supports_nvfp4_kv=True,
     ),
 )
 def create_triton_backend(config: ModelConfig):
@@ -104,7 +112,10 @@ def create_dsv4_sparse_backend(config: ModelConfig):
 
 @SUPPORTED_ATTENTION_BACKENDS.register(
     "dsa",
-    BackendInfo(supported_types=frozenset({AttnType.MLA, AttnType.DSA})),
+    BackendInfo(
+        supported_types=frozenset({AttnType.MLA, AttnType.DSA}),
+        supports_fp8_kv=True,
+    ),
 )
 def create_dsa_backend(config: ModelConfig):
     # MLA with a grouped index (index_ratio > 1) is the kpool indexer layout.
@@ -137,6 +148,10 @@ 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,
+        supports_nvfp4_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/dsa.py

@@ -201,6 +201,7 @@ def _attend(
         return glm_dsa_sparse_attn(
             q_cat, self.kvcache.latent_rows(layer_id), sel, self.sm_scale,
             counts=cnt, d_v=self.kv_lora_rank,
+            pool_scale=self.kvcache.latent_scale(layer_id),
         )
 
     def mla_forward(

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)
+        # K/V scale tensors are independent of the BF16 index tier, so selection is
+        # quantization-agnostic; only sparse K/V attention reconstructs the codes.
         return qsa_sparse_paged_attention(
             q,
             self.kvcache.k_cache(layer_id),
@@ -290,6 +302,11 @@ 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),
+            kv_quant=self.kvcache.kv_quant,
+            k_block_scale=self.kvcache.k_block_scale(layer_id),
+            v_block_scale=self.kvcache.v_block_scale(layer_id),
         )
 
     def _plan_index_writes(self, md: QSASparseMetadata, batch: Batch) -> None:

python/freetoken/attention/triton.py

@@ -151,10 +151,19 @@ def forward(
 
         k_raw = self.kvcache.k_cache(layer_id)
         v_raw = self.kvcache.v_cache(layer_id)
-        kv_heads, head_dim = k_raw.shape[-2], k_raw.shape[-1]
-        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)
+        kv_heads, stored_dim = k_raw.shape[-2], k_raw.shape[-1]
+        head_dim = q.shape[-1]
+        kv_quant = getattr(self.kvcache, "kv_quant", "none")
+        assert stored_dim == (head_dim // 2 if kv_quant == "nvfp4" else head_dim)
+        k_cache = k_raw.view(-1, kv_heads, stored_dim)
+        v_cache = v_raw.view(-1, kv_heads, stored_dim)
+        k_block_scale = self.kvcache.k_block_scale(layer_id) if kv_quant == "nvfp4" else None
+        v_block_scale = self.kvcache.v_block_scale(layer_id) if kv_quant == "nvfp4" else None
+        # 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 +190,11 @@ def forward(
                 sm_scale=scale,
                 sliding_window=spec.sliding_window,
                 sinks=spec.sinks,
+                k_scale=k_scale,
+                v_scale=v_scale,
+                kv_quant=kv_quant,
+                k_block_scale=k_block_scale,
+                v_block_scale=v_block_scale,
             )
         if (
             (not metadata.is_decode)
@@ -201,6 +215,11 @@ 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,
+                kv_quant=kv_quant,
+                k_block_scale=k_block_scale,
+                v_block_scale=v_block_scale,
             )
         return paged_attention(
             q=q,
@@ -213,6 +232,11 @@ def forward(
             sm_scale=scale,
             sliding_window=spec.sliding_window,
             sinks=spec.sinks,
+            k_scale=k_scale,
+            v_scale=v_scale,
+            kv_quant=kv_quant,
+            k_block_scale=k_block_scale,
+            v_block_scale=v_block_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