diff --git a/docs/cli.md b/docs/cli.md index 2738dc4ff..6eb283897 100644 --- a/docs/cli.md +++ b/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. diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 8bc52f92a..9bba8eda6 100644 --- a/python/freetoken/attention/__init__.py +++ b/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): diff --git a/python/freetoken/attention/base.py b/python/freetoken/attention/base.py index 021729929..bb9c83855 100644 --- a/python/freetoken/attention/base.py +++ b/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, diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 071ee0d45..003663ff7 100644 --- a/python/freetoken/attention/triton.py +++ b/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: diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 531a4310c..55a968b36 100644 --- a/python/freetoken/engine/config.py +++ b/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) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 5365cd659..22b685d73 100644 --- a/python/freetoken/engine/engine.py +++ b/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) diff --git a/python/freetoken/kernel/aot_models.py b/python/freetoken/kernel/aot_models.py index c5eec307d..16fc6b7e7 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/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 diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index bfb59f7df..a19afd440 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,6 +6,60 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import ( + KV_TILE_SCALE, + kv_load_e4m3_tile_f32 as _kv_load_f32, + kv_load_e4m3_tile_scaled16 as _kv_load_s16, +) + + +@triton.jit +def _kv_load_gfx12_bf16(ptrs, mask, TRANSPOSE: tl.constexpr): + # Each uint32 contains four OCP E4M3 codes. Force e64 for both word selectors. + # The non-temporal hint suits KV tiles that are consumed once per decode pass. + codes = tl.load(ptrs, mask=mask, other=0, cache_modifier=".cg") + lo = tl.inline_asm_elementwise( + "v_cvt_pk_f32_fp8_e64 $0, $1 op_sel:[0,0]", "=v,v", [codes], + dtype=tl.uint64, is_pure=True, pack=1) + hi = tl.inline_asm_elementwise( + "v_cvt_pk_f32_fp8_e64 $0, $1 op_sel:[1,0]", "=v,v", [codes], + dtype=tl.uint64, is_pure=True, pack=1) + # Every finite E4M3 value is exactly representable in BF16; no rounding is needed. + a = (lo >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + b = (lo >> 48).to(tl.uint16).to(tl.bfloat16, bitcast=True) + c = (hi >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + d = (hi >> 48).to(tl.uint16).to(tl.bfloat16, bitcast=True) + if TRANSPOSE: + a, b, c, d = tl.trans(a), tl.trans(b), tl.trans(c), tl.trans(d) + values = tl.interleave(tl.interleave(a, c), tl.interleave(b, d)) + if TRANSPOSE: + values = tl.trans(values) + return values + + +@triton.jit +def _kv_dequant_scale(scale_ptr, slots, stride, kv_head, mask): + return tl.load(scale_ptr + slots * stride + kv_head, mask=mask, other=0.0) * KV_TILE_SCALE + + +def _kv_scale_args(k, v, k_scale, v_scale): + if (k_scale is None) != (v_scale is None): + raise ValueError("K and V scales must be supplied together") + if k_scale is None: + if k.dtype == torch.uint8 or v.dtype == torch.uint8: + raise ValueError("FP8 KV codes require row scales") + return k, v, False + if k.dtype != torch.uint8 or v.dtype != torch.uint8: + raise ValueError("FP8 KV scales require uint8 code buffers") + if k.stride(-1) != 1 or v.stride(-1) != 1: + raise ValueError("FP8 KV attention requires contiguous head dimensions") + for scale, cache in ((k_scale, k), (v_scale, v)): + if scale.shape != cache.shape[:2] or scale.dtype != torch.float32 or scale.stride(-1) != 1: + raise ValueError("FP8 KV scales must be FP32 [slots, kv_heads] with contiguous heads") + if scale.device != cache.device: + raise ValueError("FP8 KV codes and scales must be on the same device") + return k_scale, v_scale, True + _MIN_BLOCK_KV = 32 @@ -19,11 +73,11 @@ _STAGE2_TILE_ELEMS = 8192 -def _kv_offset_dtype(k_cache: torch.Tensor, v_cache: torch.Tensor): +def _kv_offset_dtype(*caches: torch.Tensor | None): # Bound physical offsets, not request lengths. max_offset = max( sum((size - 1) * stride for size, stride in zip(cache.shape, cache.stride())) - for cache in (k_cache, v_cache) + for cache in caches if cache is not None ) return tl.int32 if max_offset < 1 << 31 else tl.int64 @@ -50,6 +104,13 @@ def _num_sms(device_index: int) -> int: return int(torch.cuda.get_device_properties(device_index).multi_processor_count) +@functools.lru_cache(maxsize=None) +def _gfx1201(device_index: int) -> bool: + if torch.version.hip is None: + return False + return getattr(torch.cuda.get_device_properties(device_index), "gcnArchName", "").split(":")[0] == "gfx1201" + + def decode_kv_split_cap(batch: int, num_q_heads: int, num_kv_heads: int, device_index: int) -> int: """Most KV splits per request for a decode batch: about four programs per SM, within [_MIN_SPLITS, _MAX_KV_SPLITS].""" head_blocks = triton.cdiv(num_q_heads, min(16, max(1, num_q_heads // num_kv_heads))) @@ -85,6 +146,8 @@ def _paged_attention_kernel( q_ptr, k_ptr, v_ptr, + k_scale_ptr, + v_scale_ptr, o_ptr, indptr_ptr, indices_ptr, @@ -98,6 +161,8 @@ def _paged_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -107,6 +172,7 @@ def _paged_attention_kernel( SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, KV_OFFSET_DTYPE: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, ): q_tok = tl.program_id(0) q_head = tl.program_id(1) @@ -146,14 +212,30 @@ def _paged_attention_kernel( skip_tile = tl.max(mask_n.to(tl.int32), axis=0) == 0 if not skip_tile: slots = tl.load(indices_ptr + kv_start + offs_n, mask=offs_n < kv_len, other=0).to(KV_OFFSET_DTYPE) - k = tl.load( - k_ptr - + slots[:, None] * stride_ks - + kv_head * stride_kh - + offs_d[None, :], - mask=(offs_n[:, None] < kv_len) & mask_d[None, :], - other=0.0, - ).to(tl.float32) + if HAS_KV_SCALE: + # fp8 KV: the codes carry magnitude, the per-(token, head) fp32 scale + # restores it. + s_k = tl.load( + k_scale_ptr + slots * stride_kss + kv_head, + mask=offs_n < kv_len, + other=0.0, + ) + k = _kv_load_f32( + k_ptr + + slots[:, None] * stride_ks + + kv_head * stride_kh + + offs_d[None, :], + (offs_n[:, None] < kv_len) & mask_d[None, :], + ) * s_k[:, None] + else: + k = tl.load( + k_ptr + + slots[:, None] * stride_ks + + kv_head * stride_kh + + offs_d[None, :], + mask=(offs_n[:, None] < kv_len) & mask_d[None, :], + other=0.0, + ).to(tl.float32) scores = tl.sum(q[None, :] * k, axis=1) * sm_scale scores = tl.where(mask_n, scores, -float("inf")) @@ -163,14 +245,28 @@ def _paged_attention_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new) - v = tl.load( - v_ptr - + slots[:, None] * stride_vs - + kv_head * stride_vh - + offs_d[None, :], - mask=(offs_n[:, None] < kv_len) & mask_d[None, :], - other=0.0, - ).to(tl.float32) + if HAS_KV_SCALE: + s_v = tl.load( + v_scale_ptr + slots * stride_vss + kv_head, + mask=offs_n < kv_len, + other=0.0, + ) + v = _kv_load_f32( + v_ptr + + slots[:, None] * stride_vs + + kv_head * stride_vh + + offs_d[None, :], + (offs_n[:, None] < kv_len) & mask_d[None, :], + ) * s_v[:, None] + else: + v = tl.load( + v_ptr + + slots[:, None] * stride_vs + + kv_head * stride_vh + + offs_d[None, :], + mask=(offs_n[:, None] < kv_len) & mask_d[None, :], + other=0.0, + ).to(tl.float32) acc = acc * alpha + tl.sum(p[:, None] * v, axis=0) l_i = l_i * alpha + tl.sum(p, axis=0) m_i = m_new @@ -195,6 +291,8 @@ def _decode_grouped_stage1_kernel( q_ptr, k_ptr, v_ptr, + k_scale_ptr, + v_scale_ptr, sm_scale, indptr_ptr, indices_ptr, @@ -208,6 +306,8 @@ def _decode_grouped_stage1_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_mid_ob, stride_mid_oh, stride_mid_os, @@ -228,6 +328,8 @@ def _decode_grouped_stage1_kernel( DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, KV_OFFSET_DTYPE: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, + PACKED_GFX12_KV: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -273,7 +375,8 @@ def _decode_grouped_stage1_kernel( if split_end > split_start: q = tl.load(q_ptr + q_offsets, mask=mask_h[:, None] & mask_d[None, :], other=0.0) - q = q.to(k_ptr.dtype.element_ty) + if not HAS_KV_SCALE: + q = q.to(k_ptr.dtype.element_ty) for rel_start in tl.range(split_start, split_end, BLOCK_N): rel_offs = rel_start + tl.arange(0, BLOCK_N) @@ -281,24 +384,45 @@ def _decode_grouped_stage1_kernel( logical_offs = effective_start + rel_offs slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=0).to(KV_OFFSET_DTYPE) - k = tl.load( - k_ptr + slots[None, :] * stride_ks + k_base_offsets, - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, - ) + if PACKED_GFX12_KV: + s_k = tl.load(k_scale_ptr + slots * stride_kss + kv_head, mask=mask_n, other=0.0) + packed_d = tl.arange(0, BLOCK_D // 4) + k = _kv_load_gfx12_bf16( + (k_ptr + kv_head * stride_kh).to(tl.pointer_type(tl.uint32)) + + slots[None, :] * (stride_ks // 4) + packed_d[:, None], + mask_n[None, :] & (packed_d[:, None] < D // 4), True) + elif HAS_KV_SCALE: + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16(k_ptr + slots[None, :] * stride_ks + k_base_offsets, + mask_n[None, :] & mask_d[:, None]).to(q.dtype) + else: + k = tl.load(k_ptr + slots[None, :] * stride_ks + k_base_offsets, + mask=mask_n[None, :] & mask_d[:, None], other=0.0) scores = tl.dot(q, k) * sm_scale + if HAS_KV_SCALE: + scores *= s_k[None, :] scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) - v = tl.load( - v_ptr + slots[:, None] * stride_vs + v_base_offsets, - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, - ) + if PACKED_GFX12_KV: + s_v = tl.load(v_scale_ptr + slots * stride_vss + kv_head, mask=mask_n, other=0.0) + packed_dv = tl.arange(0, BLOCK_DV // 4) + v = _kv_load_gfx12_bf16( + (v_ptr + kv_head * stride_vh).to(tl.pointer_type(tl.uint32)) + + slots[:, None] * (stride_vs // 4) + packed_dv[None, :], + mask_n[:, None] & (packed_dv[None, :] < DV // 4), False) + elif HAS_KV_SCALE: + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16(v_ptr + slots[:, None] * stride_vs + v_base_offsets, + mask_n[:, None] & mask_dv[None, :]).to(q.dtype) + else: + v = tl.load(v_ptr + slots[:, None] * stride_vs + v_base_offsets, + mask=mask_n[:, None] & mask_dv[None, :], other=0.0) m_new = tl.maximum(tl.max(scores, axis=1), m_i) alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + pv = p * s_v[None, :] if HAS_KV_SCALE else p + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -416,6 +540,8 @@ def decode_paged_attention( sinks: torch.Tensor | None = None, out: torch.Tensor | None = None, max_kv_splits: int | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, ) -> torch.Tensor: """SGLang-style split-k grouped decode attention for one query per request. @@ -451,6 +577,7 @@ def decode_paged_attention( assert sinks.numel() >= num_q_heads sinks = sinks.contiguous() + k_scale_arg, v_scale_arg, has_kv_scale = _kv_scale_args(k_cache, v_cache, k_scale, v_scale) o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q group = num_q_heads // num_kv_heads @@ -465,9 +592,10 @@ def decode_paged_attention( # a grid that fits in one wave runs about one program per SM, which then needs wider tiles to keep enough loads in flight block_n = 64 if grid[0] * grid[1] * grid[2] <= _num_sms(q.device.index) else _MIN_BLOCK_KV num_stages = 3 + num_warps = 4 # stage 1 keeps the Q tile and one K and one V tile in shared memory smem = _optin_smem_bytes(q.device.index) - elem = k_cache.element_size() + elem = q.element_size() if 2 * block_n * block_d * elem > 0.8 * smem: block_n = _MIN_BLOCK_KV # 64 KiB parts (sm_75) cannot fit a head_dim 512 Q tile next to 32-row K/V tiles: @@ -480,10 +608,35 @@ def decode_paged_attention( # a full window's split fits in one tile, where a deeper pipeline only adds overhead block_n, num_stages = split_len, 2 + if has_kv_scale and head_dim == 128 and group == 8 and _gfx1201(q.device.index): + # FP8 conversion temporaries spill with three stages on gfx1201. + num_stages = 1 + if batch == 1 and num_q_heads == 32 and block_n == 32: + num_warps = 2 + + if (has_kv_scale and head_dim == 128 and batch == 1 and num_q_heads == 16 + and group == 2 and block_n == 32 and _gfx1201(q.device.index)): + # GQA2 serving favors the smaller compatibility tile once its spills are removed. + num_stages, num_warps = 1, 2 + + packed_gfx12_kv = ( + has_kv_scale and q.dtype == torch.bfloat16 and batch == 1 + and ((head_dim == 128 and group in (4, 8)) + or (head_dim == 256 and num_q_heads == 16 and group == 8)) + and _gfx1201(q.device.index) + and k_cache.data_ptr() % 4 == 0 and v_cache.data_ptr() % 4 == 0 + and all(s % 4 == 0 for s in (*k_cache.stride()[:2], *v_cache.stride()[:2])) + ) + if packed_gfx12_kv: + # A shallow pipeline keeps the expanded KV tiles in registers on RDNA4. + block_n, num_stages, num_warps = (16, 1, 4) if head_dim == 256 else (64, 1, 2) + _decode_grouped_stage1_kernel[grid]( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, sm_scale, indptr, indices, @@ -497,13 +650,15 @@ def decode_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + k_scale_arg.stride(0), + v_scale_arg.stride(0), attn_logits.stride(0), attn_logits.stride(1), attn_logits.stride(2), attn_lse.stride(0), attn_lse.stride(1), attn_lse.stride(2), - KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), + KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache, k_scale, v_scale), GROUP=group, NUM_Q_HEADS=num_q_heads, BLOCK_D=block_d, @@ -517,7 +672,9 @@ def decode_paged_attention( D=head_dim, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, - num_warps=4, + HAS_KV_SCALE=has_kv_scale, + PACKED_GFX12_KV=packed_gfx12_kv, + num_warps=num_warps, num_stages=num_stages, ) split_tile, stage2_warps = _stage2_launch(max_kv_splits, block_dv) @@ -555,6 +712,8 @@ def _extend_attention_kernel( q_ptr, k_ptr, v_ptr, + k_scale_ptr, + v_scale_ptr, o_ptr, qo_indptr_ptr, kv_indptr_ptr, @@ -569,6 +728,8 @@ def _extend_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -578,6 +739,7 @@ def _extend_attention_kernel( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, HAS_SINKS: tl.constexpr, HAS_BLOCKS: tl.constexpr, KV_OFFSET_DTYPE: tl.constexpr, @@ -640,15 +802,16 @@ def _extend_attention_kernel( skip_tile = tl.max(tl.max(final_mask.to(tl.int32), axis=1), axis=0) == 0 if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0).to(KV_OFFSET_DTYPE) - k = tl.load( - k_ptr - + slots[None, :] * stride_ks - + kv_head * stride_kh - + offs_d[:, None], - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, - ) + if HAS_KV_SCALE: + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16(k_ptr + slots[None, :] * stride_ks + kv_head * stride_kh + offs_d[:, None], + mask_n[None, :] & mask_d[:, None]).to(q.dtype) + else: + k = tl.load(k_ptr + slots[None, :] * stride_ks + kv_head * stride_kh + offs_d[:, None], + mask=mask_n[None, :] & mask_d[:, None], other=0.0) scores = tl.dot(q.to(k.dtype), k) * sm_scale + if HAS_KV_SCALE: + scores *= s_k[None, :] scores = tl.where(final_mask, scores, -float("inf")) row_max = tl.max(scores, axis=1) @@ -657,15 +820,15 @@ def _extend_attention_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - v = tl.load( - v_ptr - + slots[:, None] * stride_vs - + kv_head * stride_vh - + offs_dv[None, :], - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, - ) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + if HAS_KV_SCALE: + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16(v_ptr + slots[:, None] * stride_vs + kv_head * stride_vh + offs_dv[None, :], + mask_n[:, None] & mask_dv[None, :]).to(q.dtype) + else: + v = tl.load(v_ptr + slots[:, None] * stride_vs + kv_head * stride_vh + offs_dv[None, :], + mask=mask_n[:, None] & mask_dv[None, :], other=0.0) + pv = p * s_v[None, :] if HAS_KV_SCALE else p + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -684,6 +847,8 @@ def _extend_attention_kernel( def _extend_attend_tile( q, m_i, l_i, acc, qk_scale, k_base, v_base, k_rows, v_rows, offs_d, offs_dv, mask_d, mask_dv, mask_n, mask, MASKED: tl.constexpr, + k_scale_ptr=None, v_scale_ptr=None, slots=None, stride_kss=0, stride_vss=0, kv_head=0, + HAS_KV_SCALE: tl.constexpr = False, ): if MASKED: k_mask = mask_n[None, :] & mask_d[:, None] @@ -691,7 +856,15 @@ def _extend_attend_tile( else: k_mask = mask_d[:, None] v_mask = mask_dv[None, :] - k = tl.load(k_base + k_rows[None, :] + offs_d[:, None], mask=k_mask, other=0.0) + if HAS_KV_SCALE: + scale_mask = mask_n if MASKED else tl.full(slots.shape, True, tl.int1) + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, scale_mask) + # Match the reconstructed prefix workspace: scale before rounding to + # the activation dtype, including short prefixes that stay inline. + k = (_kv_load_s16(k_base + k_rows[None, :] + offs_d[:, None], k_mask).to(tl.float32) + * s_k[None, :]).to(q.dtype) + else: + k = tl.load(k_base + k_rows[None, :] + offs_d[:, None], mask=k_mask, other=0.0) scores = tl.dot(q.to(k.dtype), k) if MASKED: scores = tl.where(mask, scores, -float("inf")) @@ -700,7 +873,12 @@ def _extend_attend_tile( m_new = tl.where(m_new == -float("inf"), -1e20, m_new) alpha = tl.math.exp2(m_i - m_new) p = tl.math.exp2(scores * qk_scale - m_new[:, None]) - v = tl.load(v_base + v_rows[:, None] + offs_dv[None, :], mask=v_mask, other=0.0) + if HAS_KV_SCALE: + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, scale_mask) + v = (_kv_load_s16(v_base + v_rows[:, None] + offs_dv[None, :], v_mask).to(tl.float32) + * s_v[:, None]).to(q.dtype) + else: + v = tl.load(v_base + v_rows[:, None] + offs_dv[None, :], mask=v_mask, other=0.0) acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) return m_new, l_i, acc @@ -729,6 +907,8 @@ def _extend_attention_split_kernel( v_extend_ptr, k_cache_ptr, v_cache_ptr, + k_scale_ptr, + v_scale_ptr, o_ptr, qo_indptr_ptr, kv_indptr_ptr, @@ -747,6 +927,8 @@ def _extend_attention_split_kernel( stride_kch, stride_vcs, stride_vch, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -759,6 +941,7 @@ def _extend_attention_split_kernel( HAS_SINKS: tl.constexpr, HAS_BLOCKS: tl.constexpr, KV_OFFSET_DTYPE: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -811,6 +994,7 @@ def _extend_attention_split_kernel( m_i, l_i, acc = _extend_attend_tile( q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, offs_d, offs_dv, mask_d, mask_dv, None, None, False, + k_scale_ptr, v_scale_ptr, slots, stride_kss, stride_vss, kv_head, HAS_KV_SCALE, ) n_left = (u0 - lo) // BLOCK_N for i in tl.range(0, n_left + tl.cdiv(prefix_end - u1, BLOCK_N)): @@ -823,6 +1007,7 @@ def _extend_attention_split_kernel( m_i, l_i, acc = _extend_attend_tile( q, m_i, l_i, acc, qk_scale, kc_base, vc_base, slots * stride_kcs, slots * stride_vcs, offs_d, offs_dv, mask_d, mask_dv, mask_n, final_mask, True, + k_scale_ptr, v_scale_ptr, slots, stride_kss, stride_vss, kv_head, HAS_KV_SCALE, ) ke_base = k_extend_ptr + kv_head * stride_keh @@ -891,6 +1076,8 @@ def extend_paged_attention( k_extend: torch.Tensor | None = None, v_extend: torch.Tensor | None = None, block_ends: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Block-tiled causal prefill/extend attention over paged KV cache; block_ends holds per query token the end of the multimodal span it sits in (0 for none), whose later keys the row also attends.""" @@ -911,6 +1098,7 @@ def extend_paged_attention( if block_ends is not None: assert block_ends.is_cuda and block_ends.dtype == torch.int32 and block_ends.numel() == num_q_tokens + k_scale_arg, v_scale_arg, has_kv_scale = _kv_scale_args(k_cache, v_cache, k_scale, v_scale) o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q block_ends_arg = block_ends if block_ends is not None else qo_indptr @@ -922,6 +1110,10 @@ def extend_paged_attention( block_m, block_n = _select_extend_tile( head_dim, block_d, _optin_smem_bytes(q.device.index) ) + if has_kv_scale and torch.version.hip is not None: + # FP8 conversion and row scales spill at 128x64 on gfx1201. + # Leave the BF16 tile policy and decode split geometry unchanged. + block_m, block_n = min(block_m, 64), min(block_n, 32) grid = (qo_indptr.numel() - 1, num_q_heads, triton.cdiv(max_q_len, block_m)) if k_extend is not None or v_extend is not None: assert k_extend is not None and v_extend is not None @@ -936,6 +1128,8 @@ def extend_paged_attention( v_extend, k_cache, v_cache, + k_scale_arg, + v_scale_arg, o, qo_indptr, kv_indptr, @@ -954,9 +1148,11 @@ def extend_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + k_scale_arg.stride(0), + v_scale_arg.stride(0), o.stride(0), o.stride(1), - KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), + KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache, k_scale, v_scale), GROUP=num_q_heads // num_kv_heads, D=head_dim, BLOCK_D=block_d, @@ -966,6 +1162,7 @@ def extend_paged_attention( SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, HAS_BLOCKS=block_ends is not None, + HAS_KV_SCALE=has_kv_scale, num_warps=8, num_stages=1, ) @@ -975,6 +1172,8 @@ def extend_paged_attention( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, o, qo_indptr, kv_indptr, @@ -989,9 +1188,11 @@ def extend_paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + k_scale_arg.stride(0), + v_scale_arg.stride(0), o.stride(0), o.stride(1), - KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), + KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache, k_scale, v_scale), GROUP=num_q_heads // num_kv_heads, D=head_dim, BLOCK_D=block_d, @@ -1001,6 +1202,7 @@ def extend_paged_attention( SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, HAS_BLOCKS=block_ends is not None, + HAS_KV_SCALE=has_kv_scale, num_warps=8, num_stages=1, ) @@ -1020,15 +1222,19 @@ def paged_attention( sinks: torch.Tensor | None = None, out: torch.Tensor | None = None, block_n: int = 32, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Paged causal attention for one layer. ``q`` is ``[num_query_tokens, num_q_heads, head_dim]``. KV cache tensors are flattened to ``[num_slots, num_kv_heads, head_dim]``. ``indptr`` and - ``indices`` describe each request's logical KV slots in order. + ``indices`` describe each request's logical KV slots in order. ``k_scale`` / + ``v_scale`` (``[num_slots, num_kv_heads]`` fp32) mark an fp8 codes cache. """ assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda + assert (k_scale is None) == (v_scale is None), "K and V scales come as a pair" assert q.dim() == 3 and k_cache.dim() == 3 and v_cache.dim() == 3 num_tokens, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] @@ -1041,14 +1247,25 @@ def paged_attention( assert sinks.numel() >= num_q_heads sinks = sinks.contiguous() + _kv_scale_args(k_cache, v_cache, k_scale, v_scale) o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q + has_kv_scale = k_scale is not None + k_scale_arg = k_scale if has_kv_scale else k_cache + v_scale_arg = v_scale if has_kv_scale else v_cache + if has_kv_scale: + assert k_scale.shape == v_scale.shape == (k_cache.shape[0], num_kv_heads), ( + tuple(k_scale.shape), + tuple(k_cache.shape), + ) block_d = triton.next_power_of_2(head_dim) grid = (num_tokens, num_q_heads) _paged_attention_kernel[grid]( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, o, indptr, indices, @@ -1062,15 +1279,18 @@ def paged_attention( k_cache.stride(1), v_cache.stride(0), v_cache.stride(1), + k_scale_arg.stride(0), + v_scale_arg.stride(0), o.stride(0), o.stride(1), - KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache), + KV_OFFSET_DTYPE=_kv_offset_dtype(k_cache, v_cache, k_scale, v_scale), GROUP=num_q_heads // num_kv_heads, D=head_dim, BLOCK_D=block_d, BLOCK_N=block_n, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + HAS_KV_SCALE=has_kv_scale, num_warps=8 if head_dim >= 256 else 4, num_stages=2, ) diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 60cb35a51..0ea1aea91 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -160,3 +160,17 @@ def f32_to_e4m3(x): h = tl.inline_asm_elementwise( "cvt.rn.satfinite.e4m3x2.f32 $0, 0f00000000, $1;", "=h,r", [x], dtype=tl.uint16, is_pure=True, pack=1) return h.to(tl.uint8).to(tl.float8e4nv, bitcast=True) + +KV_TILE_SCALE = tl.constexpr(256.0) + + +@jit +def kv_load_e4m3_tile_f32(ptrs, mask): + return e4m3_u8_to_f32(tl.load(ptrs, mask=mask, other=0)) + + +@jit +def kv_load_e4m3_tile_scaled16(ptrs, mask): + """Decode codes divided by 256; fold that factor into the row scales.""" + w = tl.load(ptrs, mask=mask, other=0).to(tl.uint16) + return ((w + (w & 0x80)) << 7).to(tl.float16, bitcast=True) diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py new file mode 100644 index 000000000..bde96506b --- /dev/null +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -0,0 +1,205 @@ +"""E4M3 KV storage with one FP32 scale per token, layer, head and K/V row. + +Quantize and scatter in one launch. Codes use uint8 on every GPU so allocation +and Triton signatures agree even on devices without native E4M3 support. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.e4m3_compat import ( + e4m3_f32_to_u8, + e4m3_u8_to_f32, + round_e4m3, +) + +FP8 = torch.float8_e4m3fn +KV_SCALE_DTYPE = torch.float32 +KV_QUANT_FP8 = "fp8" + + +@triton.jit +def _reconstruct_kv_prefix_kernel( + kc, vc, ks, vs, indptr, indices, prefixes, ko, vo, + KCS: tl.constexpr, VCS: tl.constexpr, + KSS: tl.constexpr, VSS: tl.constexpr, + KOS: tl.constexpr, VOS: tl.constexpr, + CAPACITY: tl.constexpr, HEADS: tl.constexpr, D: tl.constexpr, + BLOCK: tl.constexpr, +): + request = tl.program_id(1) + offset = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + row = offset // (HEADS * D) + head = offset // D % HEADS + dim = offset % D + start = tl.load(indptr + request) + end = tl.load(indptr + request + 1) + length = tl.minimum(tl.load(prefixes + request), end - start) + if tl.program_id(0) * BLOCK < length * HEADS * D: + slot = tl.load(indices + start + row, row < length, other=-1).to(tl.int64) + valid = (row < length) & (slot >= 0) & (slot < CAPACITY) + ck = tl.load(kc + slot * KCS[0] + head * KCS[1] + dim * KCS[2], valid, other=0) + cv = tl.load(vc + slot * VCS[0] + head * VCS[1] + dim * VCS[2], valid, other=0) + sk = tl.load(ks + slot * KSS[0] + head * KSS[1], valid, other=0.) + sv = tl.load(vs + slot * VSS[0] + head * VSS[1], valid, other=0.) + tl.store(ko + slot * KOS[0] + head * KOS[1] + dim * KOS[2], e4m3_u8_to_f32(ck) * sk, valid) + tl.store(vo + slot * VOS[0] + head * VOS[1] + dim * VOS[2], e4m3_u8_to_f32(cv) * sv, valid) + + +def reconstruct_kv_prefix( + k_cache: torch.Tensor, v_cache: torch.Tensor, + k_scale: torch.Tensor, v_scale: torch.Tensor, + indptr: torch.Tensor, indices: torch.Tensor, prefix_lens: torch.Tensor, + k_out: torch.Tensor, v_out: torch.Tensor, +) -> None: + """Reconstruct historical rows once, retaining their physical slot addresses.""" + assert k_cache.ndim == 3 + assert k_cache.shape == v_cache.shape == k_out.shape == v_out.shape + assert k_cache.dtype == v_cache.dtype == torch.uint8 + assert k_out.dtype == v_out.dtype == torch.bfloat16 + assert k_scale.shape == v_scale.shape == k_cache.shape[:2] + assert k_scale.dtype == v_scale.dtype == torch.float32 + assert indptr.ndim == indices.ndim == prefix_lens.ndim == 1 + assert indptr.numel() == prefix_lens.numel() + 1 + assert all(x.dtype in (torch.int32, torch.int64) and x.stride(0) == 1 + for x in (indptr, indices, prefix_lens)) + assert k_cache.is_cuda and all(x.device == k_cache.device for x in + (v_cache, k_scale, v_scale, indptr, indices, prefix_lens, k_out, v_out)) + if not indices.numel() or not prefix_lens.numel(): + return + capacity, heads, dim = k_cache.shape + _reconstruct_kv_prefix_kernel[(triton.cdiv(indices.numel() * heads * dim, 1024), prefix_lens.numel())]( + k_cache, v_cache, k_scale, v_scale, indptr, indices, prefix_lens, k_out, v_out, + k_cache.stride(), v_cache.stride(), k_scale.stride(), v_scale.stride(), + k_out.stride(), v_out.stride(), capacity, heads, dim, 1024, + num_warps=4, + ) + + +def kv_codes_dtype() -> torch.dtype: + """Raw E4M3 code dtype, independent of the target GPU.""" + return torch.uint8 + + +def alloc_codes(shape: tuple[int, ...], device: torch.device) -> torch.Tensor: + """Initialize dummy and unwritten slots to finite E4M3 zero.""" + return torch.zeros(shape, dtype=kv_codes_dtype(), device=device) + + +def codes_to_f32(codes: torch.Tensor) -> torch.Tensor: + """Decode through PyTorch for diagnostics and independent kernel checks.""" + if codes.dtype is not FP8: + codes = codes.view(FP8) + return codes.to(torch.float32) + + +@triton.jit +def _kv_quant_scatter_kernel( + k_src, + v_src, + k_dst, + v_dst, + k_scale, + v_scale, + idx_ptr, + stride_xs, # K source row pitch, in elements (the qkv slice is wider than one row) + stride_vx, # V source row pitch. K and V need not share one: a .clamp() on one side + # leaves it densely packed, so reusing K's pitch reads V off its rows. + stride_kd, # K cache row pitch, in elements (== HEADS * D) + stride_vd, + stride_ks, # scale row pitch, in elements (== HEADS) + stride_vs, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """One program per (token, kv_head): quantize the row and write it to slot + ``idx_ptr[token]`` (the same ``out_loc`` the bf16 store scatters through).""" + t = tl.program_id(0) + h = tl.program_id(1) + pos = tl.load(idx_ptr + t).to(tl.int64) # int64: slots * row can pass 2**31 + d = tl.arange(0, BLOCK_D) + mask = d < D + + off = h * D + d + xk = tl.load(k_src + t * stride_xs + off, mask=mask, other=0.0).to(tl.float32) + xv = tl.load(v_src + t * stride_vx + off, mask=mask, other=0.0).to(tl.float32) + + # 448 == e4m3 finite max; 1e-10 is the amax floor of the activation quant in + # kernel/triton/fp8_block_linear.py (literals keep the kernel self-contained). + sk = tl.maximum(tl.max(tl.abs(xk), axis=0), 1e-10) / 448.0 + sv = tl.maximum(tl.max(tl.abs(xv), axis=0), 1e-10) / 448.0 + qk = tl.clamp(xk / sk, -448.0, 448.0) + qv = tl.clamp(xv / sv, -448.0, 448.0) + + # Round to the E4M3 grid before packing the code bits (round-to-nearest-even). + out_k = e4m3_f32_to_u8(round_e4m3(qk)) + out_v = e4m3_f32_to_u8(round_e4m3(qv)) + + tl.store(k_dst + pos * stride_kd + h * D + d, out_k, mask=mask) + tl.store(v_dst + pos * stride_vd + h * D + d, out_v, mask=mask) + tl.store(k_scale + pos * stride_ks + h, sk) + tl.store(v_scale + pos * stride_vs + h, sv) + + +def quantize_kv_to_cache( + k: torch.Tensor, + v: torch.Tensor, + out_loc: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + k_scale: torch.Tensor, + v_scale: torch.Tensor, +) -> None: + """Quantize fresh K/V rows into an fp8 KV pool. + + ``k``/``v`` : ``[T, num_kv_heads * head_dim]`` compute-dtype rows -- exactly what + the attention backends hand to ``store_kv`` (a slice of the qkv projection, + so the row pitch may be wider than the row itself). + ``out_loc`` : ``[T]`` device slot index per row (int32 or int64). + ``*_cache`` : ``[num_slots, num_kv_heads, head_dim]`` of :func:`kv_codes_dtype`. + ``*_scale`` : ``[num_slots, num_kv_heads]`` fp32, indexed by the SAME slot. + """ + tokens = k.shape[0] + if tokens == 0: + return + assert k.dim() == 2 and v.shape == k.shape, (k.shape, v.shape) + assert k.stride(1) == 1 and v.stride(1) == 1, "K/V rows must be contiguous" + heads, dim = k_cache.shape[1], k_cache.shape[2] + assert k.shape[1] == heads * dim, (tuple(k.shape), tuple(k_cache.shape)) + assert k_cache.shape == v_cache.shape, (tuple(k_cache.shape), tuple(v_cache.shape)) + assert k_scale.shape == v_scale.shape == (k_cache.shape[0], heads), ( + tuple(k_scale.shape), + tuple(k_cache.shape), + ) + assert k_cache.dtype == v_cache.dtype == kv_codes_dtype(), ( + k_cache.dtype, + kv_codes_dtype(), + ) + assert k_scale.dtype == v_scale.dtype == KV_SCALE_DTYPE + assert k_cache.stride(1) == v_cache.stride(1) == dim + assert k_cache.stride(2) == v_cache.stride(2) == 1 + assert k_scale.stride(1) == v_scale.stride(1) == 1 + assert out_loc.shape == (tokens,) and out_loc.dtype in (torch.int32, torch.int64) + assert out_loc.stride(0) == 1 + assert all(t.device == k.device for t in (v, out_loc, k_cache, v_cache, k_scale, v_scale)) + _kv_quant_scatter_kernel[(tokens, heads)]( + k, + v, + k_cache, + v_cache, + k_scale, + v_scale, + out_loc, + k.stride(0), + v.stride(0), + k_cache.stride(0), + v_cache.stride(0), + k_scale.stride(0), + v_scale.stride(0), + D=dim, + BLOCK_D=triton.next_power_of_2(dim), + num_warps=1, + ) diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 8d038e2fb..70f9f4cfe 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -80,6 +80,22 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: return MHAKVCache +def _validate_kv_quant_pool(model_config, kv_quant: str) -> None: + if kv_quant == "none": + return + from freetoken.attention import AttnType + + # Linear layers own recurrent state, not paged KV. Quantize only the full group. + specs = [s for s in model_config.kv_cache_group_specs() if s.num_layers > 0] + if ( + kv_quant != "fp8" + or not specs + or any(s.attn_type != AttnType.FULL for s in specs) + or (getattr(model_config, "has_linear_attention", False) and len(specs) != 1) + ): + raise ValueError("FP8 KV requires MHA/GQA full-attention storage") + + def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dtype): """Build the engine's KV pool for ``num_pages`` USABLE pages (the dummy page and every secondary tier -- window pool, index slab, state rings -- are derived here or inside @@ -89,6 +105,8 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt from .dsv4.v4_pool import DSV4PagedKVCache model_config = config.model_config + kv_quant = getattr(config, "kv_quant", "none") + _validate_kv_quant_pool(model_config, kv_quant) if resolve_pool_class(model_config) is DSV4PagedKVCache: # DSV4 is driven by the generic CacheManager over the shared page table; the pool is # the only DSV4-specific piece (the swa_pool plug-in: window tier + cmp/idx/state @@ -137,6 +155,7 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt dtype=dtype, num_req_slots=config.max_running_req + 1, # + 1 for the dummy request row num_speculative_tokens=getattr(config, "decode_rows_per_req", 1) - 1, + kv_quant=kv_quant, ) @@ -149,7 +168,9 @@ def create_kvcache_pool( num_swa_tokens: int | None = None, num_req_slots: int | None = None, num_speculative_tokens: int = 0, + kv_quant: str = "none", ) -> BaseKVCachePool: + _validate_kv_quant_pool(model_config, kv_quant) if model_config.has_swa_attention: from .hybrid_swa_pool import HybridSWAKVCache @@ -278,6 +299,7 @@ def create_kvcache_pool( device=device, dtype=dtype, layer_ids=layer_ids, + kv_quant=kv_quant, ) diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index 7d6c3dc84..b8498a792 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -16,9 +16,18 @@ class CacheRebuildRejected(Exception): this is recoverable, unlike a failure after the free.""" +def kv_storage_bytes_per_elem(config) -> int: + quant = getattr(config, "kv_quant", "none") + if quant == "fp8": + return 1 + if quant != "none": + raise ValueError(f"unknown kv_quant {quant!r}") + return config.dtype.itemsize + + def spec_kv_bytes_per_token(spec, config) -> int: - """One paged-KV group's bytes per token: (1|2 slabs) x head_dim x local kv heads x dtype - x layers, plus the bf16 DSA index-key slab when the spec carries indexer dims. Pure + """One paged-KV group's bytes per token, including FP8 row scales when enabled, + plus the bf16 DSA index-key slab when the spec carries indexer dims. Pure 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.) @@ -29,9 +38,14 @@ def spec_kv_bytes_per_token(spec, config) -> int: (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_storage_bytes_per_elem(config) * spec.num_layers ) + if getattr(config, "kv_quant", "none") == "fp8": + per_token += ( + (1 if spec.mla else 2) * spec.num_layers + * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) * 4 + ) return per_token + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio @@ -78,20 +92,31 @@ def kv_cost(cls, config, **kwargs) -> tuple[int, int, int, int]: ``num_swa_pages=`` (a rebuild's target window, priced before the config override is written).""" + @classmethod + def runtime_cost(cls, config, **kwargs) -> tuple[int, int, int, int]: + """Pool cost plus geometry-dependent attention workspace, when required.""" + return cls.kv_cost(config, **kwargs) + + @property + def prefill_workspace_shape(self) -> tuple[int, int, int] | None: + return None + @classmethod def solve_num_pages(cls, config, available_memory: int) -> int: """Largest usable page count fitting ``available_memory`` (bytes measured by the engine: memory_ratio x baseline free minus resident weights and any sibling pool's fixed cost), honoring ``config.num_page_override``.""" - cache_per_page, fixed_cache_size, _, _ = cls.kv_cost(config) + cache_per_page, fixed_cache_size, _, _ = cls.runtime_cost(config) num_pages = config.num_page_override if num_pages is None: num_pages = (available_memory - fixed_cache_size) // cache_per_page assert num_pages > 1, "Not enough memory for KV cache, try reducing --num-pages" real_kv_size = num_pages * cache_per_page + fixed_cache_size + if getattr(config, "kv_quant", "none") == "fp8" and real_kv_size > available_memory: + raise ValueError("FP8 KV pool and prefill workspace exceed the available memory budget") logger.info( f"Allocating {num_pages * config.page_size} tokens for KV cache, " - f"K + V = {mem_GB(real_kv_size)}" + f"pool and workspace budget = {mem_GB(real_kv_size)}" ) return num_pages @@ -125,7 +150,7 @@ def validate_rebuild( cost_kwargs = { k: v for k, v in targets.items() if v is not None and k in cost_params } - cache_per_page, fixed_cache_size, _, _ = type(self).kv_cost(config, **cost_kwargs) + cache_per_page, fixed_cache_size, _, _ = type(self).runtime_cost(config, **cost_kwargs) budget = net_cache_budget_bytes( config.memory_ratio, baseline_free, weights_bytes, fixed_cache_size + extra_fixed_bytes, @@ -180,6 +205,16 @@ def device(self) -> torch.device: ... @abstractmethod def dtype(self) -> torch.dtype: ... + @property + def store_dtype(self) -> torch.dtype: + return self.dtype + + def k_scale(self, index: int) -> torch.Tensor | None: + return None + + def v_scale(self, index: int) -> torch.Tensor | None: + return None + @property @abstractmethod def num_layers(self) -> int: ... diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b96..c828db39e 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/python/freetoken/kvcache/mha_pool.py @@ -9,6 +9,17 @@ from .base import BaseKVCachePool +def _kv_store_dtype(dtype: torch.dtype, kv_quant: str) -> torch.dtype: + """Storage dtype of the KV buffer for a quantization mode.""" + if kv_quant == "none": + return dtype + if kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import kv_codes_dtype + + return kv_codes_dtype() + raise ValueError(f"unknown kv_quant {kv_quant!r}") + + class MHAKVCache(BaseKVCachePool): """ Base class for key-value caches. @@ -20,6 +31,12 @@ class MHAKVCache(BaseKVCachePool): that hold no paged KV; passing the full-attention layer ids here allocates one storage slab per KV layer (not per model layer) and remaps the global id to its dense slot, avoiding a multiple-x over-allocation of unused slabs. + + ``kv_quant="fp8"`` stores rows as e4m3 codes and every + ``(token, slab, layer, kv head)`` row carries one fp32 scale (see + :mod:`freetoken.kernel.triton.kv_quant`). The codes buffer keeps the exact same + shape as the 16-bit one, so ``k_cache``/``v_cache`` and every index into them are + unchanged -- only the element type, and ``store_kv``'s write path, differ. """ def __init__( @@ -32,10 +49,16 @@ def __init__( dtype: torch.dtype, device: torch.device, layer_ids: Sequence[int] | None = None, + kv_quant: str = "none", ) -> None: + _kv_store_dtype(dtype, kv_quant) + if kv_quant == "fp8" and dtype != torch.bfloat16: + raise ValueError("FP8 KV currently requires BF16 compute") tp_info = get_tp_info() local_kv_heads = div_even(num_kv_heads, tp_info.size, allow_replicate=True) self._num_layers = num_layers + self.kv_quant = kv_quant + self._compute_dtype = dtype if layer_ids is None: num_storage_layers = num_layers self._layer_map: list[int] | None = None @@ -47,14 +70,35 @@ def __init__( raise ValueError(f"KV layer id {global_id} outside [0, {num_layers})") layer_map[global_id] = dense self._layer_map = layer_map - self._kv_buffer = torch.empty( - (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, - ) + self._device = device + self._alloc(num_pages, page_size, num_storage_layers, local_kv_heads, head_dim) + + def _alloc( + self, + num_pages: int, + page_size: int, + num_storage_layers: int, + local_kv_heads: int, + head_dim: int, + ) -> None: + """Allocate storage and scales; zero FP8 dummy rows to avoid NaN codes.""" + shape = (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim) + if self.kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import alloc_codes + + self._kv_buffer = alloc_codes(shape, self._device) + self._scale_buffer = torch.zeros( + (2, num_storage_layers, num_pages * page_size, local_kv_heads), + device=self._device, + dtype=torch.float32, + ) + else: + self._kv_buffer = torch.empty( + shape, device=self._device, dtype=self._compute_dtype + ) + self._scale_buffer = None self._k_buffer = self._kv_buffer[0] self._v_buffer = self._kv_buffer[1] - self._device = device self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) def rebuild(self, num_pages: int) -> None: @@ -65,22 +109,32 @@ def rebuild(self, num_pages: int) -> None: refreshed. Object identity is preserved so cached backend references stay valid. """ _, num_storage_layers, _old_pages, page_size, local_kv_heads, head_dim = self._kv_buffer.shape - dtype = self._kv_buffer.dtype device = self._device self._k_buffer = None self._v_buffer = None self._kv_buffer = None + self._scale_buffer = None if device.type == "cuda": torch.cuda.synchronize(device) torch.cuda.empty_cache() - self._kv_buffer = torch.empty( - (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), - device=device, - dtype=dtype, - ) - self._k_buffer = self._kv_buffer[0] - self._v_buffer = self._kv_buffer[1] - self._storage_shape = (num_pages * page_size, local_kv_heads, head_dim) + self._alloc(num_pages, page_size, num_storage_layers, local_kv_heads, head_dim) + + @classmethod + def runtime_cost(cls, config, **kwargs) -> tuple[int, int, int, int]: + per_page, fixed, tokens, reserve = cls.kv_cost(config, **kwargs) + if getattr(config, "kv_quant", "none") != "fp8": + return per_page, fixed, tokens, reserve + specs = [s for s in config.model_config.kv_cache_group_specs() if s.num_layers > 0] + spec = specs[0] if len(specs) == 1 else config.model_config + heads = div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) + workspace_per_page = config.page_size * 2 * heads * spec.head_dim * 2 + # One workspace is shared by all layers; both it and storage include the dummy page. + total = per_page + workspace_per_page + return total, fixed + total, tokens, reserve + + @property + def prefill_workspace_shape(self) -> tuple[int, int, int] | None: + return self._storage_shape if self.kv_quant == "fp8" else None @classmethod def kv_cost(cls, config) -> tuple[int, int, int, int]: @@ -101,7 +155,11 @@ def rebuild_from_config( def unit_bytes(self) -> tuple[int, int]: buf = self._kv_buffer tokens = int(buf.shape[2]) * int(buf.shape[3]) - return int(buf.numel() * buf.element_size()) // tokens, 0 + kv = int(buf.numel() * buf.element_size()) // tokens + if self._scale_buffer is not None: + sc = self._scale_buffer + kv += int(sc.numel() * sc.element_size()) // tokens + return kv, 0 def _dense(self, layer_id: int) -> int: if self._layer_map is None: @@ -117,6 +175,16 @@ def k_cache(self, index: int) -> torch.Tensor: def v_cache(self, index: int) -> torch.Tensor: return self._v_buffer[self._dense(index)] + def k_scale(self, index: int) -> torch.Tensor | None: + if self._scale_buffer is None: + return None + return self._scale_buffer[0][self._dense(index)] + + def v_scale(self, index: int) -> torch.Tensor | None: + if self._scale_buffer is None: + return None + return self._scale_buffer[1][self._dense(index)] + def store_kv( self, k: torch.Tensor, @@ -124,9 +192,22 @@ def store_kv( out_loc: torch.Tensor, layer_id: int, ) -> None: + dense = self._dense(layer_id) + if self.kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import quantize_kv_to_cache + + quantize_kv_to_cache( + k=k, + v=v, + out_loc=out_loc, + k_cache=self._k_buffer[dense].view(self._storage_shape), + v_cache=self._v_buffer[dense].view(self._storage_shape), + k_scale=self._scale_buffer[0][dense], + v_scale=self._scale_buffer[1][dense], + ) + return from freetoken.kernel import store_cache - dense = self._dense(layer_id) store_cache( k_cache=self._k_buffer[dense].view(self._storage_shape), v_cache=self._v_buffer[dense].view(self._storage_shape), @@ -141,6 +222,15 @@ def device(self) -> torch.device: @property def dtype(self) -> torch.dtype: + """The COMPUTE dtype (kvcache/base.py): what ``store_kv`` receives and what + backends size their scratch with -- still 16-bit on an fp8 pool.""" + return self._compute_dtype + + @property + def store_dtype(self) -> torch.dtype: + """Element type of ``_kv_buffer``: fp8/uint8 codes when quantized, else the + compute dtype. Reading the buffer itself needs THIS, and a backend that cannot + apply the row scales never gets a quantized pool (supports_fp8_kv gate).""" return self._kv_buffer.dtype @property diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 3d75143e2..f0418a197 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -441,6 +441,14 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Print one decode scheduler status line every N decode forwards.", ) + parser.add_argument( + "--kv-cache-dtype", + dest="kv_quant", + choices=("auto", "bf16", "fp8"), + default="auto", + help="KV storage dtype. FP8 currently requires BF16 compute, TP=1, full attention and the Triton backend.", + ) + kv_capacity_group = parser.add_mutually_exclusive_group() kv_capacity_group.add_argument( "--num-pages", diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py new file mode 100644 index 000000000..7e80d05fc --- /dev/null +++ b/tests/engine/test_kv_quant_config.py @@ -0,0 +1,74 @@ +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.attention import AttnType +from freetoken.distributed import DistributedInfo +from freetoken.engine.config import EngineConfig +from freetoken.models.config import KVCacheGroupSpec + + +def _config(*, dtype=torch.bfloat16, tp=1, kind=AttnType.FULL, **kwargs): + spec = KVCacheGroupSpec(name="full", layer_ids=tuple(range(48)), num_kv_heads=4, + head_dim=128, attn_type=kind, sliding_window=None) + config = EngineConfig(model_path="/tmp/fp8-kv-test", dtype=dtype, + tp_info=DistributedInfo(rank=0, size=tp), **kwargs) + object.__setattr__(config, "model_config", SimpleNamespace( + has_swa_attention=kind == AttnType.SWA, has_linear_attention=False, + is_moe=False, num_layers=48, expert_quant="none", + kv_cache_group_specs=lambda: (spec,), + )) + return config + + +@pytest.fixture(autouse=True) +def _hardware(monkeypatch): + from freetoken.engine import engine + + monkeypatch.setattr(engine, "is_sm100_family", lambda: True) + monkeypatch.setattr(engine, "is_sm90_family", lambda: False) + monkeypatch.setattr(engine, "_flashinfer_available", lambda: True) + monkeypatch.setattr(engine, "_sgl_flash_attn_available", lambda: True) + + +def test_fp8_auto_selects_a_scale_aware_backend(): + from freetoken.engine.engine import _adjust_config + + config = _config(kv_quant="fp8") + _adjust_config(config) + assert config.attention_backend == "triton" + + +@pytest.mark.parametrize("backend", ["fi", "fa", "trtllm", "triton,fi"]) +def test_fp8_rejects_backends_that_do_not_read_row_scales(backend): + from freetoken.engine.engine import _adjust_config + + with pytest.raises(ValueError, match="requires --attention-backend triton"): + _adjust_config(_config(kv_quant="fp8", attention_backend=backend)) + + +@pytest.mark.parametrize("kwargs,match", [ + ({"dtype": torch.float16}, "bfloat16"), + ({"tp": 2}, "TP=1"), + ({"speculative_algorithm": "mtp"}, "speculative decoding"), + ({"kind": AttnType.SWA}, "MHA/GQA"), + ({"kind": AttnType.MLA}, "MHA/GQA"), + ({"kind": AttnType.QSA}, "MHA/GQA"), + ({"kv_quant": "int8"}, "unknown --kv-cache-dtype"), +]) +def test_unsupported_fp8_configs_fail_before_weights_load(kwargs, match): + from freetoken.engine.engine import _adjust_config + + args = {"kv_quant": "fp8", **kwargs} + with pytest.raises(ValueError, match=match): + _adjust_config(_config(**args)) + + +def test_qwen3_pool_cost_includes_both_scale_slabs(): + from freetoken.kvcache.mha_pool import MHAKVCache + + plain = _config() + quant = _config(kv_quant="fp8") + assert MHAKVCache.kv_cost(plain)[0] == 98304 + assert MHAKVCache.kv_cost(quant)[0] == 50688 diff --git a/tests/kernels/test_kv_fp8.py b/tests/kernels/test_kv_fp8.py new file mode 100644 index 000000000..deed99852 --- /dev/null +++ b/tests/kernels/test_kv_fp8.py @@ -0,0 +1,263 @@ +"""FP8 (e4m3) KV quantization: the store kernel against an independent oracle. + +Expectations come from a brute-force nearest-code search over an e4m3 table decoded +from first principles (sign / exponent / mantissa), NOT from the kernels' own +rounding helpers -- so a drift in ``round_e4m3`` or in the new ``e4m3_f32_to_u8`` +encoder fails here instead of being blessed by itself. +""" + +from __future__ import annotations + +import pytest +import torch + +if not torch.cuda.is_available(): # pragma: no cover + pytest.skip("CUDA required", allow_module_level=True) + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.kernel.triton.kv_quant import ( + KV_SCALE_DTYPE, + alloc_codes, + codes_to_f32, + kv_codes_dtype, + quantize_kv_to_cache, +) + +DEV = torch.device("cuda") +FP8_MAX = 448.0 +CODE_448 = 0x7E + + +@pytest.mark.parametrize("prefix", [0, 1, 65, 257]) +@pytest.mark.parametrize("dim", [63, 128]) +def test_prefix_reconstruction_strides_and_physical_slots(prefix, dim): + from freetoken.kernel.triton.kv_quant import reconstruct_kv_prefix + + torch.manual_seed(2202) + capacity, heads = 513, 4 + kc = torch.randint(0, 127, (capacity, heads * 2, dim + 3), device=DEV, dtype=torch.uint8)[:, ::2, :dim] + vc = torch.randint(128, 255, (capacity, heads, dim * 2), device=DEV, dtype=torch.uint8)[:, :, ::2] + ks = torch.rand(capacity, heads * 3, device=DEV)[:, ::3] + vs = torch.rand(capacity, heads * 2, device=DEV)[:, ::2] * 0.01 + ko = torch.full((capacity, heads * 2, dim), -42., device=DEV, dtype=torch.bfloat16)[:, ::2] + vo = torch.full((capacity, heads, dim * 2), 42., device=DEV, dtype=torch.bfloat16)[:, :, ::2] + indices = torch.randperm(capacity, device=DEV)[:prefix + 17].int() + indptr = torch.tensor([0, indices.numel()], device=DEV, dtype=torch.int32) + lengths = torch.tensor([prefix], device=DEV, dtype=torch.int32) + reconstruct_kv_prefix(kc, vc, ks, vs, indptr, indices, lengths, ko, vo) + slots = indices[:prefix].long() + for codes, scales, output in [(kc, ks, ko), (vc, vs, vo)]: + expected = (codes.contiguous().view(torch.float8_e4m3fn).float() * scales[..., None]).bfloat16() + torch.testing.assert_close(output[slots], expected[slots], atol=0, rtol=0) + untouched = torch.ones(capacity, device=DEV, dtype=torch.bool) + untouched[slots] = False + assert (ko[untouched] == -42).all() + assert (vo[untouched] == 42).all() + + +def _init_tp() -> None: + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +@pytest.fixture(autouse=True) +def _tp(): + _init_tp() + + +def _e4m3_table() -> dict[int, float]: + """Every finite e4m3fn value, decoded by hand from its bit fields. + + code = S EEEE MMM, exponent bias 7: normal (E>0) -> (-1)^S * 2^(E-7) * (1+M/8); + subnormal (E==0) -> (-1)^S * 2^-6 * (M/8). E==15 & M==7 (0x7F/0xFF) is the NaN + pattern and is dropped, which caps the format at +-448. + """ + values: dict[int, float] = {} + for code in range(256): + sign = -1.0 if (code >> 7) & 1 else 1.0 + exp, mant = (code >> 3) & 0x0F, code & 0x07 + if exp == 15 and mant == 7: + continue + values[code] = sign * ( + (2.0**-6) * (mant / 8.0) if exp == 0 else (2.0 ** (exp - 7)) * (1.0 + mant / 8.0) + ) + return values + + +E4M3_VALUES = _e4m3_table() +CODES = torch.tensor(sorted(E4M3_VALUES), dtype=torch.int32) +GRID = torch.tensor([E4M3_VALUES[c] for c in CODES.tolist()], dtype=torch.float64) +NEGATIVE_ZERO = 0x80 + + +def _as_bytes(t: torch.Tensor) -> torch.Tensor: + """A code buffer as raw bytes (the fp8 view on sm_89+, uint8 below it).""" + return t if t.dtype == torch.uint8 else t.view(torch.uint8) + + +def _canonical_zero(codes: torch.Tensor) -> torch.Tensor: + """Fold -0.0 (0x80) onto +0.0 (0x00). + + The two store paths legitimately disagree on zero's sign bit: sm_89+ converts with + ``x.to(fp8e4nv)`` (keeps it), the emulated path rounds first and + ``round_e4m3(-0.0)`` is documented to return +0.0. They decode to the same number, + so a test that compares bytes must not care which one it got. + """ + return torch.where(codes == NEGATIVE_ZERO, torch.zeros_like(codes), codes) + + +def _ref_codes(x: torch.Tensor) -> torch.Tensor: + """Independent encoder: nearest grid value by brute force, ties resolved to the + EVEN code (RNE). ``x`` is any shape; returns int32 codes.""" + grid = GRID.to(x.device) + codes = CODES.to(x.device) + dist = (x.to(torch.float64).reshape(-1, 1) - grid.unsqueeze(0)).abs() + near = dist == dist.min(dim=-1, keepdim=True).values + big = 1 << 30 + codes = codes.unsqueeze(0).expand_as(dist) + any_code = torch.where(near, codes, torch.full_like(codes, big)) + even = near & (codes % 2 == 0) + even_code = torch.where(even, codes, torch.full_like(codes, big)) + has_even = (even_code < big).any(dim=-1) + return torch.where(has_even, even_code.min(dim=-1).values, any_code.min(dim=-1).values) + + +def _store(rows_k: torch.Tensor, rows_v: torch.Tensor): + """Quantize ``[T, heads, dim]`` rows into a fresh code buffer, returning + ``(k_codes, v_codes, k_scales, v_scales)``.""" + _init_tp() + tokens, heads, dim = rows_k.shape + k_cache = alloc_codes((tokens, heads, dim), DEV) + v_cache = alloc_codes((tokens, heads, dim), DEV) + k_scale = torch.zeros((tokens, heads), dtype=KV_SCALE_DTYPE, device=DEV) + v_scale = torch.zeros((tokens, heads), dtype=KV_SCALE_DTYPE, device=DEV) + quantize_kv_to_cache( + k=rows_k.reshape(tokens, -1), + v=rows_v.reshape(tokens, -1), + out_loc=torch.arange(tokens, dtype=torch.int32, device=DEV), + k_cache=k_cache, + v_cache=v_cache, + k_scale=k_scale, + v_scale=v_scale, + ) + torch.cuda.synchronize() + return k_cache, v_cache, k_scale, v_scale + + +def test_scale_is_amax_over_e4m3_max(): + torch.manual_seed(0) + k = (torch.randn(4, 2, 64, device=DEV, dtype=torch.bfloat16) * 3.0).to(torch.bfloat16) + v = torch.randn_like(k) + _, _, k_scale, v_scale = _store(k, v) + # The kernel widens to fp32 before the amax/divide, so the reference must too: + # a bf16 intermediate would round the expected scale and hide a precision bug. + torch.testing.assert_close( + k_scale, k.abs().to(torch.float32).amax(dim=-1) / FP8_MAX, rtol=1e-6, atol=0 + ) + torch.testing.assert_close( + v_scale, v.abs().to(torch.float32).amax(dim=-1) / FP8_MAX, rtol=1e-6, atol=0 + ) + + +def test_codes_match_the_reference_quantizer_and_reconstruction_is_close(): + torch.manual_seed(1) + tokens, heads, dim = 8, 3, 128 + # Feed the rows as a qkv slice, the way the attention backends really do: the + # row pitch is then wider than the row, which the store kernel must honour. + qkv = torch.randn(tokens, heads * dim * 3, device=DEV, dtype=torch.bfloat16) + qkv[:, heads * dim : 2 * heads * dim] *= 5.0 # K: large magnitude + qkv[:, 2 * heads * dim :] *= 0.01 # V: subnormal end of the e4m3 grid + _, k_rows, v_rows = qkv.split(heads * dim, dim=-1) + k = k_rows.view(tokens, heads, dim) + v = v_rows.view(tokens, heads, dim).clamp(-FP8_MAX, FP8_MAX) + k_cache, v_cache, k_scale, v_scale = _store(k, v) + + for rows, cache, scale in ((k, k_cache, k_scale), (v, v_cache, v_scale)): + f32 = rows.to(torch.float32) + ref_scale = f32.abs().amax(dim=-1, keepdim=True) / FP8_MAX + expected = _ref_codes((f32 / ref_scale).clamp(-FP8_MAX, FP8_MAX)) + got = _canonical_zero(_as_bytes(cache).reshape(-1).to(torch.int32)) + expected = _canonical_zero(expected) + assert torch.equal(got, expected), ( + f"{int((got != expected).sum())} code mismatches of {got.numel()}" + ) + deq = codes_to_f32(cache) * scale.unsqueeze(-1) + err = (deq - f32).abs().max(dim=-1).values + assert torch.all(err <= 0.08 * f32.abs().amax(dim=-1)), float(err.max()) + + +def test_encoder_inverts_the_grid_through_the_scale_one_path(): + """Pack every e4m3 grid value into a row that also holds 448.0: the row scale is + then exactly 1.0, so each stored byte IS the encoder's answer for that value.""" + dim, per_row = 256, 255 + pairs = sorted(E4M3_VALUES.items()) + tokens = -(-len(pairs) // per_row) + rows = torch.zeros(tokens, 1, dim, dtype=torch.float32) + expected = torch.zeros(tokens, dim, dtype=torch.uint8) + for t in range(tokens): + rows[t, 0, 0] = FP8_MAX # the amax anchor + expected[t, 0] = CODE_448 + for j in range(per_row): + i = t * per_row + j + if i >= len(pairs): + break + code, value = pairs[i] + rows[t, 0, j + 1] = value + expected[t, j + 1] = code + + k_cache, _, k_scale, _ = _store( + rows.to(DEV, dtype=torch.bfloat16), + torch.zeros(tokens, 1, dim, dtype=torch.bfloat16, device=DEV), + ) + assert torch.equal(k_scale[:, 0], torch.ones_like(k_scale[:, 0])) + got = _canonical_zero(_as_bytes(k_cache)[:, 0, :]) + want = _canonical_zero(expected.to(DEV)) + bad = got != want + assert not bad.any().item(), ( + f"{int(bad.sum())} of {want.numel()} grid values round-tripped wrong; " + f"first at {bad.nonzero()[0].tolist()}: expected " + f"{want[bad][0].item():#x} got {got[bad][0].item():#x}" + ) + + +def test_zero_row_stays_finite_and_exact(): + k = torch.zeros(2, 2, 32, device=DEV, dtype=torch.bfloat16) + k_cache, _, k_scale, _ = _store(k, k.clone()) + assert torch.isfinite(k_scale).all() + assert (k_scale > 0).all(), "an all-zero row must still store a usable scale" + assert (codes_to_f32(k_cache) == 0).all() + + +def test_codes_are_plain_bytes_and_the_kernel_decode_matches_torch(): + """The KV codec never puts an fp8 type in front of Triton. + + Codes live in a uint8 buffer on EVERY architecture and the kernel widens them with + the software decoder, while the expectation below is torch's OWN e4m3 cast of those + very bytes. That pins the one claim the design rests on: byte for byte, the + software decode reads what a native fp8 unit would -- which is what lets the + quantized cache behave identically on GPUs where the fp8 type is illegal. + """ + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 + + rows = torch.randn(5, 2, 64, device=DEV, dtype=torch.bfloat16) * 3.0 + k_cache, _, _, _ = _store(rows, rows.clone()) + assert kv_codes_dtype() is torch.uint8, "keep the fp8 type out of kernel signatures" + assert k_cache.dtype is torch.uint8 and k_cache.element_size() == 1 + want = codes_to_f32(k_cache) # torch reinterprets these bytes as e4m3 and casts + + @triton.jit + def read_out(codes_ptr, out_ptr, n, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + got = kv_load_e4m3_tile_f32(codes_ptr + offs, offs < n) + tl.store(out_ptr + offs, got, mask=offs < n) + + n = k_cache.numel() + out = torch.zeros(n, dtype=torch.float32, device=DEV) + read_out[(triton.cdiv(n, 256),)](k_cache, out, n, BLOCK=256) + flat = want.reshape(-1) + assert torch.equal(out, flat), ( + f"{int((out != flat).sum())} of {n} codes decode differently from torch's cast" + ) diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index d86467b5c..0f795601d 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -52,7 +52,8 @@ def _reference_paged_attention( @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs a GPU") @pytest.mark.parametrize("path", ["paged", "decode", "extend", "split_extend"]) @pytest.mark.parametrize("offset_bits", [31, 32]) -def test_triton_attention_high_physical_slots(path: str, offset_bits: int): +@pytest.mark.parametrize("kv_quant", ["bf16", "fp8"]) +def test_triton_attention_high_physical_slots(path: str, offset_bits: int, kv_quant: str): from freetoken.kernel.triton.attention import ( decode_paged_attention, extend_paged_attention, @@ -61,11 +62,11 @@ def test_triton_attention_high_physical_slots(path: str, offset_bits: int): torch.manual_seed(1347) device = torch.device("cuda") - head_dim, num_kv_heads, num_q_heads = 64, 2, 4 + head_dim, num_kv_heads, num_q_heads = (128, 2, 8) if kv_quant == "fp8" else (64, 2, 4) stride = num_kv_heads * head_dim low, count = 128, 66 high = (1 << 32) // stride + low - required = (high + count) * stride * 2 + required = (high + count) * (stride * (1 if kv_quant == "fp8" else 2) + num_kv_heads * 4) if torch.cuda.mem_get_info()[0] < required + (1 << 30): # Release inactive cached blocks before deciding this case cannot fit. torch.cuda.empty_cache() @@ -74,38 +75,55 @@ def test_triton_attention_high_physical_slots(path: str, offset_bits: int): # Share K/V storage to keep the real allocation near 8 GiB. Int32 overflow # aliases initialized low rows, so a regression fails numerically, not by OOB. - cache = torch.empty(high + count, num_kv_heads, head_dim, device=device, dtype=torch.bfloat16) - cache[low : low + count].fill_(-2) - cache[high : high + count].normal_(mean=3, std=2) + dtype = torch.uint8 if kv_quant == "fp8" else torch.bfloat16 + cache = torch.empty(high + count, num_kv_heads, head_dim, device=device, dtype=dtype) + scales = None + if kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import quantize_kv_to_cache + + scales = torch.empty(high + count, num_kv_heads, device=device, dtype=torch.float32) + cache[low : low + count].fill_(0) + values = torch.randn(count, stride, device=device, dtype=torch.bfloat16) * 2 + 3 + slots = torch.arange(high, high + count, device=device, dtype=torch.int32) + quantize_kv_to_cache(values, values, slots, cache, cache, scales, scales) + else: + cache[low : low + count].fill_(-2) + cache[high : high + count].normal_(mean=3, std=2) if offset_bits == 31: # Keep the old negative offset inside allocated guard storage too. guard = (1 << 31) // stride cache = cache[guard:] + if scales is not None: + scales = scales[guard:] high -= guard - q = torch.randn(1, num_q_heads, head_dim, device=device, dtype=cache.dtype) + q = torch.randn(1, num_q_heads, head_dim, device=device, dtype=torch.bfloat16) indices = torch.arange(high, high + count, device=device, dtype=torch.int32) indptr = torch.tensor([0, count], device=device, dtype=torch.int32) q_to_req = torch.zeros(1, device=device, dtype=torch.int32) q_positions = torch.tensor([count - 1], device=device, dtype=torch.int64) sm_scale = head_dim**-0.5 + selected = cache[indices.long()] + if scales is not None: + selected = selected.view(torch.float8_e4m3fn).float() * scales[indices.long(), :, None] expected = _reference_paged_attention( - q, cache, cache, indptr, indices, q_to_req, q_positions, sm_scale, None, + q, selected, selected, indptr, torch.arange(count, device=device), q_to_req, q_positions, sm_scale, None, ) + scale_args = {"k_scale": scales, "v_scale": scales} if path == "paged": - actual = paged_attention(q, cache, cache, indptr, indices, q_to_req, q_positions, sm_scale) + actual = paged_attention(q, cache, cache, indptr, indices, q_to_req, q_positions, sm_scale, **scale_args) elif path == "decode": mid = torch.empty(1, num_q_heads, 64, head_dim, device=device, dtype=torch.float32) lse = torch.empty(1, num_q_heads, 64, device=device, dtype=torch.float32) - actual = decode_paged_attention(q, cache, cache, indptr, indices, q_positions, mid, lse, sm_scale) + actual = decode_paged_attention(q, cache, cache, indptr, indices, q_positions, mid, lse, sm_scale, **scale_args) else: qo_indptr = torch.tensor([0, 1], device=device, dtype=torch.int32) prefix_lens = torch.tensor([count - 1], device=device, dtype=torch.int32) # 65 cached tokens cover both the full-tile and masked split-prefix loads. - extend = cache[high + count - 1 : high + count] if path == "split_extend" else None + extend = selected[-1:].to(q.dtype) if path == "split_extend" else None actual = extend_paged_attention( q, cache, cache, qo_indptr, indptr, indices, prefix_lens, 1, sm_scale, - k_extend=extend, v_extend=extend, + k_extend=extend, v_extend=extend, **scale_args, ) torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) @@ -116,6 +134,12 @@ def test_triton_backend_passes_attention_sinks_to_paged_kernel(monkeypatch): from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata class FakeKVCache: + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + def __init__(self): self.device = torch.device("cpu") self.k = torch.zeros(4, 1, 4) @@ -775,6 +799,12 @@ def test_triton_backend_stores_kv_and_matches_reference(monkeypatch): from freetoken.attention.triton import TritonAttentionBackend class FakeKVCache: + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + def __init__(self, device: torch.device, head_dim: int): self.device = device self.dtype = torch.float32 @@ -845,6 +875,12 @@ def test_triton_backend_applies_the_batch_block_ends_only_when_the_spec_asks(mon from freetoken.attention.triton import TritonAttentionBackend class FakeKVCache: + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + def __init__(self, device: torch.device, head_dim: int): self.device = device self.dtype = torch.bfloat16 @@ -895,6 +931,12 @@ def test_triton_backend_replay_metadata_uses_capture_buffers(monkeypatch): from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata class FakeKVCache: + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + def __init__(self, device: torch.device): self.device = device diff --git a/tests/kernels/test_triton_attention_fp8.py b/tests/kernels/test_triton_attention_fp8.py new file mode 100644 index 000000000..fc698fcd6 --- /dev/null +++ b/tests/kernels/test_triton_attention_fp8.py @@ -0,0 +1,355 @@ +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.attention import ( + decode_paged_attention, + extend_paged_attention, + paged_attention, +) +from freetoken.kernel.triton.kv_quant import quantize_kv_to_cache + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU required") + + +def _data(tokens, heads=4, dim=128): + torch.manual_seed(73) + k = torch.randn(tokens, heads * dim * 3, device="cuda", dtype=torch.bfloat16)[:, heads * dim:2 * heads * dim] + v = torch.randn(tokens, heads * dim, device="cuda", dtype=torch.bfloat16) * 0.13 + slots = torch.randperm(tokens + 64, device="cuda")[:tokens].to(torch.int32) + kc = torch.zeros(tokens + 64, heads, dim, device="cuda", dtype=torch.uint8) + vc = torch.zeros_like(kc) + ks = torch.zeros(tokens + 64, heads, device="cuda", dtype=torch.float32) + vs = torch.zeros_like(ks) + quantize_kv_to_cache(k, v, slots, kc, vc, ks, vs) + return k.view(tokens, heads, dim), v.view(tokens, heads, dim), slots, kc, vc, ks, vs + + +def _read(codes, scales, slots): + return codes.view(torch.float8_e4m3fn).float()[slots.long()] * scales[slots.long(), :, None] + + +def _oracle(q, k, v, prefix): + group = q.shape[1] // k.shape[1] + k = k.repeat_interleave(group, dim=1).transpose(0, 1) + v = v.repeat_interleave(group, dim=1).transpose(0, 1) + score = torch.einsum("thd,hnd->htn", q.float(), k.float()) * q.shape[-1]**-0.5 + mask = torch.arange(k.shape[1], device=q.device)[None, :] <= prefix + torch.arange(q.shape[0], device=q.device)[:, None] + return torch.einsum("htn,hnd->thd", score.masked_fill(~mask[None], -torch.inf).softmax(-1), v.float()) + + +@pytest.mark.parametrize("tokens,window,padding,offset", [ + (17, None, 0, 0), (512, None, 0, 0), (8193, None, 0, 0), + (32768, None, 0, 0), (513, 137, 0, 0), + (65, None, 4, 4), (65, None, 0, 1), (65, 17, 1, 0), +]) +@pytest.mark.parametrize("splits", [1, 8, 64]) +@pytest.mark.parametrize("q_heads,kv_heads,dim", [(16, 8, 128), (32, 8, 128), (32, 4, 128), (16, 2, 256)]) +def test_fp8_decode_keeps_adaptive_split_range(tokens, window, padding, offset, splits, q_heads, kv_heads, dim): + _, _, slots, kc, vc, ks, vs = _data(tokens, heads=kv_heads, dim=dim) + if padding or offset: + # Independent K/V layouts catch byte-vs-word addressing and view-offset mistakes. + k_storage = torch.empty((tokens + 64) * kv_heads * (dim + padding) + offset, + device="cuda", dtype=torch.uint8) + k_view = k_storage[offset:].view(tokens + 64, kv_heads, dim + padding)[..., :dim] + k_view.copy_(kc) + kc = k_view + v_storage = torch.empty(tokens + 64, kv_heads, dim + 4, device="cuda", dtype=torch.uint8) + v_view = v_storage[..., :dim] + v_view.copy_(vc) + vc = v_view + q = torch.randn(1, q_heads, dim, device="cuda", dtype=torch.bfloat16) + out = decode_paged_attention( + q, kc, vc, torch.tensor([0, tokens], device="cuda", dtype=torch.int32), slots, + torch.tensor([tokens-1], device="cuda", dtype=torch.int32), + torch.empty(1, q_heads, 64, dim, device="cuda"), torch.empty(1, q_heads, 64, device="cuda"), + dim**-0.5, max_kv_splits=splits, k_scale=ks, v_scale=vs, sliding_window=window, + ) + visible_slots = slots[-window:] if window else slots + reference = _oracle(q, _read(kc, ks, visible_slots), _read(vc, vs, visible_slots), len(visible_slots)-1) + torch.testing.assert_close(out.float(), reference, atol=0.002, rtol=0.02) + + +@pytest.mark.parametrize("transpose", [False, True]) +def test_fp8_rdna4_packed_decode_preserves_every_finite_code(transpose): + import triton + import triton.language as tl + + from freetoken.kernel.triton.attention import _gfx1201, _kv_load_gfx12_bf16 + + if not _gfx1201(torch.cuda.current_device()): + pytest.skip("RDNA4 packed KV conversion requires gfx1201") + + @triton.jit + def decode_codes(x, y, TRANSPOSE: tl.constexpr): + rows = tl.arange(0, 64) + words = tl.arange(0, 32) + if TRANSPOSE: + ptrs = x + rows[None, :] * 32 + words[:, None] + else: + ptrs = x + rows[:, None] * 32 + words[None, :] + values = _kv_load_gfx12_bf16(ptrs, tl.full(ptrs.shape, True, tl.int1), TRANSPOSE) + if TRANSPOSE: + values = tl.trans(values) + tl.store(y + rows[:, None] * 128 + tl.arange(0, 128)[None, :], values) + + codes = torch.arange(256, device="cuda", dtype=torch.uint8).repeat(32) + actual = torch.empty(codes.numel(), device="cuda", dtype=torch.bfloat16) + decode_codes[(1,)](codes.view(torch.uint32), actual, transpose, num_warps=2) + expected = codes.view(torch.float8_e4m3fn).bfloat16() + finite = expected.isfinite() + # Integer comparison also checks negative zero and byte order within each word. + assert torch.equal(actual.view(torch.int16)[finite], expected.view(torch.int16)[finite]) + + +@pytest.mark.parametrize("path", ["paged", "extend", "split-extend"]) +def test_fp8_prefill_reads_scales_and_keeps_live_chunk(path): + prefix, query = 65, 129 + k, v, slots, kc, vc, ks, vs = _data(prefix + query, heads=2) + q = torch.randn(query, 8, 128, device="cuda", dtype=torch.bfloat16) + qptr = torch.tensor([0, query], device="cuda", dtype=torch.int32) + kptr = torch.tensor([0, prefix+query], device="cuda", dtype=torch.int32) + prefix_t = torch.tensor([prefix], device="cuda", dtype=torch.int32) + k_ref, v_ref = _read(kc, ks, slots), _read(vc, vs, slots) + if path == "paged": + out = paged_attention(q, kc, vc, kptr, slots, + torch.zeros(query, device="cuda", dtype=torch.int32), + prefix+torch.arange(query, device="cuda", dtype=torch.int32), + 128**-0.5, k_scale=ks, v_scale=vs) + else: + live = path == "split-extend" + out = extend_paged_attention(q, kc, vc, qptr, kptr, slots, prefix_t, query, 128**-0.5, + k_extend=k[prefix:] if live else None, + v_extend=v[prefix:] if live else None, + k_scale=ks, v_scale=vs) + if live: + k_ref[prefix:], v_ref[prefix:] = k[prefix:].float(), v[prefix:].float() + torch.testing.assert_close(out.float(), _oracle(q, k_ref, v_ref, prefix), atol=0.003, rtol=0.03) + + +def test_fp8_multi_request_prefill_reads_independent_row_scales(): + prefixes = (65, 37) + queries = (4, 7) + lengths = tuple(p + n for p, n in zip(prefixes, queries)) + k, v, slots, kc, vc, ks, vs = _data(sum(lengths), heads=2) + q = torch.randn(sum(queries), 8, 128, device="cuda", dtype=torch.bfloat16) + qptr = torch.tensor([0, queries[0], sum(queries)], device="cuda", dtype=torch.int32) + kptr = torch.tensor([0, lengths[0], sum(lengths)], device="cuda", dtype=torch.int32) + prefix_t = torch.tensor(prefixes, device="cuda", dtype=torch.int32) + live_k = torch.cat((k[prefixes[0]:lengths[0]], k[lengths[0]+prefixes[1]:])) + live_v = torch.cat((v[prefixes[0]:lengths[0]], v[lengths[0]+prefixes[1]:])) + out = extend_paged_attention(q, kc, vc, qptr, kptr, slots, prefix_t, max(queries), + 128**-0.5, k_extend=live_k, v_extend=live_v, + k_scale=ks, v_scale=vs) + decoded_k = _read(kc, ks, slots) + decoded_v = _read(vc, vs, slots) + refs = [] + q_start = k_start = live_start = 0 + for prefix, query, length in zip(prefixes, queries, lengths): + ref_k = decoded_k[k_start:k_start+length].clone() + ref_v = decoded_v[k_start:k_start+length].clone() + ref_k[prefix:] = live_k[live_start:live_start+query] + ref_v[prefix:] = live_v[live_start:live_start+query] + refs.append(_oracle(q[q_start:q_start+query], ref_k, ref_v, prefix)) + q_start += query + k_start += length + live_start += query + torch.testing.assert_close(out.float(), torch.cat(refs), atol=0.003, rtol=0.03) + + +def test_fp8_multi_request_decode_keeps_rows_separate(): + lengths = (65, 37) + _, _, slots, kc, vc, ks, vs = _data(sum(lengths), heads=2) + q = torch.randn(2, 8, 128, device="cuda", dtype=torch.bfloat16) + kptr = torch.tensor([0, lengths[0], sum(lengths)], device="cuda", dtype=torch.int32) + positions = torch.tensor([n - 1 for n in lengths], device="cuda", dtype=torch.int32) + mid = torch.empty(2, 8, 64, 128, device="cuda") + lse = torch.empty(2, 8, 64, device="cuda") + out = decode_paged_attention(q, kc, vc, kptr, slots, positions, mid, lse, + 128**-0.5, k_scale=ks, v_scale=vs) + decoded_k = _read(kc, ks, slots) + decoded_v = _read(vc, vs, slots) + refs = [_oracle(q[:1], decoded_k[:lengths[0]], decoded_v[:lengths[0]], lengths[0] - 1), + _oracle(q[1:], decoded_k[lengths[0]:], decoded_v[lengths[0]:], lengths[1] - 1)] + torch.testing.assert_close(out.float(), torch.cat(refs), atol=0.003, rtol=0.03) + + +@pytest.mark.parametrize("q_heads,kv_heads,dim", [(8, 2, 128), (16, 2, 256)]) +def test_fp8_graph_replays_changed_inputs_and_slot_mapping(q_heads, kv_heads, dim): + k, v, slots, kc, vc, ks, vs = _data(8, heads=kv_heads, dim=dim) + k, v = k.reshape(8, -1).contiguous(), v.reshape(8, -1).contiguous() + q = torch.randn(1, q_heads, dim, device="cuda", dtype=torch.bfloat16) + kptr = torch.tensor([0, 8], device="cuda", dtype=torch.int32) + pos = torch.tensor([7], device="cuda", dtype=torch.int32) + mid = torch.empty(1, q_heads, 64, dim, device="cuda") + lse = torch.empty(1, q_heads, 64, device="cuda") + out = torch.empty_like(q) + + def forward(): + quantize_kv_to_cache(k, v, slots, kc, vc, ks, vs) + return decode_paged_attention(q, kc, vc, kptr, slots, pos, mid, lse, dim**-0.5, + out=out, k_scale=ks, v_scale=vs) + + forward() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + forward() + for step in range(5): + k.normal_(std=step+1) + v.normal_(std=0.1*(step+1)) + q.normal_() + slots.copy_(torch.randperm(72, device="cuda")[:8].int()) + graph.replay() + reference = _oracle(q, _read(kc, ks, slots), _read(vc, vs, slots), 7) + torch.testing.assert_close(out.float(), reference, atol=0.005, rtol=0.03) + + +def test_fp8_first_prefill_preserves_bf16_attention_and_populates_codes(): + from types import SimpleNamespace + + from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.kvcache.mha_pool import MHAKVCache + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + torch.manual_seed(104) + tokens, heads, dim = 129, 2, 128 + q = torch.randn(tokens, 8, dim, device='cuda', dtype=torch.bfloat16) + k = torch.randn(tokens, heads * dim, device='cuda', dtype=torch.bfloat16) + v = torch.randn(tokens, heads * dim, device='cuda', dtype=torch.bfloat16) + slots = torch.randperm(tokens + 16, device='cuda')[:tokens].int() + indptr = torch.tensor([0, tokens], device='cuda', dtype=torch.int32) + metadata = TritonMetadata( + cu_seqlens_q_gpu=indptr, indptr=indptr, indices=slots, + q_to_req=torch.zeros(tokens, device='cuda', dtype=torch.int32), + q_positions=torch.arange(tokens, device='cuda', dtype=torch.int32), + is_decode=False, prefix_lens=torch.zeros(1, device='cuda', dtype=torch.int32), + max_q_len=tokens, has_prefix=False, + ) + batch = SimpleNamespace(attn_metadata=metadata, out_loc=slots, mm_block_ends=None) + outputs = [] + for mode in ('none', 'fp8'): + pool = MHAKVCache(heads, 1, dim, tokens+16, 1, torch.bfloat16, torch.device('cuda'), kv_quant=mode) + backend = object.__new__(TritonAttentionBackend) + backend.kvcache = pool + backend.prefill_tile_min_q = 128 + outputs.append(backend.forward(q, k, v, 0, batch)) + if mode == 'fp8': + assert pool.k_cache(0).dtype == torch.uint8 + assert (pool.k_scale(0)[slots.long()] > 0).all() + assert (pool.v_scale(0)[slots.long()] > 0).all() + assert torch.equal(outputs[0], outputs[1]) + + +@pytest.mark.parametrize("page_size", [1, 8]) +@pytest.mark.parametrize("query", [4, 129]) +@pytest.mark.parametrize("q_heads,dim", [(8, 128), (16, 256)]) +def test_fp8_workspace_reuse_rebuild_and_graph(page_size, query, q_heads, dim): + from types import SimpleNamespace + from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.kvcache.mha_pool import MHAKVCache + + if try_get_tp_info() is None: + set_tp_info(0, 1) + torch.manual_seed(2203) + prefix, heads = 65, 2 + pool = MHAKVCache(heads, 2, dim, 256 // page_size, page_size, + torch.bfloat16, torch.device("cuda"), kv_quant="fp8") + backend = object.__new__(TritonAttentionBackend) + backend.kvcache, backend.device = pool, pool.device + backend.prefill_tile_min_q = 128 + backend.prefill_workspace = None + backend.init_workspace() + original_pointer = backend.prefill_workspace.data_ptr() + + for pages in (256 // page_size, 512 // page_size, 224 // page_size): + if pool._kv_buffer.shape[2] != pages: + torch.cuda.synchronize() + backend.reset_workspace() + assert backend.prefill_workspace is None + pool.rebuild(pages) + backend.init_workspace() + capacity = pages * page_size + assert backend.prefill_workspace.shape == (2, capacity, heads, dim) + pointer = backend.prefill_workspace.data_ptr() + q = torch.randn(query, q_heads, dim, device="cuda", dtype=torch.bfloat16) + k = torch.randn(prefix + query, heads * dim, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) * .1 + slots = torch.randperm(capacity, device="cuda")[:prefix + query].int() + metadata = TritonMetadata( + cu_seqlens_q_gpu=torch.tensor([0, query], device="cuda", dtype=torch.int32), + indptr=torch.tensor([0, prefix + query], device="cuda", dtype=torch.int32), + indices=slots, q_to_req=torch.zeros(query, device="cuda", dtype=torch.int32), + q_positions=torch.arange(prefix, prefix + query, device="cuda", dtype=torch.int32), + is_decode=False, prefix_lens=torch.tensor([prefix], device="cuda", dtype=torch.int32), + max_q_len=query, has_prefix=True, + ) + batch = SimpleNamespace(attn_metadata=metadata, out_loc=slots[prefix:], mm_block_ends=None) + + def forward(): + pool.store_kv(k[:prefix], v[:prefix], slots[:prefix], 0) + return backend.forward(q, k[prefix:], v[prefix:], 0, batch) + + forward() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = forward() + for step in range(3): + k.normal_(std=step + 1) + v.normal_(std=.1 * (step + 1)) + q.normal_() + slots.copy_(torch.randperm(capacity, device="cuda")[:prefix + query].int()) + graph.replay() + kc = pool.k_cache(0).view(capacity, heads, dim) + vc = pool.v_cache(0).view(capacity, heads, dim) + kr = torch.cat([_read(kc, pool.k_scale(0), slots[:prefix]).bfloat16(), k[prefix:].view(query, heads, dim)]) + vr = torch.cat([_read(vc, pool.v_scale(0), slots[:prefix]).bfloat16(), v[prefix:].view(query, heads, dim)]) + torch.testing.assert_close(out.float(), _oracle(q, kr, vr, prefix), atol=.003, rtol=.03) + # A different layer overwrites the same workspace; the next replay must refresh it. + pool.store_kv(k[:prefix] * 2, v[:prefix] * 3, slots[:prefix], 1) + backend.forward(q, k[prefix:], v[prefix:], 1, batch) + assert backend.prefill_workspace.data_ptr() == pointer + torch.cuda.synchronize() + del graph, out + if capacity == 256: + assert pointer == original_pointer + + +@pytest.mark.parametrize("path", ["paged", "decode", "extend", "split-extend"]) +def test_fp8_attention_large_scale_pitch(path): + # An int32 scale offset wraps to initialized row zero, yielding a numerical + # failure rather than an invalid access. Code offsets themselves fit int32. + torch.cuda.empty_cache() + if torch.cuda.mem_get_info()[0] < 18 * 2**30: + pytest.skip("Large scale-pitch regression needs 18 GiB of free GPU memory") + torch.manual_seed(1347) + heads, dim, count = 2, 128, 66 + scales = torch.empty_strided((65, heads), (1 << 26, 1), device="cuda", dtype=torch.float32) + scales[0].fill_(0.01) + scales[64].copy_(torch.tensor([0.25, 0.5], device="cuda")) + codes = torch.zeros(65, heads, dim, device="cuda", dtype=torch.uint8) + codes[64].copy_(torch.randn(heads, dim, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8)) + slots = torch.full((count,), 64, device="cuda", dtype=torch.int32) + indptr = torch.tensor([0, count], device="cuda", dtype=torch.int32) + pos = torch.tensor([count-1], device="cuda", dtype=torch.int32) + q = torch.randn(1, 8, dim, device="cuda", dtype=torch.bfloat16) + selected = codes[64].view(torch.float8_e4m3fn).float() * scales[64, :, None] + reference = selected.repeat_interleave(4, dim=0)[None] + kwargs = {"k_scale": scales, "v_scale": scales} + if path == "paged": + out = paged_attention(q, codes, codes, indptr, slots, torch.zeros_like(pos), pos, dim**-0.5, **kwargs) + elif path == "decode": + out = decode_paged_attention(q, codes, codes, indptr, slots, pos, + torch.empty(1, 8, 64, dim, device="cuda"), + torch.empty(1, 8, 64, device="cuda"), dim**-0.5, **kwargs) + else: + live = selected[None].to(q.dtype) if path == "split-extend" else None + out = extend_paged_attention(q, codes, codes, torch.tensor([0, 1], device="cuda", dtype=torch.int32), + indptr, slots, pos, 1, dim**-0.5, + k_extend=live, v_extend=live, **kwargs) + torch.testing.assert_close(out.float(), reference, atol=0.003, rtol=0.03) diff --git a/tests/kvcache/test_hybrid_linear_paged_pools.py b/tests/kvcache/test_hybrid_linear_paged_pools.py index 8b60801cb..25eb4b377 100644 --- a/tests/kvcache/test_hybrid_linear_paged_pools.py +++ b/tests/kvcache/test_hybrid_linear_paged_pools.py @@ -92,6 +92,53 @@ def test_factory_kpool1_builds_plain_dsa_pool(): assert type(pool) is DSAKVCache +@pytest.mark.parametrize("kv_quant", ["none", "fp8"]) +def test_full_group_quantization_preserves_hybrid_layer_mapping_and_cost(kv_quant): + from dataclasses import replace + from types import SimpleNamespace + + from freetoken.kvcache.mha_pool import MHAKVCache + + cfg = _glm5_like_config() + linear, full = cfg.attention_groups + cfg = replace( + cfg, num_qo_heads=16, num_kv_heads=2, head_dim=256, + attention_groups=( + replace(linear, variant="gdn", output_gate="silu"), + FullAttentionGroupConfig(name="full", layer_ids=full.layer_ids, + num_kv_heads=2, head_dim=256, rotary_config=cfg.rotary_config), + ), + ) + pool = create_kvcache_pool(cfg, 7, 8, torch.bfloat16, torch.device("cpu"), kv_quant=kv_quant) + config = SimpleNamespace(model_config=cfg, tp_info=SimpleNamespace(size=1), + dtype=torch.bfloat16, page_size=8, kv_quant=kv_quant) + # Only the three full layers own KV; the other nine have no code or scale slab. + bytes_per_token = 3 * 2 * 2 * (256 + 4 if kv_quant == "fp8" else 256 * 2) + assert MHAKVCache.kv_cost(config)[0] == bytes_per_token * 8 + for pages in (7, 11, 3): + if pages != 7: + pool.rebuild(pages) + assert pool.unit_bytes() == (bytes_per_token, 0) + assert pool._kv_buffer.numel() * pool._kv_buffer.element_size() == ( + 3 * 2 * 2 * 256 * pages * 8 * (1 if kv_quant == "fp8" else 2)) + for lid in full.layer_ids: + pool.k_cache(lid).fill_(lid) + pool.v_cache(lid).fill_(lid + 1) + if kv_quant == "fp8": + pool.k_scale(lid).fill_(lid / 10) + pool.v_scale(lid).fill_((lid + 1) / 10) + for lid in full.layer_ids: + assert torch.all(pool.k_cache(lid) == lid) + assert torch.all(pool.v_cache(lid) == lid + 1) + if kv_quant == "fp8": + assert pool.k_scale(lid).shape == (pages * 8, 2) + torch.testing.assert_close(pool.k_scale(lid), torch.full((pages * 8, 2), lid / 10)) + torch.testing.assert_close(pool.v_scale(lid), torch.full((pages * 8, 2), (lid + 1) / 10)) + for lid in linear.layer_ids: + with pytest.raises(KeyError): + pool.k_cache(lid) + + def test_cost_model_kpool_shadow_slab_quarter_cost(): """The shadow slab stores one row per index_ratio tokens: the kpool spec's index bytes are 1/ratio of the plain DSA slab (rings/scratch are per-request @@ -114,4 +161,4 @@ def test_cost_model_kpool_shadow_slab_quarter_cost(): assert ( spec_kv_bytes_per_token(spec1, econf) - spec_kv_bytes_per_token(spec, econf) == index_full - index_full // 4 - ) \ No newline at end of file + ) diff --git a/tests/kvcache/test_mha_pool_fp8.py b/tests/kvcache/test_mha_pool_fp8.py new file mode 100644 index 000000000..956ab2771 --- /dev/null +++ b/tests/kvcache/test_mha_pool_fp8.py @@ -0,0 +1,212 @@ +"""The fp8 KV pool: code buffer + per-(token, head) scales, sized and rebuilt together. + +The pool-side half of --kv-cache-dtype fp8. The interesting failures are the silent +ones: a scale buffer that misses the layer_ids remap, a rebuild that resizes the codes +but not the scales, or a ``unit_bytes`` that drifts from ``kv_cost`` (which is what the +VRAM budget, the cache sliders and the runtime rebuild all divide by). +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.models.config import KVCacheGroupSpec + +if not torch.cuda.is_available(): # pragma: no cover + pytest.skip("CUDA required", allow_module_level=True) + +from freetoken.kernel.triton.kv_quant import kv_codes_dtype + +DEV = torch.device("cuda") +HEADS, DIM, LAYERS, PAGES, PAGE_SIZE = 4, 64, 3, 6, 8 + + +def _init_tp() -> None: + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + + +@pytest.fixture(autouse=True) +def _tp(): + _init_tp() + + +def _pool(kv_quant="fp8", num_pages=PAGES, layer_ids=None, num_layers=LAYERS): + from freetoken.kvcache.mha_pool import MHAKVCache + + _init_tp() + return MHAKVCache( + num_kv_heads=HEADS, + num_layers=num_layers, + head_dim=DIM, + num_pages=num_pages, + page_size=PAGE_SIZE, + dtype=torch.bfloat16, + device=DEV, + layer_ids=layer_ids, + kv_quant=kv_quant, + ) + + +def test_kv_store_dtype_selection(): + from freetoken.kvcache.mha_pool import _kv_store_dtype + + assert _kv_store_dtype(torch.bfloat16, "none") is torch.bfloat16 + assert _kv_store_dtype(torch.bfloat16, "fp8") is kv_codes_dtype() + with pytest.raises(ValueError, match="kv_quant"): + _kv_store_dtype(torch.bfloat16, "q6") + + +def test_fp8_pool_keeps_geometry_and_adds_scale_views(): + pool = _pool() + assert pool.kv_quant == "fp8" + # Same shape as the 16-bit pool -- only the element type changed. + assert pool._kv_buffer.shape == (2, LAYERS, PAGES, PAGE_SIZE, HEADS, DIM) + assert pool._kv_buffer.dtype == kv_codes_dtype() + assert pool.store_dtype == kv_codes_dtype() + # ``dtype`` stays the COMPUTE dtype (kvcache/base.py). Backends size their scratch + # with it -- reporting codes here hands e4m3 to a 16-bit tl.dot, which does not fail + # until CUDA-graph capture (exactly how QSA's indexer died in the field). + assert pool.dtype == torch.bfloat16 + slots = PAGES * PAGE_SIZE + for layer in range(LAYERS): + assert pool.k_scale(layer).shape == (slots, HEADS) + assert pool.v_scale(layer).shape == (slots, HEADS) + assert pool.k_scale(layer).dtype == torch.float32 + # A 16-bit pool exposes no scales at all. + assert _pool(kv_quant="none").k_scale(0) is None + + +def test_layer_ids_remap_applies_to_scales_too(): + layer_ids = (1, 3) + pool = _pool(layer_ids=layer_ids, num_layers=4) + assert pool._kv_buffer.shape[1] == 2 + rows = torch.full((1, HEADS * DIM), 2.0, device=DEV, dtype=torch.bfloat16) + pool.store_kv(rows, rows * 3, torch.tensor([9], device=DEV, dtype=torch.int32), 3) + torch.testing.assert_close(pool.k_scale(3)[9], torch.full((HEADS,), 2 / 448, device=DEV)) + torch.testing.assert_close(pool.v_scale(3)[9], torch.full((HEADS,), 6 / 448, device=DEV)) + assert not pool.k_scale(1).any() + assert not pool.v_scale(1).any() + with pytest.raises(KeyError): + pool.k_scale(0) + + +def test_store_kv_scatters_codes_and_scales(): + from freetoken.kernel.triton.kv_quant import codes_to_f32 + + tokens = 5 + torch.manual_seed(7) + rows = torch.randn(tokens, HEADS * DIM, device=DEV, dtype=torch.bfloat16) * 2.0 + out_loc = torch.tensor( + [3, 0, PAGES * PAGE_SIZE - 1, 40, 17], device=DEV, dtype=torch.int32 + ) + pool = _pool() + pool.store_kv(rows, rows.clone(), out_loc, layer_id=0) + torch.cuda.synchronize() + + codes = codes_to_f32(pool.k_cache(0).view(-1, HEADS, DIM)) + scale = pool.k_scale(0) + f32 = rows.view(tokens, HEADS, DIM).to(torch.float32) + deq = codes[out_loc.long()] * scale[out_loc.long()].unsqueeze(-1) + amax = f32.abs().amax(dim=-1, keepdim=True) + # The scales are per (token, head), so the bound is relative to each row's max. + err = ((deq - f32).abs() / amax.clamp_min(1e-6)).max() + assert float(err) < 0.08, float(err) + # Rows nobody wrote must stay exactly zero -- the scatter touched only out_loc. + untouched = torch.ones(PAGES * PAGE_SIZE, dtype=torch.bool, device=DEV) + untouched[out_loc.long()] = False + assert (codes[untouched] == 0).all() + assert (scale[untouched] == 0).all() + + +def test_rebuild_resizes_codes_and_scales_together(): + pool = _pool() + before = id(pool) + pool.rebuild(11) + assert id(pool) == before # identity preserved (backends cache the object) + assert pool._kv_buffer.shape == (2, LAYERS, 11, PAGE_SIZE, HEADS, DIM) + slots = 11 * PAGE_SIZE + assert pool.k_scale(0).shape == (slots, HEADS) + assert pool.k_cache(0).shape[0] == 11 + # Per-token cost is page-count invariant, codes and scales alike. + assert pool.unit_bytes() == _pool().unit_bytes() + + +def _sizing_config(kv_quant): + from freetoken.attention import AttnType + + spec = KVCacheGroupSpec( + name="full", + layer_ids=tuple(range(LAYERS)), + num_kv_heads=HEADS, + head_dim=DIM, + sliding_window=None, + attn_type=AttnType.FULL, + ) + mc = SimpleNamespace( + has_swa_attention=False, + has_linear_attention=False, + num_layers=LAYERS, + num_kv_heads=HEADS, + head_dim=DIM, + kv_cache_group_specs=lambda: (spec,), + ) + return SimpleNamespace( + model_config=mc, + page_size=PAGE_SIZE, + dtype=torch.bfloat16, + tp_info=SimpleNamespace(size=1), + kv_quant=kv_quant, + ) + + +@pytest.mark.parametrize("kv_quant", ["none", "fp8"]) +def test_unit_bytes_matches_the_cost_model_that_sized_the_pool(kv_quant): + """The budget solve and the live pool must agree byte for byte.""" + from freetoken.kvcache.base import spec_kv_bytes_per_token + from freetoken.kvcache.mha_pool import MHAKVCache + + config = _sizing_config(kv_quant) + (spec,) = config.model_config.kv_cache_group_specs() + per_token = spec_kv_bytes_per_token(spec, config) + assert MHAKVCache.kv_cost(config)[0] == per_token * PAGE_SIZE + + pool = _pool(kv_quant=kv_quant) + assert pool.unit_bytes() == (per_token, 0) + + +def test_fp8_lands_just_above_half_the_bytes(): + """Half the code bytes, plus a scale sidecar of 4 B per (token, slab, layer, head).""" + plain, quantized = _pool("none").unit_bytes()[0], _pool("fp8").unit_bytes()[0] + scales = 2 * LAYERS * HEADS * 4 + assert quantized == plain // 2 + scales, (plain, quantized, scales) + + +def test_fp8_runtime_budget_includes_workspace_and_dummy(): + from freetoken.kvcache.mha_pool import MHAKVCache + from freetoken.kvcache.base import CacheRebuildRejected + + config = _sizing_config("fp8") + config.num_page_override = None + config.memory_ratio = 1.0 + persistent = MHAKVCache.kv_cost(config)[0] + workspace = PAGE_SIZE * 2 * HEADS * DIM * 2 + per_page, fixed, _, _ = MHAKVCache.runtime_cost(config) + assert per_page == persistent + workspace + assert fixed == per_page + budget = (PAGES + 1) * per_page + assert MHAKVCache.solve_num_pages(config, budget) == PAGES + config.num_page_override = PAGES + 1 + with pytest.raises(ValueError, match="budget"): + MHAKVCache.solve_num_pages(config, budget) + pool = _pool() + before = pool._kv_buffer.data_ptr() + with pytest.raises(CacheRebuildRejected): + pool.validate_rebuild(config, num_pages=PAGES + 1, target_moe=0, per_expert_bytes=1, + baseline_free=budget, weights_bytes=0, current_num_pages=PAGES) + assert pool._kv_buffer.data_ptr() == before + assert MHAKVCache.runtime_cost(_sizing_config("none")) == MHAKVCache.kv_cost(_sizing_config("none"))