feat(kvcache): 8-bit KV cache (q8_0 / fp8_e4m3) behind --kv-cache-dtype

python/freetoken/attention/triton.py

@@ -144,6 +144,7 @@ def forward(
             extend_paged_attention,
             paged_attention,
         )
+        from freetoken.kvcache.quant import BLOCK as QBLOCK
 
         metadata = batch.attn_metadata
         assert isinstance(metadata, TritonMetadata)
@@ -155,6 +156,12 @@ 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 8-bit pool hands its per-block scales alongside the slabs; an unquantized one
+        # has none, and the kernels compile their bf16 path unchanged.
+        k_scale = v_scale = None
+        if getattr(self.kvcache, "quant", None) is not None and self.kvcache.quant.enabled:
+            k_scale = self.kvcache.k_scale(layer_id).view(-1, kv_heads, head_dim // QBLOCK)
+            v_scale = self.kvcache.v_scale(layer_id).view(-1, kv_heads, head_dim // QBLOCK)
 
         spec = attn_spec or AttentionSpec()
         indices = metadata.indices
@@ -181,6 +188,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 +210,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 +224,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

@@ -80,6 +80,16 @@ class EngineConfig:
     # KV capacity in tokens; resolved into num_page_override by _adjust_config once page_size
     # is final. Mutually exclusive with num_page_override.
     num_token_override: int | None = None
+    # KV element storage (--kv-cache-dtype): "auto" keeps the compute dtype, "q8_0" and
+    # "fp8_e4m3" store 8 bits plus a per-block scale. Resolved through
+    # freetoken.kvcache.quant.resolve_kv_quant by the pools and the cost model.
+    kv_cache_dtype: str = "auto"
+
+    @cached_property
+    def kv_quant(self):
+        from freetoken.kvcache.quant import resolve_kv_quant
+
+        return resolve_kv_quant(self.kv_cache_dtype)
 
     @cached_property
     def hf_config(self):

python/freetoken/engine/engine.py

@@ -154,6 +154,44 @@ def _resolve_auto_attention_backend(
     )
 
 
+def _validate_kv_cache_dtype(config, model_config) -> None:
+    """Gate --kv-cache-dtype against what the quantized path actually implements.
+
+    8-bit KV storage lives in the triton attention kernels and the MHA/hybrid-SWA pools.
+    Every other backend reads the KV slabs through its own kernels (flashinfer's
+    ``kv_data_type``, trtllm's fp8 path) which this has not been wired into, and the
+    MLA/DSA/DSV4/BSA pools have their own slab layouts. Reject those combinations here,
+    at config time, rather than letting a wrong-dtype tensor reach a kernel.
+    """
+    quant = getattr(config, "kv_quant", None)
+    if quant is None or not quant.enabled:
+        return
+
+    from freetoken.kvcache.quant import BLOCK
+
+    backends = [p.strip() for p in config.attention_backend.split(",")]
+    if any(b != "triton" for b in backends):
+        raise ValueError(
+            f"--kv-cache-dtype {quant.name} needs the triton attention backend, but the "
+            f"resolved backend is {config.attention_backend!r}. Pass "
+            "--attention-backend triton, or drop --kv-cache-dtype."
+        )
+
+    specs = [s for s in model_config.kv_cache_group_specs() if s.num_layers > 0]
+    if any(s.mla or s.index_head_dim > 0 for s in specs):
+        raise ValueError(
+            f"--kv-cache-dtype {quant.name} does not support MLA/DSA latent KV pools "
+            "(their slabs alias K and V and carry an index tier); use --kv-cache-dtype auto."
+        )
+    bad = [s for s in specs if s.head_dim % BLOCK]
+    if bad:
+        names = ", ".join(f"{s.name} (head_dim {s.head_dim})" for s in bad)
+        raise ValueError(
+            f"--kv-cache-dtype {quant.name} needs every head_dim to be a multiple of "
+            f"{BLOCK}, the quantization block; this model has {names}."
+        )
+
+
 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 +1228,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