feat(kvcache): add opt-in fp8 kv cache for r9700

docs/cli.md

@@ -74,8 +74,30 @@ ft serve --model ... --gpu GPU-9e8d7c6b  # the same card by UUID (a unique prefi
 | `--embed-weights` | gpu | Where the input token embedding table lives. `host` keeps it in pinned host RAM and each lookup reads its rows over PCIe, so its VRAM goes to the expert cache and KV. Refused when the table is tied to the LM head |
 | `--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` | auto | `auto`/`bf16`: compute-dtype KV storage; `fp8`: E4M3 KV storage with FP32 row scales. FP8 requires BF16 compute, TP=1, MHA/GQA full-attention KV and Triton |
 | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU |
 
+FP8 KV is independent of weight quantization. For Qwen3 and Qwen3.5/3.6, use
+`--dtype bfloat16 --kv-cache-dtype fp8 --attention-backend triton`.
+Each cached K/V row stores E4M3 codes in a byte buffer and one FP32 scale per
+token, layer and KV head. At head dimension 128 this saves 48.44% of KV memory
+relative to BF16 for persistent storage. Single-request prefix prefill with at
+least 128 new tokens reconstructs history once per layer into a shared BF16
+workspace, then uses the mainline BF16 attention tiles. The workspace has one
+layer's K/V shape, including the dummy page; it is reused across all layers.
+For Qwen3-30B-A3B at 33059 physical token rows this adds 64.57 MiB, leaving a
+46.35% saving for persistent KV plus workspace. Startup sizing, automatic MoE
+sizing, and live pool resizing include this workspace and the FP32 scales.
+Hybrid linear/full-attention models quantize only their full-attention KV layers;
+linear recurrent state retains its existing precision and storage. SWA, MLA,
+sparse KV pools and speculative decoding are rejected. The Triton decode split
+policy remains adaptive, with the mainline maximum of 64.
+On gfx1201, single-request BF16 decode uses native packed FP8 reads for D=128
+with GQA4/8 and for D=256 with 16 query heads and 2 KV heads. D=128 GQA2 retains
+its tuned compatibility path. Other shapes and devices use the generic path.
+The inline FP8 prefill fallback uses at most 64x32 tiles on ROCm; prefill and
+the stage-2 decode reduction are unchanged by the native decode path.
+
 ### MoE offload
 
 See [models.md](models.md#moe-strategies) for what each strategy does, [performance/prefill.md](performance/prefill.md) for prefill speed, and [performance/decode.md](performance/decode.md) for decode speed.

python/freetoken/attention/__init__.py

@@ -35,6 +35,7 @@ class BackendInfo:
     consumes_attn_spec: bool = False
     # How the backend captures a graph of `rows_per_req` query rows per request: "device" keeps all metadata on device, "host_plan" plans on the host before each replay, "none" cannot.
     uniform_extend_graph: Literal["none", "host_plan", "device"] = "none"
+    supports_fp8_kv: bool = False
 
 
 SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend")
@@ -89,6 +90,7 @@ def create_fa_backend(config: ModelConfig):
         supported_types=frozenset({AttnType.FULL, AttnType.SWA}),
         consumes_attn_spec=True,
         uniform_extend_graph="device",
+        supports_fp8_kv=True,
     ),
 )
 def create_triton_backend(config: ModelConfig):

python/freetoken/attention/base.py

@@ -55,6 +55,12 @@ def get_last_indices(self, bs: int) -> torch.Tensor: ...
 
 
 class BaseAttnBackend(ABC):
+    def init_workspace(self) -> None:
+        """Allocate persistent attention scratch before capture or serving."""
+
+    def reset_workspace(self) -> None:
+        """Release scratch at the synchronized runtime-cache rebuild boundary."""
+
     @abstractmethod
     def forward(
         self,
@@ -93,6 +99,14 @@ def reset_capture(self) -> None:
 
 
 class HybridBackend(BaseAttnBackend):
+    def init_workspace(self) -> None:
+        self.prefill_backend.init_workspace()
+        self.decode_backend.init_workspace()
+
+    def reset_workspace(self) -> None:
+        self.prefill_backend.reset_workspace()
+        self.decode_backend.reset_workspace()
+
     def __init__(
         self,
         prefill_backend: BaseAttnBackend,

python/freetoken/attention/triton.py

@@ -6,10 +6,13 @@
 import torch
 from freetoken.core import Batch, get_global_ctx
 from freetoken.kernel.triton.attention import decode_kv_split_cap
+from freetoken.utils import init_logger
 
 from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata
 from .utils import BaseCaptureData
 
+logger = init_logger(__name__)
+
 if TYPE_CHECKING:
     from freetoken.models import ModelConfig
 
@@ -61,6 +64,7 @@ class TritonMetadata(BaseAttnMetadata):
     attn_logits: torch.Tensor | None = None
     attn_lse: torch.Tensor | None = None
     swa_indices: torch.Tensor | None = None
+    has_prefix: bool = True
 
     def get_last_indices(self, bs: int) -> torch.Tensor:
         return self.cu_seqlens_q_gpu[1 : 1 + bs] - 1
@@ -83,6 +87,16 @@ def __init__(self, config: ModelConfig):
             default=int(getattr(config, "head_dim", 1)),
         )
         self.kv_heads = sorted({group.num_kv_heads for group in kv_groups}) or [int(getattr(config, "num_kv_heads", 1))]
+        self.prefill_workspace: torch.Tensor | None = None
+
+    def init_workspace(self) -> None:
+        shape = self.kvcache.prefill_workspace_shape
+        if shape is not None:
+            self.prefill_workspace = torch.empty((2, *shape), device=self.device, dtype=torch.bfloat16)
+            logger.info(f"FP8 prefill workspace: {self.prefill_workspace.numel() * 2} bytes, shared across layers")
+
+    def reset_workspace(self) -> None:
+        self.prefill_workspace = None
 
     def _split_cap(self, bs: int) -> int:
         """The most decode KV splits any layer group takes at this batch size."""
@@ -142,6 +156,8 @@ 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)
+        k_scale = self.kvcache.k_scale(layer_id)
+        v_scale = self.kvcache.v_scale(layer_id)
 
         spec = attn_spec or AttentionSpec()
         block_ends = batch.mm_block_ends if spec.bidirectional_mm_blocks else None
@@ -166,12 +182,29 @@ 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)
             and q.dtype in (torch.float16, torch.bfloat16)
             and (q.shape[-1] <= 256 or metadata.max_q_len >= self.prefill_tile_min_q)
         ):
+            k_extend = k.view(q.shape[0], kv_heads, head_dim)
+            v_extend = v.view(q.shape[0], kv_heads, head_dim)
+            if k_scale is not None and not metadata.has_prefix:
+                # No cached rows are read on the first chunk; retain main's BF16 kernel.
+                k_cache, v_cache = k_extend, v_extend
+                k_scale = v_scale = None
+            elif k_scale is not None and metadata.max_q_len >= 128 and metadata.prefix_lens.numel() == 1:
+                from freetoken.kernel.triton.kv_quant import reconstruct_kv_prefix
+
+                assert self.prefill_workspace is not None, "FP8 prefill workspace must be initialized before serving"
+                wk, wv = self.prefill_workspace.unbind(0)
+                reconstruct_kv_prefix(k_cache, v_cache, k_scale, v_scale,
+                                      metadata.indptr, indices, metadata.prefix_lens, wk, wv)
+                k_cache, v_cache = wk, wv
+                k_scale = v_scale = None
             return extend_paged_attention(
                 q=q,
                 k_cache=k_cache,
@@ -184,9 +217,11 @@ def forward(
                 sm_scale=scale,
                 sliding_window=spec.sliding_window,
                 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_extend=k_extend,
+                v_extend=v_extend,
                 block_ends=block_ends,
+                k_scale=k_scale,
+                v_scale=v_scale,
             )
         if block_ends is not None:
             raise NotImplementedError("bidirectional multimodal blocks need the extend kernel path")
@@ -201,6 +236,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:
@@ -253,6 +290,7 @@ def prepare_metadata(self, batch: Batch) -> None:
             prefix_lens=prefix_lens,
             max_q_len=max(seqlens_q),
             swa_indices=swa_indices,
+            has_prefix=any(cached_lens),
         )
 
     def init_capture_graph(self, max_seq_len: int, bs_list: List[int], rows_per_req: int = 1) -> None:

python/freetoken/engine/config.py

@@ -27,6 +27,7 @@ class EngineConfig:
     hf_overrides: Mapping[str, Any] = field(default_factory=dict)
     max_running_req: int = 4
     attention_backend: str = "auto"
+    kv_quant: str = "none"
     moe_strategy: str = "auto"
     # old name of moe_strategy; __post_init__ folds it in
     moe_backend: str | None = field(default=None, repr=False)

python/freetoken/engine/engine.py

@@ -167,7 +167,15 @@ def _backend_requirements_met(name: str) -> bool:
     return True
 
 
-def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str:
+def _resolve_kv_quant(value: str | None) -> str:
+    key = (value or "none").lower()
+    aliases = {"auto": "none", "bf16": "none", "none": "none", "fp8": "fp8"}
+    if key not in aliases:
+        raise ValueError(f"unknown --kv-cache-dtype {value!r}; expected auto, bf16 or fp8")
+    return aliases[key]
+
+
+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:
@@ -197,6 +205,10 @@ def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str:
             continue
         if not _backend_parts_serve(name, required):
             continue
+        if kv_quant == "fp8" and not all(
+            attention_backend_info(p).supports_fp8_kv for p in name.split(",")
+        ):
+            continue
         if not _backend_requirements_met(name):
             continue
         return name
@@ -255,6 +267,8 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att
     backend_parts = [p.strip() for p in config.attention_backend.split(",")]
     for part in backend_parts:
         info = attention_backend_info(part)
+        if getattr(config, "kv_quant", "none") == "fp8" and not info.supports_fp8_kv:
+            raise ValueError("--kv-cache-dtype fp8 requires --attention-backend triton")
         missing = required - info.supported_types
         if missing:
             valid = [
@@ -598,6 +612,7 @@ def __init__(self, config: EngineConfig):
             config.attention_backend, config.model_config
         )
         self.linear_backend = LinearAttnBackend(self.device, self.linear_state_pool) if self.linear_state_pool is not None else None
+        self.attn_backend.init_workspace()
 
         # ======================= Sampler initialization ========================
         self.sampler = Sampler(self.device, config.model_config.vocab_size)
@@ -751,7 +766,7 @@ def _resolve_auto_moe_cache_size(self, config: EngineConfig, banks, method=None,
         """
         from freetoken.engine.cache_budget import expert_bytes_per_slot, resolve_moe_cache_auto
 
-        cache_per_page, fixed_cache_size, page_tokens, min_reserve = self._pool_cls.kv_cost(config)
+        cache_per_page, fixed_cache_size, page_tokens, min_reserve = self._pool_cls.runtime_cost(config)
         fixed_cache_size += state_pool_bytes(config)  # sibling GDN state pool, engine-summed
         spec_fixed, draft_page = self._spec_budget(config)
         fixed_cache_size += spec_fixed
@@ -1289,6 +1304,7 @@ def rebuild_runtime_cache(
         if self.spec is not None:
             # the draft KV is keyed by the target's slots, so it can stay only when the target KV pool does
             self.spec.detach(keep_kv=num_pages is None and num_swa_pages is None)
+        self.attn_backend.reset_workspace()
         # 2. Resize caches in place (each frees its old GPU tensors before allocating).
         # Pin the new window first (validated above) so any KV-pool rebuild below sizes the window
         # to it (_dsv4_pool_sizes / _swa_paged_num_tokens read config.swa_num_pages_override).
@@ -1314,6 +1330,7 @@ def rebuild_runtime_cache(
             self.linear_state_pool.rebuild(num_mamba_slots + 1)
         # 3. Refresh max_seq_len (+ generic page table) for the new token budget.
         self._refresh_seq_state(config)
+        self.attn_backend.init_workspace()
         aligned_max_seq_len = self.page_table.shape[1]
         # 4. Re-capture CUDA graphs against the new tensors (reset_capture above re-armed
         #    the backend; _sync_get_memory empties the cache so freed memory is reclaimed).
@@ -2014,10 +2031,22 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
     if _dtype == torch.float16 and getattr(getattr(model_config, "quant", None), "weight_scales", False):
         # GGUF NVFP4 results are stored in the activation dtype before weight_scale_2 (~1e-4) multiplies them, which overflows fp16.
         raise ValueError("--dtype float16 with NVFP4 GGUF weight scales is unsupported (the unscaled results overflow float16); use bfloat16.")
+    kv_quant = _resolve_kv_quant(getattr(config, "kv_quant", "none"))
+    override("kv_quant", kv_quant)
+    if kv_quant == "fp8":
+        from freetoken.kvcache import _validate_kv_quant_pool
+
+        _validate_kv_quant_pool(model_config, kv_quant)
+        if config.dtype != torch.bfloat16:
+            raise ValueError("--kv-cache-dtype fp8 currently requires --dtype bfloat16")
+        if config.tp_info.size != 1:
+            raise ValueError("--kv-cache-dtype fp8 currently requires TP=1")
+        if config.speculative_algorithm is not None:
+            raise ValueError("--kv-cache-dtype fp8 does not support speculative decoding yet")
     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

@@ -33,7 +33,7 @@
 
 from .index import num_splits_for
 
-KV_CACHE_DTYPE_BYTES = 2  # every current model allocates bf16 paged KV
+KV_CACHE_DTYPE_BYTES = 2  # raw store_cache rows; FP8 KV uses its Triton quantize/scatter kernel
 EMBED_DTYPE_BYTES = 2  # embedding weights stay bf16 on the indexing() path