From b2e584515a0f7e99bf13de905440385a3fc12d52 Mon Sep 17 00:00:00 2001 From: mkornreich <52514824+mkornreich@users.noreply.github.com> Date: Mon, 24 Aug 2026 00:57:37 -0400 Subject: [PATCH] Add fp8 KV-cache quantization (--kv-dtype) New --kv-dtype auto|bfloat16|fp8_e4m3|fp8_e5m2 stores the paged KV cache in fp8 while queries stay bf16 and the FlashInfer fp8 kernels dequantize on read (scale 1.0 -- post-rope k/v magnitudes sit inside e4m3 range). Halves KV bytes/token, ~doubling the token budget that fits in memory, so long prompts that overflowed the bf16 ceiling fit; e4m3 preserves quality in testing. Threaded via EngineConfig.kv_dtype/resolved_kv_dtype (the cost model in spec_kv_bytes_per_token prices the KV dtype's itemsize), MHAKVCache.store_kv (casts bf16->fp8 before the same-width store), and the fi backend (splits q_data_type/kv_data_type in plan() via a new Context.compute_dtype, dropping the deprecated data_type alias). Default (no flag) path is unchanged. Scope: the fp8 store-quantize lives in MHAKVCache.store_kv and the dequant-on-read in the FlashInfer backend, so fp8 KV is full-attention (MHA) + FlashInfer only for now; the engine rejects fp8 on an SWA/MLA/other pool or a non-fi backend up front with a clear error rather than silently corrupting KV. Benchmark (Qwen3-1.7B, RTX 5070 8GB, memory_ratio 0.85, bf16 vs fp8_e4m3): KV budget (same 2.90 GiB): 27,142 -> 54,284 tokens (2.00x) needle @ 10/50/90% depth (14.7k ctx): 3/3 -> 3/3 (retrieval preserved) decode throughput: 98.6 -> 98.3 tok/s (unaffected) prefill throughput: 10,032 -> 9,880 tok/s (-1.5%, quantize-on-store) Co-Authored-By: Claude Opus 4.8 --- python/freetoken/attention/fi.py | 11 +++++++---- python/freetoken/core.py | 3 +++ python/freetoken/engine/config.py | 11 +++++++++++ python/freetoken/engine/engine.py | 29 +++++++++++++++++++++++++++- python/freetoken/kvcache/base.py | 4 +++- python/freetoken/kvcache/mha_pool.py | 7 +++++++ python/freetoken/server/args.py | 21 ++++++++++++++++++++ 7 files changed, 80 insertions(+), 6 deletions(-) diff --git a/python/freetoken/attention/fi.py b/python/freetoken/attention/fi.py index c9e58538..6c8f9d16 100644 --- a/python/freetoken/attention/fi.py +++ b/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, ) diff --git a/python/freetoken/core.py b/python/freetoken/core.py index ef0a539c..c26df06d 100644 --- a/python/freetoken/core.py +++ b/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 diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f3..fba943fe 100644 --- a/python/freetoken/engine/config.py +++ b/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: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 22c4a6c7..f3d5a27a 100644 --- a/python/freetoken/engine/engine.py +++ b/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) diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index ae8cf9ec..5655dc30 100644 --- a/python/freetoken/kvcache/base.py +++ b/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 diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b9..c6f92321 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/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), diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 4954c5f5..b84111e2 100644 --- a/python/freetoken/server/args.py +++ b/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"]