feat(models): support deepseek v4.1 with native fp8-fp4 kv storage

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`, `nvfp4`, or `fp8-fp4`. Generic [FP8](#fp8-kv-cache) and [NVFP4](#nvfp4-kv-cache) use their own scale layouts; [DeepSeek-V4.1](deepseek-v41.md#native-fp8-fp4-kv-storage) uses `fp8-fp4` to preserve its native FP8 window, FP4 compressed KV and MXFP4 index keys. |
 | `--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-strategies) for what each strategy does.
@@ -101,12 +125,17 @@ See [models.md](models.md#moe-strategies) for what each strategy does.
 
 ### Image input
 
-Experimental. Needs a checkpoint whose family registers a vision encoder ([models.md](models.md#image-input) lists them and how each one
-maps the flags below); a request carrying images is rejected otherwise. Images are accepted on all three protocols (OpenAI `image_url`,
-Anthropic `image` blocks, Responses `input_image`) as an http(s) URL or base64. Images inside a tool
-result (an Anthropic `tool_result` block from Claude Code's Read, a Responses `function_call_output`
-from Codex's view_image) are moved to the user turn that follows the tool message, as vLLM does,
-because chat templates render tool messages as plain text.
+Experimental. Needs a checkpoint whose family registers a vision encoder;
+[models.md](models.md#image-input) lists the families and their image options.
+A request carrying images is rejected otherwise. Images are accepted on all three protocols (OpenAI
+`image_url`, Anthropic `image` blocks, Responses `input_image`) as an http(s) URL
+or base64. Tool-result image support depends on the model's chat encoding:
+DeepSeek-V4.1 accepts ordered text and images inside Anthropic `tool_result`
+blocks and Responses `function_call_output` without moving them. For Qwen VL and
+other chat templates that render tool messages as plain text, images are moved
+to a user turn after the tool message, as vLLM does. This supports tool images
+from Claude Code's Read and Codex's view_image while keeping the tool's text
+under its original call ID.
 `GET /v1/stats` reports what the server accepts as `model.input_modalities` (`["text"]` or `["text", "image"]`),
 so a client can gate its attachment controls without reading the checkpoint config.
 
@@ -115,12 +144,19 @@ so a client can gate its attachment controls without reading the checkpoint conf
 | `--text-model-only` | off | Serve a multimodal checkpoint text-only: no encoder tower is built (its VRAM goes to the KV/expert pools) and every multimodal input is rejected. Same as `--mm-disable` with every encoder kind |
 | `--mm-disable` | none | Encoder towers to leave unbuilt (`vision`, `audio`); every input they would serve is rejected |
 | `--mm-encoder-weights` | host | Where the encoder tower's block weights live. `host` streams them from pinned host banks two blocks at a time behind the compute, so the GPU holds two blocks instead of the whole tower; small images pay the copy time, large ones hide it behind the compute. `gpu` keeps them resident. An encoder without a block stack stays resident either way |
-| `--image-min-tokens`, `--image-max-tokens` | processor defaults | Per-image token budget: the image processor resizes every image to take between these many tokens, converted to the family's own units by its processor. A family with fixed budgets honors the maximum only and refuses one below its smallest budget at start-up |
+| `--image-min-tokens`, `--image-max-tokens` | processor defaults | Per-image token budget, converted to the family's own units by its processor. A family with fixed budgets honors the maximum only and refuses one below its smallest budget at start-up. DeepSeek-V4.1 supports the maximum and rejects the minimum flag, as described below |
 | `--mm-processor-kwargs` | none | JSON object of extra keyword arguments for the checkpoint's image processor call, for knobs the token budget does not cover; applied after the budget, so an explicit key wins |
 | `--mm-embed-cache-device` | cpu | Where encoded image embeddings live between prefill chunks. `cpu` keeps them out of the VRAM budget; `cuda` skips the copy back |
 | `--allowed-media-domains` | any | Comma-separated hostname allowlist for image URLs; requests for other domains are rejected with a 400. Empty allows any domain |
 | `--allowed-local-media-path` | off | Directory `file://` image refs may be read from; unset rejects local files |
 
+DeepSeek-V4.1's `--image-max-tokens` limits the entire image span, including start,
+row-newline and end tokens. It rejects `--image-min-tokens`; use
+`--mm-processor-kwargs '{"vision_min_pixels": 295936}'` to set the minimum pixel
+area before the maximum-token resize. Its supported processor options are
+`vision_min_pixels`, `vision_max_n_token` and `vision_max_wh_ratio`;
+`vision_max_n_token` takes precedence over `--image-max-tokens` when both are set.
+
 ## ft shell
 
 ```bash
@@ -202,3 +238,33 @@ profile that `ft serve --moe-strategy 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, QSA, and MLA/DSA (including GLM-5.3-Flash). Head dimensions must be
+divisible by 16. DSV4 and BSA pools are rejected at startup. `auto` selects
+Triton, QSA sparse, or DSA attention for supported models. For MLA/DSA use
+`--attention-backend auto` or `--attention-backend dsa`. Only the latent slab is
+quantized; indexer keys, kpool tails/gates, and recurrent states retain their
+existing precision. A 512-element latent row occupies 292 bytes instead of 1024
+bytes in BF16, excluding those other tiers.
+
+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. Paged MHA prefill uses fresh compute-dtype K/V
+while cached prefixes are restored, as in the FP8 path. MLA/DSA stores fresh
+latent rows first and reads the quantized cache in both prefill and decode.
+
+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.