Add fp8 KV-cache quantization (--kv-dtype)

python/freetoken/attention/fi.py

@@ -57,7 +57,8 @@ class FIMetadata(BaseAttnMetadata):
     page_size:          Literal[1] # currently only support page_size=1
     pos_encoding_mode:  str
     seq_lens_cpu:       torch.Tensor  # on cpu
-    dtype:              torch.dtype
+    dtype:              torch.dtype    # KV-cache storage dtype (may be fp8)
+    q_dtype:            torch.dtype    # query dtype (model activation dtype, bf16)
     wrapper:            BatchPrefillWithPagedKVCacheWrapper | BatchDecodeWithPagedKVCacheWrapper
     initialized:        bool = False
     # fmt: on
@@ -86,6 +87,8 @@ def __init__(self, config: ModelConfig) -> None:
 
         self.config = config
         self.kvcache = get_global_ctx().kv_cache
+        # Query dtype (bf16), distinct from self.kvcache.dtype when KV is fp8.
+        self.q_dtype = get_global_ctx().compute_dtype
         self.device = self.kvcache.device
         # fa2 split-KV prefill needs ``tmp_v <= qo_heads_local * padded_batch_size *
         # cta_tile_q * head_dim * 4`` bytes of scratch, where flashinfer's scheduler
@@ -165,8 +168,7 @@ def _initialize_metadata_once(self, metadata: FIMetadata) -> None:
                 page_size=metadata.page_size,
                 pos_encoding_mode=metadata.pos_encoding_mode,
                 seq_lens=metadata.seq_lens_cpu,
-                data_type=metadata.dtype,
-                q_data_type=metadata.dtype,
+                q_data_type=metadata.q_dtype,
                 kv_data_type=metadata.dtype,
                 non_blocking=True,
             )
@@ -182,7 +184,7 @@ def _initialize_metadata_once(self, metadata: FIMetadata) -> None:
                 page_size=metadata.page_size,
                 pos_encoding_mode=metadata.pos_encoding_mode,
                 seq_lens=metadata.seq_lens_cpu,
-                q_data_type=metadata.dtype,
+                q_data_type=metadata.q_dtype,
                 kv_data_type=metadata.dtype,
                 non_blocking=True,
                 causal=True,
@@ -256,6 +258,7 @@ def prepare_metadata(self, batch: Batch) -> None:
             pos_encoding_mode="NONE",
             seq_lens_cpu=seq_len_cpu,
             dtype=self.kvcache.dtype,
+            q_dtype=self.q_dtype,
             wrapper=self.decode_wrappers if batch.is_decode else self.prefill_wrapper,
         )
 

python/freetoken/core.py

@@ -173,6 +173,9 @@ class Context:
     moe_backend: BaseMoeBackend = field(init=False)
     moe_offload_cache: OffloadMoeCache | None = None
     kv_cache: BaseKVCachePool = field(init=False)
+    # Query/activation dtype (bf16). Distinct from kv_cache.dtype under fp8 KV; the engine
+    # sets it before building the attention backend, which reads it for plan()'s q_data_type.
+    compute_dtype: torch.dtype = field(init=False, default=torch.bfloat16)
     # Per-request recurrent state for GatedDeltaNet layers; set by the engine for
     # hybrid linear-attention models, otherwise None.
     linear_state_pool: LinearStatePool | None = None

python/freetoken/engine/config.py

@@ -58,6 +58,11 @@ class EngineConfig:
     cuda_graph_max_bs: int | None = None
     page_size: int = 1
     memory_ratio: float = 0.9
+    # KV-cache storage dtype (--kv-dtype). None -> use the model/activation ``dtype`` (bf16).
+    # fp8 (e4m3/e5m2) halves KV bytes/token, ~doubling the token budget that fits in memory,
+    # at some attention precision cost. Read via ``resolved_kv_dtype``; only the paged MHA/SWA
+    # KV slabs switch dtype -- queries stay in ``dtype`` and the fp8 kernels dequantize on read.
+    kv_dtype: torch.dtype | None = None
     # 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
     # ceil(ratio * max_running_req) extra slots.
@@ -91,6 +96,12 @@ def model_config(self) -> ModelConfig:
         parse_config = _load_attr(spec.module, spec.parse_config)
         return parse_config(self.hf_config)
 
+    @property
+    def resolved_kv_dtype(self) -> torch.dtype:
+        """Storage dtype for the paged KV cache: the explicit ``--kv-dtype`` if set, else the
+        model dtype. Queries always stay in ``dtype``; this only sizes/quantizes the KV slabs."""
+        return self.kv_dtype if self.kv_dtype is not None else self.dtype
+
     @property
     def max_seq_len(self) -> int:
         if self.max_seq_len_override is not None:

python/freetoken/engine/engine.py

@@ -302,11 +302,24 @@ def __init__(self, config: EngineConfig):
         self.stream = torch.cuda.Stream()
         torch.cuda.set_stream(self.stream)
         self.dtype = config.dtype
+        # Paged KV storage dtype (may be fp8 via --kv-dtype); queries/activations stay self.dtype.
+        self.kv_dtype = config.resolved_kv_dtype
         self.config = config  # retained for runtime cache rebuild (rebuild_runtime_cache)
         # KV pool family fixed at construction from the model config: its classmethods own the
         # page-token geometry and cost arithmetic the engine needs BEFORE the pool exists
         # (num_pages sizing, --moe-cache-auto); the instance owns rebuild/validation after.
         self._pool_cls = resolve_pool_class(config.model_config)
+        # fp8 KV cache is implemented for the full-attention (MHA) pool only: the fp8 store-quantize
+        # lives in MHAKVCache.store_kv, and other pools (SWA/MLA/DSA/...) would byte-copy bf16 into
+        # fp8 slabs and corrupt KV. Reject the unsupported combination up front with a clear error.
+        if self.kv_dtype != self.dtype:
+            from freetoken.kvcache.mha_pool import MHAKVCache
+
+            if self._pool_cls is not MHAKVCache:
+                raise NotImplementedError(
+                    f"--kv-dtype {self.kv_dtype} (fp8 KV cache) is only supported for full-attention "
+                    f"(MHA) models; this model uses {self._pool_cls.__name__}."
+                )
         self.ctx = Context(config.page_size)
         set_global_ctx(self.ctx)
 
@@ -345,7 +358,7 @@ def __init__(self, config: EngineConfig):
         self.num_pages = self._pool_cls.solve_num_pages(config, available_memory)
         num_tokens = self.num_pages * config.page_size
         self.ctx.kv_cache = self.kv_cache = create_kv_pool(
-            config, self.num_pages, device=self.device, dtype=self.dtype
+            config, self.num_pages, device=self.device, dtype=self.kv_dtype
         )
 
         # ======================= Linear (GatedDeltaNet) state initialization ========================
@@ -379,9 +392,23 @@ def __init__(self, config: EngineConfig):
         self.kv_cache.attach_page_table(self.page_table)
 
         # ======================= Attention & MoE backend initialization ========================
+        # Query/activation dtype for attention (the fp8-KV path needs it distinct from the KV
+        # storage dtype); backends read it off the global context in their constructors.
+        self.ctx.compute_dtype = self.dtype
         self.ctx.attn_backend = self.attn_backend = create_attention_backend(
             config.attention_backend, config.model_config
         )
+        # The fp8 dequant-on-read (q_data_type bf16 / kv_data_type fp8 in plan()) lives in the
+        # FlashInfer backend; other backends would read the fp8 slabs as bf16. Require fi for fp8 KV.
+        if self.kv_dtype != self.dtype:
+            from freetoken.attention.fi import FlashInferBackend
+
+            if not isinstance(self.attn_backend, FlashInferBackend):
+                raise NotImplementedError(
+                    f"--kv-dtype {self.kv_dtype} (fp8 KV cache) requires the FlashInfer (fi) "
+                    f"attention backend; resolved backend is {type(self.attn_backend).__name__} "
+                    f"(pass --attention-backend fi)."
+                )
         if config.model_config.is_moe:
             self.ctx.moe_backend = self.moe_backend = create_moe_backend(config.moe_backend)
 

python/freetoken/kvcache/base.py

@@ -22,11 +22,13 @@ def spec_kv_bytes_per_token(spec, config) -> int:
     per-spec arithmetic -- pool families compose it over THEIR OWN groups; no family
     branching here. (2 bytes/elem == the torch.bfloat16 dsa_pool.DSAKVCache._alloc
     hardcodes; keep the two in lockstep if the slab dtype ever changes.)"""
+    # KV slabs use the (possibly fp8) KV dtype; duck-typed test configs lack the property.
+    kv_dtype = getattr(config, "resolved_kv_dtype", None) or config.dtype
     per_token = (
         (1 if spec.mla else 2)  # MLA latent groups store one slab (V aliases K)
         * spec.head_dim
         * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True)
-        * config.dtype.itemsize
+        * kv_dtype.itemsize
         * spec.num_layers
     )
     return per_token + spec.index_head_dim * spec.num_index_layers * 2

python/freetoken/kvcache/mha_pool.py

@@ -127,6 +127,13 @@ def store_kv(
         from freetoken.kernel import store_cache
 
         dense = self._dense(layer_id)
+        # fp8 KV cache: quantize the incoming bf16 k/v to the pool dtype here (scale 1.0 --
+        # post-RoPE k/v magnitudes sit well inside e4m3/e5m2 range), so store_cache stays a
+        # same-width byte copy and the fp8 attention kernels dequantize on read.
+        buf_dtype = self._kv_buffer.dtype
+        if k.dtype != buf_dtype:
+            k = k.to(buf_dtype)
+            v = v.to(buf_dtype)
         store_cache(
             k_cache=self._k_buffer[dense].view(self._storage_shape),
             v_cache=self._v_buffer[dense].view(self._storage_shape),

python/freetoken/server/args.py

@@ -213,6 +213,17 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         help="Data type for model weights and activations. 'auto' will use FP16 for FP32/FP16 models and BF16 for BF16 models.",
     )
 
+    parser.add_argument(
+        "--kv-dtype",
+        type=str,
+        default="auto",
+        choices=["auto", "bfloat16", "fp8_e4m3", "fp8_e5m2"],
+        help="KV-cache storage dtype. 'auto' matches the model dtype (bf16). fp8_e4m3/fp8_e5m2 "
+        "halve KV bytes/token (~2x the token budget that fits) with negligible quality loss; "
+        "queries stay bf16 and the FlashInfer fp8 kernels dequantize on read. fp8 currently "
+        "supports full-attention (MHA) models on the FlashInfer backend only.",
+    )
+
     parser.add_argument(
         "--tensor-parallel-size",
         "--tp-size",
@@ -680,6 +691,16 @@ def _infer_reasoning_parser(model_path: str) -> str | None:
         "float32": torch.float32,
     }
     kwargs["dtype"] = DTYPE_MAP[dtype_str] if isinstance(dtype_str, str) else dtype_str
+
+    # KV-cache dtype: "auto" -> None (falls back to the model dtype in resolved_kv_dtype).
+    KV_DTYPE_MAP = {
+        "auto": None,
+        "bfloat16": torch.bfloat16,
+        "fp8_e4m3": torch.float8_e4m3fn,
+        "fp8_e5m2": torch.float8_e5m2,
+    }
+    kwargs["kv_dtype"] = KV_DTYPE_MAP[kwargs["kv_dtype"]]
+
     kwargs["tp_info"] = DistributedInfo(0, kwargs["tensor_parallel_size"])
     del kwargs["tensor_parallel_size"]