diff --git a/docs/cli.md b/docs/cli.md index ff4af382d..32a6b6f3e 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -71,8 +71,32 @@ ft serve --model ... --gpu GPU-9e8d7c6b # the same card by UUID (a unique prefi | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) | | `--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` | bf16 | `bf16` or `fp8`: store the KV cache as e4m3 codes plus one fp32 scale per (token, kv head), roughly doubling the tokens that fit in the same VRAM; see [FP8 KV cache](#fp8-kv-cache) | | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU | +### FP8 KV cache + +`ft serve --kv-cache-dtype fp8` halves the bytes per cached token (8-bit codes instead +of 16), so a card that held N tokens holds close to 2N. Each `(token, kv head)` row +keeps its own fp32 scale, which costs ~3% back at `head_dim=128`. Requirements and +trade-offs: + +- Needs the **triton** attention backend; `--attn auto` selects it (and refuses an + explicit `fi`/`fa`/`trtllm`, which cannot be shown to apply these scales). +- Works on the plain paged, hybrid-SWA and QSA sparse (Qwen3.8-Flash-Next) KV pools. + On QSA the block-selection index keys stay 16-bit; only the selected K/V rows are + read back as codes. MLA/DSA latent KV, DeepSeek-V4's tiered pool and the block-sparse + MiniMax-M3 pool stay 16-bit; asking for fp8 there fails at startup rather than + silently ignoring the flag. +- The same bytes on every GPU FreeToken targets: the codes sit in a plain byte buffer + and are decoded in software, so the cache holds identical data and produces identical + numbers on any card (the fp8 type is deliberately kept out of the kernels, which is + also what makes the feature work on the RTX 30 series). +- Accuracy is checkpoint-dependent. Expect it to matter most on long contexts and on + models with outlier key channels; keep `bf16` when a run must be bit-reproducible. +- `ft ctl stats` / `/v1/cache/status` report the smaller `kv_bytes_per_token`, and + `ft ctl cache --kv N` moves the same (now cheaper) pool. + ### MoE offload See [models.md](models.md#moe-backends) for what each backend does. diff --git a/docs/models.md b/docs/models.md index c9499163f..e41375bfa 100644 --- a/docs/models.md +++ b/docs/models.md @@ -41,3 +41,9 @@ for them; other checkpoints of the same architectures work too. authoritative model args are read from there. - Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM. - Multimodal checkpoints are served text-only. +- `--kv-cache-dtype fp8` (see [cli.md](cli.md#fp8-kv-cache)) covers the plain paged, + hybrid-SWA and QSA sparse KV pools — gpt-oss, Qwen3/3.5/3.6, GLM-4.x, Gemma-4, + MiniMax-M2.5, Muse-Glimmer, Llama/Qwen2/Mistral, Qwen3.8-Flash-Next (on QSA only the + selected K/V rows are read back as codes; block selection keeps 16-bit index keys). + MLA/DSA (GLM-5.2), DeepSeek-V4's tiered pool and MiniMax-M3's block-sparse pool stay + 16-bit and reject it. diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 8a410ffe4..012759ad3 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -33,6 +33,11 @@ class BackendInfo: # Whether forward() honors a per-call AttentionSpec (window/sm_scale/sinks). # Non-consumers raise on a non-None spec instead of silently dropping it. consumes_attn_spec: bool = False + # Whether forward() reads an fp8 KV pool (codes + per-token/per-head scales). + # Backends that hand the cache to an external kernel must opt out until that + # kernel is proven to apply our scale layout; the engine then refuses (or auto- + # avoids) them for --kv-cache-dtype fp8. + supports_fp8_kv: bool = False SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend") @@ -84,6 +89,7 @@ def create_fa_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.FULL, AttnType.SWA}), consumes_attn_spec=True, + supports_fp8_kv=True, ), ) def create_triton_backend(config: ModelConfig): @@ -137,6 +143,9 @@ def create_m3_sparse_backend(config: ModelConfig): "qsa_sparse", BackendInfo( supported_types=frozenset({AttnType.QSA}), + # The attend kernel dequantizes on load (kernel/triton/qsa/attend.py); the + # compressed index keys it scores against are a separate, always-16-bit tier. + supports_fp8_kv=True, # 64-token pages: a 4-token compress group never straddles a page, so the # compressed row of a group is page_base // 4 + block-in-page. page_sizes=(64,), diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py index 4a28dc852..68153b1d4 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -108,7 +108,17 @@ def __init__(self, config: ModelConfig) -> None: f"qsa_sparse backend needs a QSA pool, got {type(self.kvcache).__name__}" ) self.device = self.kvcache.device + # The pool's COMPUTE dtype, never its store dtype (the contract lives in + # kvcache/base.py). These buffers feed the indexer -- qsa_index_norm_rope and + # qsa_mqa_paged -- whose tl.dot has no fp8 path, so an e4m3 q_index does not + # fail here, it fails at CUDA-graph capture with "Unsupported rhs dtype + # fp8e4nv". --kv-cache-dtype fp8 quantizes only the KV tiers; the index tiers + # stay 16-bit by design (kvcache/qsa_pool.py). self.dtype = self.kvcache.dtype + assert self.dtype.itemsize == 2, ( + f"QSA block selection needs a 16-bit compute dtype, got {self.dtype} -- " + "the KV pool must report its compute dtype, not e4m3 codes" + ) self.index_head_dim = self.kvcache.index_head_dim self.ratio = self.kvcache.index_ratio self.ring_capacity = self.kvcache.ring_capacity @@ -282,6 +292,8 @@ def qsa_forward( self._update_index_cache(index, md, slot) indices = self._select(index, md, slot) + # Scale tensors only exist on an fp8 pool (k_scale returns None otherwise); the + # index tier stays bf16 either way, so _select above is quantization-agnostic. return qsa_sparse_paged_attention( q, self.kvcache.k_cache(layer_id), @@ -290,6 +302,8 @@ def qsa_forward( md.block_table, md.token_to_req, torch.empty_like(q), + k_scale=self.kvcache.k_scale(layer_id), + v_scale=self.kvcache.v_scale(layer_id), ) def _plan_index_writes(self, md: QSASparseMetadata, batch: Batch) -> None: diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 9eed1e1d2..c731c2636 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -155,6 +155,11 @@ def forward( assert head_dim == q.shape[-1] k_cache = k_raw.view(-1, kv_heads, head_dim) v_cache = v_raw.view(-1, kv_heads, head_dim) + # An fp8 KV pool hands us its per-(token, head) scales; a 16-bit pool returns + # None and every kernel below keeps its original (scale-free) code path. + k_scale = self.kvcache.k_scale(layer_id) + v_scale = self.kvcache.v_scale(layer_id) + assert (k_scale is None) == (v_scale is None), "K and V scales come as a pair" spec = attn_spec or AttentionSpec() indices = metadata.indices @@ -181,6 +186,8 @@ def forward( sm_scale=scale, sliding_window=spec.sliding_window, sinks=spec.sinks, + k_scale=k_scale, + v_scale=v_scale, ) if ( (not metadata.is_decode) @@ -201,6 +208,8 @@ def forward( sinks=spec.sinks, k_extend=k.view(q.shape[0], kv_heads, head_dim), v_extend=v.view(q.shape[0], kv_heads, head_dim), + k_scale=k_scale, + v_scale=v_scale, ) return paged_attention( q=q, @@ -213,6 +222,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: diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index bcbe6bcf2..08d833ea4 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -59,6 +59,12 @@ class EngineConfig: cuda_graph_bs: List[int] | None = None cuda_graph_max_bs: int | None = None page_size: int = 1 + # KV-cache storage quantization: "none" stores the compute dtype, "fp8" stores e4m3 + # codes plus one fp32 scale per (token, slab, layer, kv head) -- about 2x the tokens + # per GiB, at a small accuracy cost. --kv-cache-dtype; resolved from "auto" by + # _adjust_config, which also refuses it on a pool family or attention backend that + # cannot read the scales. + kv_quant: str = "none" memory_ratio: float = 0.9 # Hybrid GDN models default to the HybridRadixCache (cross-request GDN-state prefix reuse); # `--cache-type naive` opts out. linear_state_cache_ratio sizes the GDN snapshot cache as diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index b5a6fa3b0..cc87d95bf 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -114,10 +114,38 @@ def _backend_requirements_met(name: str) -> bool: return True -def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str: +# --kv-cache-dtype spellings -> the stored EngineConfig.kv_quant value. +KV_QUANT_ALIASES = {"auto": "none", "bf16": "none", "none": "none", "fp8": "fp8"} + + +def _resolve_kv_quant(value: str | None) -> str: + """Normalize a --kv-cache-dtype spelling to EngineConfig.kv_quant.""" + key = (value or "auto").strip().lower() + if key not in KV_QUANT_ALIASES: + raise ValueError( + f"unknown --kv-cache-dtype {value!r}; expected one of " + f"{', '.join(sorted(KV_QUANT_ALIASES))}" + ) + return KV_QUANT_ALIASES[key] + + +def _backend_supports_kv_quant(name: str, kv_quant: str) -> bool: + """Whether every comma part of an attention-backend string can read a quantized + KV pool (an unquantized pool needs nothing from the backend).""" + if kv_quant == "none": + return True + return all( + attention_backend_info(part.strip()).supports_fp8_kv for part in name.split(",") + ) + + +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: + whose packages are installed, whose every comma part serves ALL required + types, and which can decode a quantized KV cache when one is configured. + Reproduces the historical hardware tree for FULL-only models: sm_100 -> trtllm, sm_90+sgl_kernel -> "fa,fi", flashinfer -> fi, else triton.""" candidates: list[tuple[str, bool]] = [] if AttnType.DSV4 in required: @@ -144,10 +172,18 @@ def _resolve_auto_attention_backend(required: frozenset[AttnType]) -> str: continue if not _backend_requirements_met(name): continue + if not _backend_supports_kv_quant(name, kv_quant): + continue return name raise RuntimeError( "No attention backend can serve attention types " - f"{sorted(t.value for t in required)} on this machine." + f"{sorted(t.value for t in required)} on this machine" + + ( + f" with a {kv_quant} KV cache" + if kv_quant != "none" + else "" + ) + + "." ) @@ -192,6 +228,23 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att f"SWA models require, got {config.attention_backend!r}." ) + # A quantized KV pool is only readable by a backend that applies its per-(token, + # head) scales; one that hands the cache to an external kernel would silently + # attend to raw e4m3 codes. Rejected here, before any weight is resident. + kv_quant = getattr(config, "kv_quant", "none") + if not _backend_supports_kv_quant(config.attention_backend, kv_quant): + fp8_backends = [ + name + for name in ("trtllm", "fi", "fa", "triton") + if required <= attention_backend_info(name).supported_types + and attention_backend_info(name).supports_fp8_kv + ] + raise ValueError( + f"--kv-cache-dtype {kv_quant} needs an attention backend that decodes the KV " + f"scales; {config.attention_backend!r} does not. Valid for this model: " + f"{', '.join(fp8_backends) or 'none'} (or use --kv-cache-dtype bf16)." + ) + # An explicitly-selected backend may require a package that isn't installed. Auto # never resolves to one of these when its package is missing, so this only fires for # explicit --attention-backend choices. @@ -1301,6 +1354,27 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # lists, then validate whatever is now selected (explicit or auto) -- every # comma part must serve every required type, with packages/arch available. required_attn_types = _required_attn_types(model_config) + # Resolve KV quantization BEFORE the backend tree: a quantized pool narrows both + # which pool families are usable and which backend auto may pick. + kv_quant = _resolve_kv_quant(getattr(config, "kv_quant", "none")) + override("kv_quant", kv_quant) + if kv_quant != "none": + # fp8 codes are wired through the pools that hand their rows to a Triton + # kernel: the plain paged and hybrid-SWA ones, plus the QSA sparse pool, whose + # index tier stays bf16 -- only the selected tokens come back as codes. + # Everything else (MLA's absorbed cache, DSA/DSV4/BSA sparse) has kernels that + # assert on 16-bit rows, and kvcache/__init__.py rejects fp8 for those families + # at pool creation. + quant_unsupported = required_attn_types - { + AttnType.FULL, AttnType.SWA, AttnType.QSA, + } + if quant_unsupported: + raise ValueError( + f"--kv-cache-dtype {kv_quant} is implemented for the plain paged, " + "hybrid-SWA and QSA sparse KV pools; this model also needs " + f"{', '.join(sorted(t.value for t in quant_unsupported))} attention " + "(use --kv-cache-dtype bf16)." + ) _dtype = getattr(config, "dtype", None) # duck-typed test configs omit it if ( required_attn_types & {AttnType.BSA, AttnType.QSA} @@ -1328,7 +1402,7 @@ def override(attr: str, value: Any): # this is dangerous, use with caution 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 0ce249de8..664b8043f 100644 --- a/python/freetoken/kernel/aot_models.py +++ b/python/freetoken/kernel/aot_models.py @@ -11,9 +11,11 @@ (a drifted derivation misses the prebuilt cache by spec name and falls back to JIT, which needs nvcc): -- store: ``element_size = num_kv_heads * head_dim * 2`` (bf16 KV row), one per - paged-KV attention group (kvcache/mha_pool.py, kvcache/hybrid_swa_pool.py). - DSV4 writes its MLA latent via torch scatter and contributes nothing. +- store: ``element_size = num_kv_heads * head_dim * dtype_bytes``, one per + paged-KV attention group (kvcache/mha_pool.py, kvcache/hybrid_swa_pool.py) and + per KV width: 2 for the 16-bit cache, 1 for an fp8 one (``--kv-cache-dtype + fp8``, kvcache/mha_pool.py). DSV4 writes its MLA latent via torch scatter and + contributes nothing. - index: ``element_size = hidden_size * 2`` (bf16 embedding row) paired with the runtime ``num_splits_for`` rule (layers/embedding.py -> kernel/index.py). DSV4 (plain nn.Embedding) and GGUF embeddings (GGUFEmbedding) bypass it. @@ -33,7 +35,8 @@ from .index import num_splits_for -KV_CACHE_DTYPE_BYTES = 2 # every current model allocates bf16 paged KV +KV_CACHE_DTYPE_BYTES = 2 # the default paged KV is the 16-bit compute dtype +FP8_KV_CACHE_DTYPE_BYTES = 1 # --kv-cache-dtype fp8 stores one e4m3 code per element EMBED_DTYPE_BYTES = 2 # embedding weights stay bf16 on the indexing() path @@ -391,8 +394,9 @@ def expert_bank_row_bytes(fmt: str, hidden_size: int, moe_intermediate_size: int ) -def store_element_sizes(model: AotModel) -> set[int]: - return {kv * hd * KV_CACHE_DTYPE_BYTES for kv, hd in model.kv_groups} +def store_element_sizes(model: AotModel, dtype_bytes: int = KV_CACHE_DTYPE_BYTES) -> set[int]: + """Store-kernel row sizes for one model's paged-KV groups at a given bytes/elem.""" + return {kv * hd * dtype_bytes for kv, hd in model.kv_groups} def index_variants(model: AotModel) -> set[tuple[int, int]]: @@ -414,9 +418,17 @@ def fast_index_copy_feature_sizes(model: AotModel) -> set[int]: def aggregate_store_element_sizes() -> tuple[int, ...]: + """Every store row size the runtime can ask for. + + Both KV widths ship: the 16-bit default and the fp8 (``--kv-cache-dtype fp8``) + code buffer, whose rows are exactly half as wide. A missing size is not a + correctness bug -- it is a kernel-cache miss that falls back to JIT and fails + the ``FREETOKEN_DISABLE_JIT=1`` release gate. + """ sizes: set[int] = set() for model in SUPPORTED_MODELS: - sizes.update(store_element_sizes(model)) + for dtype_bytes in (KV_CACHE_DTYPE_BYTES, FP8_KV_CACHE_DTYPE_BYTES): + sizes.update(store_element_sizes(model, dtype_bytes)) return tuple(sorted(sizes)) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index c2358d84f..31ae84847 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,6 +6,8 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 as _kv_load_f32 + _MAX_KV_SPLITS = 8 _MIN_BLOCK_KV = 32 @@ -47,6 +49,8 @@ def _paged_attention_kernel( q_ptr, k_ptr, v_ptr, + k_scale_ptr, + v_scale_ptr, o_ptr, indptr_ptr, indices_ptr, @@ -60,6 +64,8 @@ def _paged_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -68,6 +74,7 @@ def _paged_attention_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, ): q_tok = tl.program_id(0) q_head = tl.program_id(1) @@ -107,14 +114,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) - 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. Index math stays int32 like the 16-bit path below. + 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")) @@ -124,14 +147,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 @@ -149,6 +186,8 @@ def _decode_grouped_stage1_kernel( q_ptr, k_ptr, v_ptr, + k_scale_ptr, + v_scale_ptr, sm_scale, indptr_ptr, indices_ptr, @@ -162,6 +201,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, @@ -179,6 +220,7 @@ def _decode_grouped_stage1_kernel( D: tl.constexpr, DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -224,7 +266,10 @@ 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: + # A 16-bit cache feeds tl.dot as-is; an fp8 cache is decoded up to q's own + # compute dtype below, so q must NOT be narrowed to the (1-byte) cache type. + 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) @@ -232,19 +277,39 @@ 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) - k = tl.load( - k_ptr + slots[None, :] * stride_ks + k_base_offsets, - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, - ) + if HAS_KV_SCALE: + s_k = tl.load( + k_scale_ptr + slots * stride_kss + kv_head, mask=mask_n, other=0.0 + ) + k = _kv_load_f32( + k_ptr + slots[None, :] * stride_ks + k_base_offsets, + mask_n[None, :] & mask_d[:, None], + ) + k = (k * s_k[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 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 HAS_KV_SCALE: + s_v = tl.load( + v_scale_ptr + slots * stride_vss + kv_head, mask=mask_n, other=0.0 + ) + v = _kv_load_f32( + v_ptr + slots[:, None] * stride_vs + v_base_offsets, + mask_n[:, None] & mask_dv[None, :], + ) + v = (v * s_v[:, 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) @@ -364,10 +429,18 @@ def decode_paged_attention( sliding_window: int | None = None, sinks: torch.Tensor | None = None, out: torch.Tensor | 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.""" + """SGLang-style split-k grouped decode attention for one query per request. + + ``k_scale`` / ``v_scale`` (``[num_slots, num_kv_heads]`` fp32) turn the cache into + an fp8 KV cache: every code row is multiplied by its own token/head scale. Both + must be given together; ``None`` keeps the 16-bit path byte-identical. + """ 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 batch, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] @@ -391,6 +464,15 @@ def decode_paged_attention( 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 + # Unused pointer args still need a real tensor (same convention as sinks_arg). + 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), + ) # valid_block_h = heads computed per program (drives the grid + head indexing); block_h = # power-of-two tile size for tl.arange. They differ only for non-power-of-two GQA groups # (e.g. 6), where block_h rounds up and the kernel masks the extra lanes. @@ -405,6 +487,8 @@ def decode_paged_attention( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, sm_scale, indptr, indices, @@ -418,6 +502,8 @@ 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), @@ -435,6 +521,7 @@ def decode_paged_attention( D=head_dim, DV=head_dim, SLIDING_WINDOW=sliding_window or 0, + HAS_KV_SCALE=has_kv_scale, num_warps=4, num_stages=2, ) @@ -471,6 +558,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, @@ -484,6 +573,8 @@ def _extend_attention_kernel( stride_kh, stride_vs, stride_vh, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -494,6 +585,7 @@ def _extend_attention_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -545,14 +637,27 @@ 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) - 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 = tl.load( + k_scale_ptr + slots * stride_kss + kv_head, mask=mask_n, other=0.0 + ) + k = _kv_load_f32( + k_ptr + + slots[None, :] * stride_ks + + kv_head * stride_kh + + offs_d[:, None], + mask_n[None, :] & mask_d[:, None], + ) + k = (k * s_k[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 scores = tl.where(final_mask, scores, -float("inf")) @@ -562,14 +667,27 @@ 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, - ) + if HAS_KV_SCALE: + s_v = tl.load( + v_scale_ptr + slots * stride_vss + kv_head, mask=mask_n, other=0.0 + ) + v = _kv_load_f32( + v_ptr + + slots[:, None] * stride_vs + + kv_head * stride_vh + + offs_dv[None, :], + mask_n[:, None] & mask_dv[None, :], + ) + v = (v * s_v[:, 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, + ) acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -592,6 +710,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, @@ -609,6 +729,8 @@ def _extend_attention_split_kernel( stride_kch, stride_vcs, stride_vch, + stride_kss, + stride_vss, stride_ot, stride_oh, GROUP: tl.constexpr, @@ -619,6 +741,7 @@ def _extend_attention_split_kernel( BLOCK_N: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + HAS_KV_SCALE: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -672,14 +795,27 @@ def _extend_attention_split_kernel( if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) - k = tl.load( - k_cache_ptr - + slots[None, :] * stride_kcs - + kv_head * stride_kch - + offs_d[:, None], - mask=mask_n[None, :] & mask_d[:, None], - other=0.0, - ) + if HAS_KV_SCALE: + s_k = tl.load( + k_scale_ptr + slots * stride_kss + kv_head, mask=mask_n, other=0.0 + ) + k = _kv_load_f32( + k_cache_ptr + + slots[None, :] * stride_kcs + + kv_head * stride_kch + + offs_d[:, None], + mask_n[None, :] & mask_d[:, None], + ) + k = (k * s_k[None, :]).to(q.dtype) + else: + k = tl.load( + k_cache_ptr + + slots[None, :] * stride_kcs + + kv_head * stride_kch + + offs_d[:, None], + mask=mask_n[None, :] & mask_d[:, None], + other=0.0, + ) scores = tl.dot(q.to(k.dtype), k) * sm_scale scores = tl.where(final_mask, scores, -float("inf")) @@ -689,14 +825,27 @@ def _extend_attention_split_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - v = tl.load( - v_cache_ptr - + slots[:, None] * stride_vcs - + kv_head * stride_vch - + offs_dv[None, :], - mask=mask_n[:, None] & mask_dv[None, :], - other=0.0, - ) + if HAS_KV_SCALE: + s_v = tl.load( + v_scale_ptr + slots * stride_vss + kv_head, mask=mask_n, other=0.0 + ) + v = _kv_load_f32( + v_cache_ptr + + slots[:, None] * stride_vcs + + kv_head * stride_vch + + offs_dv[None, :], + mask_n[:, None] & mask_dv[None, :], + ) + v = (v * s_v[:, None]).to(q.dtype) + else: + v = tl.load( + v_cache_ptr + + slots[:, None] * stride_vcs + + kv_head * stride_vch + + offs_dv[None, :], + mask=mask_n[:, None] & mask_dv[None, :], + other=0.0, + ) acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -773,10 +922,18 @@ def extend_paged_attention( out: torch.Tensor | None = None, k_extend: torch.Tensor | None = None, v_extend: 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-tiled causal prefill/extend attention over paged KV cache. + + ``k_scale`` / ``v_scale`` mark an fp8 KV cache (see ``decode_paged_attention``). + The ``k_extend`` / ``v_extend`` rows are the current request's own K/V and stay in + the compute dtype either way, so only the cached prefix is decoded. + """ 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_q_tokens, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] @@ -793,6 +950,15 @@ def extend_paged_attention( 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 + # Unused pointer args still need a real tensor (same convention as sinks_arg). + 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) block_dv = triton.next_power_of_2(head_dim) # Tile size is shared-memory bound: keep the fast (large) tiles on GPUs whose opt-in @@ -815,6 +981,8 @@ def extend_paged_attention( v_extend, k_cache, v_cache, + k_scale_arg, + v_scale_arg, o, qo_indptr, kv_indptr, @@ -832,6 +1000,8 @@ 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), GROUP=num_q_heads // num_kv_heads, @@ -842,6 +1012,7 @@ def extend_paged_attention( 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, num_stages=1, ) @@ -851,6 +1022,8 @@ def extend_paged_attention( q, k_cache, v_cache, + k_scale_arg, + v_scale_arg, o, qo_indptr, kv_indptr, @@ -864,6 +1037,8 @@ 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), GROUP=num_q_heads // num_kv_heads, @@ -874,6 +1049,7 @@ def extend_paged_attention( 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, num_stages=1, ) @@ -893,15 +1069,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] @@ -916,12 +1096,22 @@ def paged_attention( 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, @@ -935,6 +1125,8 @@ 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), GROUP=num_q_heads // num_kv_heads, @@ -943,6 +1135,7 @@ def paged_attention( 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 1d9f744ce..532c1d7ad 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -5,8 +5,18 @@ Affected kernels branch on :func:`e4m3_native_cx` (a compile-time constexpr): the native branch stays byte-identical on sm_89+, the emulated branch is dead-code eliminated there. When the emulated branch is active, wrappers must pass e4m3 -tensors as ``.view(torch.uint8)`` and allocate act-quant outputs as bf16 -- use -the host-side twin :func:`e4m3_native` for those decisions. +tensors as ``.view(torch.uint8)`` and allocate act-quant outputs as bf16. + +TWO independent probes answer "is fp8e4nv native here": :func:`e4m3_native` (host, +torch's device capability -- decides what a buffer is ALLOCATED as) and +:func:`e4m3_native_cx` (compile-time, triton's target -- decides which arm a kernel +compiles). They cannot be merged, because a constexpr function that referenced the +host one would not survive triton's cache-key AST walk, so on a box where the probes +disagree the host holds real fp8 tensors while kernels take the emulated arm. +:func:`warn_if_probes_disagree` says so once at startup. Code handed a buffer whose +dtype the host already settled -- the KV pool -- must not re-ask at all and picks its +decode from the pointer: :func:`kv_load_e4m3_tile_f32`. Trusting the probe over the +tensor is what produced "cannot cast int32 to fp8e4nv" at CUDA graph capture. ``FREETOKEN_FORCE_E4M3_EMU=1`` (or true/yes/on) forces the emulated path on any GPU (for A/B validation against the native fp8 unit). The flag is read ONCE at @@ -47,8 +57,9 @@ def _env_force() -> bool: def e4m3_native() -> bool: - """Host-side twin of :func:`e4m3_native_cx`: True when kernels take fp8e4nv - tensors directly. False: pass ``.view(torch.uint8)`` and bf16 act buffers.""" + """Host-side probe: does THIS device take fp8e4nv tensors directly? True: kernels + get fp8 tensors, False: pass ``.view(torch.uint8)`` and bf16 act buffers. NOT + necessarily the answer :func:`e4m3_native_cx` gives -- see this module's header.""" global _native if _env_force() != FORCE_EMU: raise RuntimeError( @@ -64,9 +75,47 @@ def e4m3_native() -> bool: # one process runs on one GPU, so its convention is that GPU's; None (-> the current device) only before the process binds _native = torch.cuda.get_device_capability(assigned_visible_gpu()) >= (8, 9) + warn_if_probes_disagree() return _native +_warned_disagree = False + + +def warn_if_probes_disagree() -> None: + """Log once when the two native-fp8e4nv probes answer differently. + + :func:`e4m3_native` decides what buffers the host ALLOCATES while + :func:`e4m3_native_cx` decides which arm a kernel compiles, and the two cannot be + unified (that function's docstring explains triton's cache-key walk). A box where + triton's probe under-reports therefore runs every e4m3 kernel through the software + decode -- bit-exact per this module's header, but slower -- and any kernel that + trusts the probe over the tensor it was handed stops compiling outright. Reads the + latched ``_native`` rather than calling back into :func:`e4m3_native`.""" + global _warned_disagree + if _warned_disagree: + return + _warned_disagree = True + if FORCE_EMU: # emulating by request is not a disagreement + return + try: + triton_native = target_info.cuda_capability_geq(8, 9) + if triton_native == bool(_native): + return + major, minor = torch.cuda.get_device_capability() + except Exception: # noqa: BLE001 -- no driver/no target yet: nothing to compare + return + from freetoken.utils import init_logger + + init_logger(__name__).warning( + "native fp8e4nv disagreement: torch reports sm_%d%d for this device but " + "triton's target probe says %s, so e4m3 kernels compile the software-decode " + "branch (bit-exact, slower). The fp8 KV cache is unaffected -- it follows the " + "buffer it was given.", + major, minor, "supported" if triton_native else "unsupported", + ) + + def e4m3_kernel_view(t: torch.Tensor) -> torch.Tensor: """An e4m3 tensor as the branched kernels expect it: unchanged when native, the uint8 view otherwise (the fp8 pointer type is illegal pre-sm_89).""" @@ -84,7 +133,17 @@ def e4m3_native_cx(): """Compile-time: does the compilation target have native fp8e4nv (sm_89+)? Delegates to ``target_info`` (reads the active driver's target, so cross-compilation tests that patch ``driver.active.get_current_target`` - resolve consistently).""" + resolve consistently). + + It CANNOT defer to :func:`e4m3_native`, however much one verdict per process is + what we want: triton hashes a constexpr function by walking its AST + (runtime/jit.py: cache_key -> record_reference), and a bare reference to a plain + python function raises "Unsupported function referenced: " + -- trying that once disabled every e4m3 kernel at once, PLE gather included. + Module attributes (``target_info.whatever``) survive the walk, plain functions do + not. The probes therefore stay separate, :func:`warn_if_probes_disagree` reports + when they disagree, and code handed a buffer the host already typed -- the KV + pool -- ignores this function and follows the pointer: kv_load_e4m3_tile_f32.""" return not FORCE_EMU and target_info.cuda_capability_geq(8, 9) @@ -122,3 +181,54 @@ def round_e4m3(x): y_norm = ((b + 524287 + lsb) & 0xFFF00000).to(tl.float32, bitcast=True) y_sub = (x + 24576.0) - 24576.0 return tl.where(tl.abs(x) >= 0.015625, y_norm, y_sub) + + +@jit +def e4m3_f32_to_u8(x): + """Encode an fp32 value that ALREADY lies on the e4m3 grid -- the output of + :func:`round_e4m3`, clamped to +-448 -- into its e4m3 byte code. This is the + encoder the pre-sm_89 emulated path needs to STORE fp8-sized data (the fp8 + type itself is unavailable there, so the bytes live in a uint8 buffer that + :func:`e4m3_u8_to_f32` decodes back). + + Normal range: read the (unbiased) exponent and the now-zero-padded fp32 + mantissa back out of the fp32 header. Subnormal range (|x| < 2^-6, grid step + 2^-9): the value is an exact multiple of 2^-9, so ``|x| * 512`` IS the mantissa + field -- the sign bit has to be carried in by hand, since that branch never + looks at the header. The same 0.015625 boundary as :func:`round_e4m3` keeps the + two consistent: ``e4m3_u8_to_f32(e4m3_f32_to_u8(round_e4m3(x)))`` is x's + single-rounded value for every input, and no code it emits is a NaN pattern (the + caller's +-448 clamp caps the code at 0x7E). ``-0.0`` encodes as 0x00 after + round_e4m3 (which documents returning +0.0 for it). + """ + u = x.to(tl.uint32, bitcast=True) + sign = ((u >> 31) & 1).to(tl.int32) + exp = ((u >> 23) & 0xFF).to(tl.int32) - 127 + mant = ((u >> 20) & 7).to(tl.int32) + normal = (sign << 7) | ((exp + 7) << 3) | mant + sub = (sign << 7) | (tl.abs(x) * 512.0).to(tl.int32) + return tl.where(tl.abs(x) >= 0.015625, normal, sub).to(tl.uint8) + + +@jit +def kv_load_e4m3_tile_f32(ptrs, mask): + """Load a tile of KV e4m3 codes and widen it to fp32. + + Straight-line on purpose: no probe, no dtype test, so there is no arm left to + prune. The pools keep their codes in a plain byte buffer on EVERY architecture + (kv_quant.kv_codes_dtype), so the fp8e4nv type never reaches Triton through here. + Both ways of choosing an arm were tried on real hardware and each broke the run: + the compile-time fp8-native answer is a second, independent verdict that can + disagree with the host that allocated the buffer, and a comparison against the + pointer's element type is NOT statically pruned -- Triton type-checks the arm that + should have been dead, and an int mask fill against an fp8 pointer is rejected + ("cannot cast int32 to fp8e4nv", raised at CUDA graph capture on sm_100). + + What remains is the decode that already runs wherever the fp8 type is unavailable, + bit-exact per this module's header: the same load-with-int-fill and software + widening as kernel/triton/ple.py and nvfp4_linear.py. Callers use this only in + their quantized branch -- the 16-bit path keeps its own tl.load, so bf16 attention + is untouched instruction for instruction -- and the dense paged kernels and the + QSA sparse one read the same pool, hence this helper lives here. + """ + return e4m3_u8_to_f32(tl.load(ptrs, mask=mask, other=0)) diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py new file mode 100644 index 000000000..e041e0466 --- /dev/null +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -0,0 +1,197 @@ +"""FP8 (e4m3) KV-cache storage: per-token, per-head symmetric quantization. + +One KV row is one ``(token, kv_head)`` slice of ``head_dim`` elements. It is stored +as ``head_dim`` e4m3 bytes plus ONE fp32 scale shared by the whole row: + + scale = max(amax(row) / 448, eps) # 448 == e4m3 finite max + code = round_e4m3(clamp(row / scale)) # RNE onto the e4m3 grid + read = code.to(f32) * scale # in the attention kernels + +Granularity rationale: an fp32 scale per (token, head) costs ``4 / head_dim`` bytes +per element (3% at head_dim 128, 6% at 64) while tracking each key's own magnitude, +which is what keeps a quantized KV from collapsing on outlier heads. A coarser +per-tensor scale needs no storage at all but has no headroom for them; a finer +per-element "scale" is the format itself. + +Architectures below sm_89 have no fp8e4nv type in Triton (see +:mod:`freetoken.kernel.triton.e4m3_compat`), so the codes live in a plain ``uint8`` +buffer on EVERY architecture and are decoded by :func:`e4m3_u8_to_f32`. That keeps one +set of bytes and one set of numbers across GPUs, and it is why :func:`kv_codes_dtype` +is a constant rather than a question: the fp8 type never appears in a kernel +signature, so nothing here can disagree with the host that allocated the buffer. +Choosing the encode/decode per target -- by an arch probe, or by testing the pointer's +element type -- is what broke this feature twice on real hardware (see +:func:`kv_load_e4m3_tile_f32`). + +The write path replaces ``kernel.store_cache`` for a quantized pool: the plain store +kernel is a raw byte copy that requires the source and the cache to share a dtype, +and quantization is exactly the step where the two diverge. Folding the scatter into +the quantization kernel keeps that to a single launch (and a single HBM round trip) +under CUDA-graph capture, where the slot ids arrive as a device tensor. +""" + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from freetoken.kernel.triton.e4m3_compat import ( + e4m3_f32_to_u8, + round_e4m3, +) + +FP8 = torch.float8_e4m3fn +KV_SCALE_DTYPE = torch.float32 +KV_QUANT_FP8 = "fp8" + + +def kv_codes_dtype() -> torch.dtype: + """Storage dtype of one quantized KV element: e4m3 bytes in a uint8 buffer. + + A constant, deliberately. Every attempt to answer this per target -- triton's + compile-time probe, then testing the pointer's element type at the load -- ended up + picking an arm that did not match the buffer the host had just allocated (see the + module header). torch still reads these bytes as fp8 whenever real numbers are + wanted: :func:`codes_to_f32`. + """ + return torch.uint8 + + +def alloc_codes(shape: tuple[int, ...], device: torch.device) -> torch.Tensor: + """A zero-filled code buffer of :func:`kv_codes_dtype` -- bytes, on every arch. + + Zero-filling matters because the pools read slots that were never written: a stale + byte decodes to a real number (0x7F/0xFF even to NaN once reinterpreted as fp8), + while 0x00 is exactly 0.0 (same reasoning as kvcache/bsa_pool.py). + """ + return torch.zeros(shape, dtype=kv_codes_dtype(), device=device) + + +def codes_to_f32(codes: torch.Tensor) -> torch.Tensor: + """Decode a code buffer to fp32 ON THE HOST (torch's own e4m3 cast). + + Works on either storage dtype -- uint8 bytes are reinterpreted as fp8 first -- so + a test or a debugging tool reads the same numbers on sm_86 as on sm_90, and does + it through torch rather than through the software decoder it is checking. + """ + 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) + + # Straight-line, like the reader: round onto the e4m3 grid in ONE step (RNE) and + # pack the bits into the byte buffer. No fp8 type on either side -- the reason this + # feature stopped compiling twice is explained in kv_codes_dtype's docstring. + 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 == KV_SCALE_DTYPE, k_scale.dtype + _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, + ) + + +__all__ = [ + "FP8", + "KV_QUANT_FP8", + "KV_SCALE_DTYPE", + "alloc_codes", + "codes_to_f32", + "kv_codes_dtype", + "quantize_kv_to_cache", +] + diff --git a/python/freetoken/kernel/triton/qsa/attend.py b/python/freetoken/kernel/triton/qsa/attend.py index 541e27c68..7e5815487 100644 --- a/python/freetoken/kernel/triton/qsa/attend.py +++ b/python/freetoken/kernel/triton/qsa/attend.py @@ -9,12 +9,16 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 + @triton.jit def _qsa_sparse_paged_gqa_splitk_kernel( q_ptr, k_cache_ptr, v_cache_ptr, + k_scale_ptr, + v_scale_ptr, indices_ptr, block_table_ptr, token_to_req_ptr, @@ -29,6 +33,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( stride_v_block, stride_v_token, stride_v_head, + stride_kss, + stride_vss, stride_indices_row, stride_table_req, stride_output_row, @@ -46,6 +52,10 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + # e4m3 KV pool (kvcache/mha_pool.py): read codes + per-token row scales instead of + # 16-bit values. The bf16 branch below stays exactly as it was, instruction for + # instruction, for the unquantized default. + HAS_KV_SCALE: tl.constexpr, ) -> None: # row * stride can overflow int32 for large row counts. row = tl.program_id(0).to(tl.int64) @@ -101,24 +111,66 @@ def _qsa_sparse_paged_gqa_splitk_kernel( valid &= (physical_page >= 0) & (physical_page < num_cache_blocks) # physical_page * block stride can overflow int32 for large caches. safe_page = tl.maximum(physical_page, 0).to(tl.int64) - keys = tl.load( - k_cache_ptr - + safe_page[None, :] * stride_k_block - + page_offset[None, :] * stride_k_token - + kv_head * stride_k_head - + dim_offsets[:, None], - mask=valid[None, :], - other=0.0, - ) - values = tl.load( - v_cache_ptr - + safe_page[:, None] * stride_v_block - + page_offset[:, None] * stride_v_token - + kv_head * stride_v_head - + dim_offsets[None, :], - mask=valid[:, None], - other=0.0, - ) + if HAS_KV_SCALE: + # The scale row is the slot the code lives in: QSA pins page_size to this + # kernel's PAGE_SIZE (attention/__init__.py registers page_sizes=(64,)), so + # slot = page * PAGE_SIZE + offset addresses k_scale/v_scale exactly. + scale_slot = safe_page * PAGE_SIZE + page_offset # safe_page is int64 + # Invalid columns mask the codes to 0.0 and the scale to 1.0, and slots that + # were never written read back 0.0 * 0.0 -- both buffers are zero-filled. + # Either way the operand stays finite, so the -inf row mask below is what + # decides such a column's fate rather than a NaN poisoning the row. + s_k = tl.load( + k_scale_ptr + scale_slot[None, :] * stride_kss + kv_head, + mask=valid[None, :], + other=1.0, + ) + s_v = tl.load( + v_scale_ptr + scale_slot[:, None] * stride_vss + kv_head, + mask=valid[:, None], + other=1.0, + ) + keys = ( + kv_load_e4m3_tile_f32( + k_cache_ptr + + safe_page[None, :] * stride_k_block + + page_offset[None, :] * stride_k_token + + kv_head * stride_k_head + + dim_offsets[:, None], + valid[None, :], + ) + * s_k + ).to(query.dtype) + values = ( + kv_load_e4m3_tile_f32( + v_cache_ptr + + safe_page[:, None] * stride_v_block + + page_offset[:, None] * stride_v_token + + kv_head * stride_v_head + + dim_offsets[None, :], + valid[:, None], + ) + * s_v + ).to(query.dtype) + else: + keys = tl.load( + k_cache_ptr + + safe_page[None, :] * stride_k_block + + page_offset[None, :] * stride_k_token + + kv_head * stride_k_head + + dim_offsets[:, None], + mask=valid[None, :], + other=0.0, + ) + values = tl.load( + v_cache_ptr + + safe_page[:, None] * stride_v_block + + page_offset[:, None] * stride_v_token + + kv_head * stride_v_head + + dim_offsets[None, :], + mask=valid[:, None], + other=0.0, + ) scores = tl.dot(query, keys) # Scaling scores avoids re-quantizing a scaled query to BF16. scores *= softmax_scale_log2 @@ -232,8 +284,10 @@ def qsa_sparse_paged_attention( block_table: torch.Tensor, token_to_req: torch.Tensor, out: torch.Tensor | None = None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, ) -> torch.Tensor: - """Run sparse GQA directly over paged BF16 K/V caches.""" + """Run sparse GQA directly over paged K/V caches (bf16, or e4m3 + row scales).""" if q.ndim != 3 or k_cache.ndim != 4 or v_cache.shape != k_cache.shape: raise ValueError("QSA sparse attention received invalid Q/K/V shapes") @@ -247,7 +301,23 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse attention requires valid grouped-query heads") head_dim = q.shape[2] assert head_dim >= 16 and (head_dim & (head_dim - 1)) == 0 - assert q.dtype == k_cache.dtype == v_cache.dtype + if (k_scale is None) != (v_scale is None): + raise ValueError("QSA sparse attention requires both KV scale tensors") + if k_scale is not None: + # The pool hands out 1-byte e4m3 codes plus one fp32 row scale per + # (slot, kv_head); the kernel rebuilds that slot as + # page * PAGE_SIZE + page_offset, which is exact only because QSA pins + # page_size to PAGE_SIZE (page_sizes=(64,) in attention/__init__.py). + if k_cache.element_size() != 1 or v_cache.element_size() != 1: + raise ValueError("QSA KV scales require 1-byte e4m3 code caches") + if k_scale.dtype is not torch.float32 or v_scale.dtype is not torch.float32: + raise ValueError("QSA KV scale tensors must be float32") + want = (k_cache.shape[0] * k_cache.shape[1], k_cache.shape[2]) + if k_scale.shape != want or v_scale.shape != want: + raise ValueError(f"QSA KV scale tensors must have shape {want}") + assert k_scale.stride(1) == v_scale.stride(1) == 1 + else: + assert q.dtype == k_cache.dtype == v_cache.dtype assert logical_indices.dtype == block_table.dtype == torch.int32 assert token_to_req.dtype == torch.int32 assert q.stride(2) == k_cache.stride(3) == v_cache.stride(3) == 1 @@ -303,6 +373,10 @@ def qsa_sparse_paged_attention( q, k_cache, v_cache, + # Never dereferenced while HAS_KV_SCALE is False -- pass the caches so the + # launch stays type-valid without a second None-handling path. + k_cache if k_scale is None else k_scale, + v_cache if v_scale is None else v_scale, logical_indices, block_table, token_to_req, @@ -317,6 +391,8 @@ def qsa_sparse_paged_attention( v_cache.stride(0), v_cache.stride(1), v_cache.stride(2), + 0 if k_scale is None else k_scale.stride(0), + 0 if v_scale is None else v_scale.stride(0), logical_indices.stride(0), block_table.stride(0), out.stride(0), @@ -334,6 +410,7 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, + HAS_KV_SCALE=k_scale is not None, num_warps=partial_warps, num_stages=2, ) diff --git a/python/freetoken/kernel/triton/qsa/score.py b/python/freetoken/kernel/triton/qsa/score.py index 49d702082..7b65da8d6 100644 --- a/python/freetoken/kernel/triton/qsa/score.py +++ b/python/freetoken/kernel/triton/qsa/score.py @@ -139,6 +139,14 @@ def qsa_mqa_paged( raise ValueError("QSA request mapping and positions must match query rows") if sequence_lengths.shape != (page_table.shape[0],): raise ValueError("QSA sequence lengths must match page-table requests") + if q.dtype not in (torch.bfloat16, torch.float16) or k_cache.dtype is not q.dtype: + # The dot below is a plain 16-bit matmul: an fp8 operand is not a slow path, it + # is a triton compile error that surfaces mid-CUDA-graph-capture. The KV cache + # may well be e4m3 codes (--kv-cache-dtype fp8) -- what must never reach here + # are those codes; the compressed index keys are their own 16-bit tier. + raise ValueError( + f"QSA scoring is 16-bit only, got query={q.dtype} keys={k_cache.dtype}" + ) score_divisor = math.sqrt(q.shape[2]) if score_scale is None else score_scale columns = logits.shape[1] if not q.shape[0] or not columns: diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 41b82f1d5..a7510d69f 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -76,6 +76,17 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: return MHAKVCache +def _reject_unsupported_quant(pool: str, kv_quant: str) -> None: + """A pool family that has no fp8 store/scale-read path must say so at startup, + not silently serve a 16-bit cache the budget priced for an fp8 one.""" + if kv_quant != "none": + raise ValueError( + f"--kv-cache-dtype {kv_quant} is not implemented for the {pool} KV pool " + "(only the plain paged / hybrid-SWA pools, served by the triton attention " + "backend); use --kv-cache-dtype bf16." + ) + + 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 @@ -85,10 +96,12 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt from .dsv4_paged_pool import DSV4PagedKVCache model_config = config.model_config + kv_quant = getattr(config, "kv_quant", "none") 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 # shadows). Sizing reads dsv4_args, never the group spec. + _reject_unsupported_quant("DSV4 paged", kv_quant) pool = DSV4PagedKVCache( sizes=_dsv4_pool_sizes(config, num_pages + 1), # +1 for dummy page args=model_config.dsv4_args, @@ -117,6 +130,7 @@ def create_kv_pool(config, num_pages: int, device: torch.device, dtype: torch.dt device=device, dtype=dtype, num_req_slots=config.max_running_req + 1, # + 1 for the dummy request row + kv_quant=kv_quant, ) @@ -128,6 +142,7 @@ def create_kvcache_pool( device: torch.device, num_swa_tokens: int | None = None, num_req_slots: int | None = None, + kv_quant: str = "none", ) -> BaseKVCachePool: if model_config.has_swa_attention: from .hybrid_swa_pool import HybridSWAKVCache @@ -140,6 +155,7 @@ def create_kvcache_pool( num_swa_tokens=num_swa_tokens, device=device, dtype=dtype, + kv_quant=kv_quant, ) from .mha_pool import MHAKVCache @@ -167,6 +183,7 @@ def create_kvcache_pool( if len(kv_specs) == 1 and kv_specs[0].attn_type == _AttnType.BSA: from .bsa_pool import BSAKVCache + _reject_unsupported_quant("block-sparse (BSA)", kv_quant) spec = kv_specs[0] assert layer_ids is None, "hybrid-linear x BSA has no pool support yet" return BSAKVCache( @@ -204,11 +221,15 @@ def create_kvcache_pool( index_ratio=spec.index_ratio, num_req_slots=num_req_slots, layer_ids=spec.layer_ids, + # Quantizes the KV tiers only -- the compressed index slab the score kernel + # reads stays the engine dtype (kvcache/qsa_pool.py). + kv_quant=kv_quant, ) if len(kv_specs) == 1 and kv_specs[0].mla: from .dsa_pool import DSAKVCache, KpoolDSAKVCache, MLAKVCache + _reject_unsupported_quant("latent-KV (MLA/DSA)", kv_quant) spec = kv_specs[0] # With a layer remap the pool allocates len(layer_ids) slabs; without one # it backs every model layer (all-MLA models, GLM-5.2). @@ -255,6 +276,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 95669e8c8..abcb916c1 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -9,6 +9,9 @@ logger = init_logger(__name__) +# One fp32 scale per (token, slab, layer, kv head) rides alongside fp8 KV codes. +FP8_KV_SCALE_BYTES = 4 + class CacheRebuildRejected(Exception): """A runtime cache rebuild was rejected BEFORE any destructive free (e.g. the @@ -16,12 +19,40 @@ class CacheRebuildRejected(Exception): this is recoverable, unlike a failure after the free.""" +def kv_storage_bytes_per_elem(config) -> int: + """Storage bytes of ONE cached KV element under the configured quantization. + + ``kv_quant == "fp8"`` stores e4m3 codes (1 byte) instead of the 16-bit compute + dtype; anything else is the compute dtype itself. Single source for the pool's + allocation, the budget math below, and the AOT kernel-shape table. + """ + 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 kv_scale_bytes_per_token(spec, config) -> int: + """Sidecar scale bytes per token of one group: the fp8 cache keeps one fp32 scale + per (token, slab, layer, kv head). 0 for the unquantized pool. + + Priced here rather than inside the pool so ``kv_cost`` and the pool's own + allocation can never disagree -- the same rule the 16-bit path follows.""" + if getattr(config, "kv_quant", "none") != "fp8": + return 0 + heads = div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) + return (1 if spec.mla else 2) * spec.num_layers * heads * FP8_KV_SCALE_BYTES + + 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 - 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.) + x layers, plus the fp8 scale sidecar when the cache is quantized, 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.) ``index_ratio`` > 1 (QSA) stores one index key per token group, not per token; that slab's ring and scratch rows are fixed-size and priced in QSAKVCache.kv_cost instead.""" @@ -29,10 +60,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 ) - return per_token + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio + return ( + per_token + + kv_scale_bytes_per_token(spec, config) + + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio + ) class BaseKVCachePool(ABC): @@ -45,6 +80,12 @@ class BaseKVCachePool(ABC): # model re-bound after a rebuild; the engine asks before it resizes. needs_rebind_on_rebuild: ClassVar[bool] = False + # KV storage quantization: "none" keeps the compute dtype, "fp8" keeps e4m3 codes plus + # one fp32 scale per (token, slab, layer, kv head). A pool that does not implement the + # quantized allocation/store/scale-view trio stays at "none"; create_kv_pool rejects a + # quantization the family does not implement, so nothing here is ever silently ignored. + kv_quant: str = "none" + # ---- sizing/cost classmethods: run BEFORE the pool exists (startup budget solve, # --moe-cache-auto). The engine measures memory and passes bytes in; each pool family # implements kv_cost for ITS OWN buffers only (the engine sums families, e.g. adds @@ -154,13 +195,36 @@ def store_kv( layer_id: int, ) -> None: ... + def k_scale(self, index: int) -> torch.Tensor | None: + """fp32 ``[num_slots, local_kv_heads]`` scale of ``k_cache(index)``, indexed by the + same slot; None when the pool is not quantized.""" + return None + + def v_scale(self, index: int) -> torch.Tensor | None: + """fp32 ``[num_slots, local_kv_heads]`` scale of ``v_cache(index)``; see k_scale.""" + return None + @property @abstractmethod def device(self) -> torch.device: ... @property @abstractmethod - def dtype(self) -> torch.dtype: ... + def dtype(self) -> torch.dtype: + """The pool's COMPUTE dtype: the dtype of the K/V rows a backend hands to + ``store_kv``, and what backends size their scratch with. A quantized (e4m3) + pool still answers 16-bit here -- code bytes handed to a scratch buffer end up + as the rhs of a ``tl.dot`` that has no fp8 path (QSA's indexer died this way at + graph capture). See :attr:`store_dtype` for what the buffer holds.""" + ... + + @property + def store_dtype(self) -> torch.dtype: + """Element type of the KV buffer: e4m3/uint8 codes on a quantized pool, else + :attr:`dtype`. Only code that touches the buffer itself needs this; attention + backends that cannot apply the row scales are refused a quantized pool up + front (``BackendInfo.supports_fp8_kv``).""" + return self.dtype @property @abstractmethod diff --git a/python/freetoken/kvcache/hybrid_swa_pool.py b/python/freetoken/kvcache/hybrid_swa_pool.py index 41c3880e3..1813128b2 100644 --- a/python/freetoken/kvcache/hybrid_swa_pool.py +++ b/python/freetoken/kvcache/hybrid_swa_pool.py @@ -23,6 +23,49 @@ class _KVGroupStorage: k_buffer: torch.Tensor v_buffer: torch.Tensor storage_shape: tuple[int, int, int] + # (2, num_layers, num_slots, local_kv_heads) fp32, or None for an unquantized group. + scale_buffer: torch.Tensor | None = None + + +def _alloc_group_storage( + *, + num_layers: int, + local_kv_heads: int, + head_dim: int, + device: torch.device, + store_dtype: torch.dtype, + outer_size: int, + inner_size: int, + quantized: bool, +) -> _KVGroupStorage: + """One group's code buffer (+ fp8 scale buffer), shared by the initial allocation + and the in-place rebuild so the two can never drift. + + A quantized buffer is zero-filled: a stale e4m3 code decodes to a real number, so + an unwritten slot (the dummy page, a padded request row) would poison attention, + while code 0x00 is exactly 0.0. One memset per allocation, same as + kvcache/bsa_pool.py. The 16-bit buffer keeps torch.empty. + """ + shape = (2, num_layers, outer_size, inner_size, local_kv_heads, head_dim) + if quantized: + from freetoken.kernel.triton.kv_quant import alloc_codes + + buffer = alloc_codes(shape, device) + scale = torch.zeros( + (2, num_layers, outer_size * inner_size, local_kv_heads), + device=device, + dtype=torch.float32, + ) + else: + buffer = torch.empty(shape, device=device, dtype=store_dtype) + scale = None + return _KVGroupStorage( + buffer=buffer, + k_buffer=buffer[0], + v_buffer=buffer[1], + storage_shape=(outer_size * inner_size, local_kv_heads, head_dim), + scale_buffer=scale, + ) class HybridSWAKVCache(BaseKVCachePool): @@ -37,14 +80,22 @@ def __init__( dtype: torch.dtype, device: torch.device, num_swa_tokens: int | None = None, + kv_quant: str = "none", ) -> None: specs = {group.name: group for group in groups if group.num_layers > 0} if set(specs) != {"full", "swa"}: raise ValueError(f"HybridSWAKVCache requires full and swa groups, got {sorted(specs)}") + from .mha_pool import _kv_store_dtype + self._num_layers = num_layers self._device = device - self._dtype = dtype + self.kv_quant = kv_quant + self._compute_dtype = dtype + # What the BUFFER holds -- fp8/uint8 codes when quantized -- reported as + # store_dtype. The dtype property keeps answering the compute dtype, same + # contract as MHAKVCache (kvcache/base.py): backends size their scratch with it. + self._store_dtype = _kv_store_dtype(dtype, kv_quant) self._full_num_tokens = num_full_pages * page_size self._swa_num_tokens = num_swa_tokens if num_swa_tokens is not None else self._full_num_tokens self._page_size = page_size @@ -63,6 +114,7 @@ def __init__( inner_size=page_size, dtype=dtype, device=device, + kv_quant=kv_quant, ) self.swa_kv_pool = self._allocate_group( specs["swa"], @@ -71,6 +123,7 @@ def __init__( inner_size=1, dtype=dtype, device=device, + kv_quant=kv_quant, ) self._storages = { "full": self.full_kv_pool, @@ -88,18 +141,20 @@ def _allocate_group( inner_size: int, dtype: torch.dtype, device: torch.device, + kv_quant: str = "none", ) -> _KVGroupStorage: + from .mha_pool import _kv_store_dtype + local_kv_heads = div_even(spec.num_kv_heads, tp_size, allow_replicate=True) - buffer = torch.empty( - (2, spec.num_layers, outer_size, inner_size, local_kv_heads, spec.head_dim), + return _alloc_group_storage( + num_layers=spec.num_layers, + local_kv_heads=local_kv_heads, + head_dim=spec.head_dim, device=device, - dtype=dtype, - ) - return _KVGroupStorage( - buffer=buffer, - k_buffer=buffer[0], - v_buffer=buffer[1], - storage_shape=(outer_size * inner_size, local_kv_heads, spec.head_dim), + store_dtype=_kv_store_dtype(dtype, kv_quant), + outer_size=outer_size, + inner_size=inner_size, + quantized=kv_quant != "none", ) @staticmethod @@ -200,6 +255,16 @@ def v_cache(self, index: int) -> torch.Tensor: ref = self.layers_mapping[index] return self._storages[ref.group].v_buffer[ref.index] + def k_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scale = self._storages[ref.group].scale_buffer + return None if scale is None else scale[0][ref.index] + + def v_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scale = self._storages[ref.group].scale_buffer + return None if scale is None else scale[1][ref.index] + def store_kv( self, k: torch.Tensor, @@ -207,13 +272,26 @@ def store_kv( out_loc: torch.Tensor, layer_id: int, ) -> None: - from freetoken.kernel import store_cache - ref = self.layers_mapping[layer_id] storage = self._storages[ref.group] indices = out_loc if ref.group == "swa": indices = self.translate_loc_from_full_to_swa(out_loc) + 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=indices, + k_cache=storage.k_buffer[ref.index].view(storage.storage_shape), + v_cache=storage.v_buffer[ref.index].view(storage.storage_shape), + k_scale=storage.scale_buffer[0][ref.index], + v_scale=storage.scale_buffer[1][ref.index], + ) + return + from freetoken.kernel import store_cache + store_cache( k_cache=storage.k_buffer[ref.index].view(storage.storage_shape), v_cache=storage.v_buffer[ref.index].view(storage.storage_shape), @@ -236,7 +314,12 @@ def device(self) -> torch.device: @property def dtype(self) -> torch.dtype: - return self._dtype + return self._compute_dtype + + @property + def store_dtype(self) -> torch.dtype: + """Element type of the KV buffers: e4m3/uint8 codes when quantized.""" + return self._store_dtype @property def num_layers(self) -> int: @@ -245,24 +328,31 @@ def num_layers(self) -> int: @staticmethod def _group_geometry(group: _KVGroupStorage) -> tuple: # Everything the realloc needs that does NOT pin the old buffer alive: layer count, - # kv heads, head_dim, device, dtype. (Plain ints + device/dtype handles, no tensor.) + # kv heads, head_dim, device, storage dtype, and whether codes are fp8. + # (Plain ints + device/dtype handles, no tensor.) _, num_layers, _old_outer, _old_inner, local_kv_heads, head_dim = group.buffer.shape - return (num_layers, local_kv_heads, head_dim, group.buffer.device, group.buffer.dtype) + return ( + num_layers, + local_kv_heads, + head_dim, + group.buffer.device, + group.buffer.dtype, + group.scale_buffer is not None, + ) @staticmethod def _alloc_group(geom: tuple, outer_size: int, inner_size: int) -> _KVGroupStorage: # Only the outer (page/token) dimension changes; the rest comes from ``geom``. - num_layers, local_kv_heads, head_dim, device, dtype = geom - buffer = torch.empty( - (2, num_layers, outer_size, inner_size, local_kv_heads, head_dim), + num_layers, local_kv_heads, head_dim, device, store_dtype, quantized = geom + return _alloc_group_storage( + num_layers=num_layers, + local_kv_heads=local_kv_heads, + head_dim=head_dim, device=device, - dtype=dtype, - ) - return _KVGroupStorage( - buffer=buffer, - k_buffer=buffer[0], - v_buffer=buffer[1], - storage_shape=(outer_size * inner_size, local_kv_heads, head_dim), + store_dtype=store_dtype, + outer_size=outer_size, + inner_size=inner_size, + quantized=quantized, ) def rebuild(self, num_full_pages: int, num_swa_tokens: int | None = None) -> None: @@ -346,10 +436,16 @@ def unit_bytes(self) -> tuple[int, int]: full = self.full_kv_pool.buffer swa = self.swa_kv_pool.buffer full_tokens = int(full.shape[2]) * int(full.shape[3]) - return ( - int(full.numel() * full.element_size()) // full_tokens, - int(swa.numel() * swa.element_size()) // self._swa_num_tokens, - ) + kv = int(full.numel() * full.element_size()) // full_tokens + swa_b = int(swa.numel() * swa.element_size()) // self._swa_num_tokens + # fp8 codes are priced with their scale sidecar, matching kv_cost exactly. + if self.full_kv_pool.scale_buffer is not None: + fs = self.full_kv_pool.scale_buffer + kv += int(fs.numel() * fs.element_size()) // full_tokens + if self.swa_kv_pool.scale_buffer is not None: + ss = self.swa_kv_pool.scale_buffer + swa_b += int(ss.numel() * ss.element_size()) // self._swa_num_tokens + return kv, swa_b # ---- SWA pool sizing (pure arithmetic; the pool family's geometry formulas) ---- diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b96..c58fab366 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"`` halves the cache: rows become 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,13 @@ def __init__( dtype: torch.dtype, device: torch.device, layer_ids: Sequence[int] | None = None, + kv_quant: str = "none", ) -> None: 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 +67,41 @@ 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 the code buffer (and, when quantized, the scale buffer). + + A quantized buffer is zero-filled -- e4m3 has NaN bit patterns, so an + unwritten slot (the dummy page, a padded request's row) must not read back as + one. The 16-bit buffer keeps ``torch.empty``: it is bytes-sized, never + interpreted, and the memset would cost real startup time on a large cache. + """ + 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 +112,15 @@ 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 kv_cost(cls, config) -> tuple[int, int, int, int]: @@ -101,7 +141,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 +161,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 +178,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 +208,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/kvcache/qsa_pool.py b/python/freetoken/kvcache/qsa_pool.py index fddcdbd35..e79e59009 100644 --- a/python/freetoken/kvcache/qsa_pool.py +++ b/python/freetoken/kvcache/qsa_pool.py @@ -63,6 +63,7 @@ def __init__( num_req_slots: int, ring_capacity: int | None = None, layer_ids: Sequence[int] | None = None, + kv_quant: str = "none", ) -> None: if index_ratio < 1 or page_size % index_ratio != 0: # slot // index_ratio only names one group when a group never straddles a page. @@ -98,6 +99,14 @@ def __init__( dtype=dtype, device=device, layer_ids=layer_ids, + # "fp8" quantizes ONLY the KV tiers: codes + one scale per (slot, kv_head) + # replace the bf16 K/V, and store_kv's fused writer replaces the + # separate qsa_store_rows. The three index tiers below stay bf16 no matter + # what is asked for -- block selection reads a different tensor + # (models/qwen3_8_flash_next.py builds index_k in the engine dtype), so + # --kv-cache-dtype fp8 leaves retrieval quality, and the score kernel's + # dtype asserts, untouched. + kv_quant=kv_quant, ) self._zero_kv_slabs() self._alloc_index_tiers(num_pages) @@ -105,8 +114,11 @@ def __init__( def _zero_kv_slabs(self) -> None: # Defense-in-depth: the attend kernels pos-mask every K/V load (the real fix for # torch.empty's recycled NaN/Inf bit patterns), but a zeroed slab keeps any future - # unmasked read finite instead of model-poisoning. One memset per (re)allocation. - self._kv_buffer.zero_() + # unmasked read finite instead of model-poisoning. One memset per (re)allocation, + # through the byte view so it works whether the slab holds bf16 values or e4m3 + # codes -- a memset is the one op both representations support, and the codes + # live in bytes on every architecture (kv_quant.kv_codes_dtype). + self._kv_buffer.view(torch.uint8).zero_() def _alloc_index_tiers(self, num_pages: int) -> None: # ZERO-initialized: the score kernel reads whole rows of blocks unmasked and relies on @@ -145,6 +157,9 @@ def rebuild(self, num_pages: int) -> None: self._kv_buffer = None self._k_buffer = None self._v_buffer = None + # Same reason as above on an fp8 pool: a grown K/V slab whose scales are gone + # would serve quantized rows at the wrong scale rather than fail. + self._scale_buffer = None raise @classmethod diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 5b4db587d..36e7e231d 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -368,6 +368,21 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Set the page size for system management.", ) + parser.add_argument( + "--kv-cache-dtype", + dest="kv_quant", + type=str, + default=ServerArgs.kv_quant, + choices=["auto", "bf16", "fp8"], + help=( + "KV-cache storage format. 'bf16' (default) stores the compute dtype; 'fp8'" + " stores e4m3 codes plus one fp32 scale per (token, kv head), roughly " + "doubling the tokens that fit in the same VRAM. Requires the triton" + " attention backend and a plain paged, hybrid-SWA or QSA sparse KV pool" + " (not MLA/DSA, DSV4, or MiniMax-M3 block-sparse models)." + ), + ) + parser.add_argument( "--attention-backend", "--attn", diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py new file mode 100644 index 000000000..9b870273c --- /dev/null +++ b/tests/engine/test_kv_quant_config.py @@ -0,0 +1,184 @@ +"""Config-time gates for ``--kv-cache-dtype`` (EngineConfig.kv_quant). + +fp8 KV is only half a feature: the pool has to store it AND the attention backend has +to read the scales. Everything here must fail while the config is still a dataclass -- +after weights are resident, a wrong combination has already cost a load and (worse) +fi/fa/trtllm would happily attend over raw e4m3 codes and produce plausible garbage. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.attention import AttnType +from freetoken.models.config import KVCacheGroupSpec + + +def _spec(name, attn_type, *, mla=False, sliding_window=None, index_head_dim=0, index_ratio=1): + return KVCacheGroupSpec( + name=name, + layer_ids=(0, 1), + num_kv_heads=1, + head_dim=64, + sliding_window=sliding_window, + mla=mla, + index_head_dim=index_head_dim, + num_index_layers=2 if index_head_dim else 0, + index_ratio=index_ratio, + attn_type=attn_type, + ) + + +def _model_config(kind): + mc = SimpleNamespace( + model_type=kind, + single_stream_only=False, + is_moe=False, + expert_quant="none", + has_swa_attention=False, + has_linear_attention=False, + num_layers=4, + rotary_config=SimpleNamespace(max_position=1024), + ) + specs = { + "full": (_spec("full", AttnType.FULL),), + "swa": ( + _spec("full", AttnType.FULL), + _spec("swa", AttnType.SWA, sliding_window=128), + ), + "mla": (_spec("full", AttnType.MLA, mla=True),), + "dsa": (_spec("full", AttnType.DSA, mla=True, index_head_dim=128),), + "dsv4": (_spec("dsv4", AttnType.DSV4, sliding_window=128),), + "bsa": (_spec("full", AttnType.BSA, index_head_dim=128),), + "qsa": (_spec("full", AttnType.QSA, index_head_dim=128, index_ratio=4),), + }[kind] + if kind == "swa": + mc.has_swa_attention = True + if kind == "dsv4": + mc.dsv4_args = SimpleNamespace(window_size=128) + if kind == "qsa": + mc.has_linear_attention = True + mc.kv_cache_group_specs = lambda: specs + return mc + + +def _config(kind, **overrides): + from freetoken.distributed import DistributedInfo + from freetoken.engine.config import EngineConfig + + config = EngineConfig( + model_path="/tmp/freetoken-test-model", + tp_info=DistributedInfo(rank=0, size=1), + dtype=torch.bfloat16, + **overrides, + ) + object.__setattr__(config, "model_config", _model_config(kind)) + return config + + +def _patch_fast_machine(monkeypatch): + """A machine where every fast external backend is available, so an fp8 result can + only come from the gate and not from a missing package.""" + from freetoken.engine import engine + + monkeypatch.setattr(engine, "is_sm100_family", lambda: False) + monkeypatch.setattr(engine, "is_sm90_family", lambda: True) + monkeypatch.setattr(engine, "_flashinfer_available", lambda: True) + monkeypatch.setattr(engine, "_sgl_flash_attn_available", lambda: True) + + +def test_kv_quant_spellings(): + from freetoken.engine.engine import _resolve_kv_quant + + assert _resolve_kv_quant("auto") == "none" + assert _resolve_kv_quant("bf16") == "none" + assert _resolve_kv_quant("FP8") == "fp8" + assert _resolve_kv_quant(None) == "none" + with pytest.raises(ValueError, match="kv-cache-dtype"): + _resolve_kv_quant("q8") + + +def test_only_the_backends_that_read_scales_declare_fp8_support(): + from freetoken.attention import SUPPORTED_ATTENTION_BACKENDS, attention_backend_info + + fp8 = { + name + for name in SUPPORTED_ATTENTION_BACKENDS.supported_names() + if attention_backend_info(name).supports_fp8_kv + } + # triton serves the plain paged / hybrid-SWA pools and applies the scales in + # kernel/triton/attention.py; qsa_sparse dequantizes the rows it selects in + # kernel/triton/qsa/attend.py. Nothing else may join this set: fi/fa/trtllm (and the + # dsa/dsv4_sparse kernels) have no scale path and would attend over raw e4m3 codes, + # producing plausible garbage instead of an error. + assert fp8 == {"triton", "qsa_sparse"} + + +def test_auto_avoids_the_fast_backends_for_fp8(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + # Same machine where a 16-bit cache auto-selects the sm90 "fa,fi" tree... + plain = _config("full", attention_backend="auto") + _adjust_config(plain) + assert plain.kv_quant == "none" + assert plain.attention_backend == "fa,fi" + + quantized = _config("full", attention_backend="auto", kv_quant="fp8") + _adjust_config(quantized) + assert quantized.attention_backend == "triton" + + +@pytest.mark.parametrize("backend", ["fi", "fa", "trtllm", "fi,triton", "triton,fi"]) +def test_explicit_unsupported_backend_is_rejected(monkeypatch, backend): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + monkeypatch.setattr( + "freetoken.engine.engine.is_sm100_family", lambda: True + ) # let trtllm clear its own arch gate first + config = _config("full", attention_backend=backend, kv_quant="fp8", page_size=1) + with pytest.raises(ValueError, match="kv-cache-dtype fp8"): + _adjust_config(config) + + +def test_explicit_triton_is_accepted(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("full", attention_backend="triton", kv_quant="fp8") + _adjust_config(config) + assert config.kv_quant == "fp8" + + +@pytest.mark.parametrize("kind", ["mla", "dsa", "dsv4", "bsa"]) +def test_pool_families_without_a_scale_read_path_are_rejected(monkeypatch, kind): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) # the rejection must not depend on this box's wheels + config = _config(kind, attention_backend="auto", kv_quant="fp8") + with pytest.raises(ValueError, match="kv-cache-dtype fp8"): + _adjust_config(config) + + +def test_qsa_keeps_fp8_available(monkeypatch): + """The QSA pool is the block-sparse family whose K/V rows do reach a kernel that can + dequantize them; the compressed index keys selection scores against are a separate, + always-16-bit tier, so the gate has nothing left to refuse here.""" + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("qsa", attention_backend="auto", kv_quant="fp8") + _adjust_config(config) + assert config.kv_quant == "fp8" + assert "qsa_sparse" in config.attention_backend + + +def test_swa_pool_keeps_fp8_available(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("swa", attention_backend="auto", kv_quant="fp8") + _adjust_config(config) + assert config.kv_quant == "fp8" and config.attention_backend == "triton" diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py index 8bcb8e581..c5ca36553 100644 --- a/tests/kernels/test_e4m3_compat.py +++ b/tests/kernels/test_e4m3_compat.py @@ -16,6 +16,9 @@ capability patched to sm_80/86/89/120 and every triton launch forced into warmup (compile-only), the full wrapper->kernel paths -- uint8 views, PDL gates, e4m3 branches -- must compile for the foreign arch. +4. A cache-key guard: the @constexpr_function probes must not reference plain + python functions. triton's AST walk rejects that, and it disables every kernel + that branches on e4m3 native-ness at once, so the failure is not local. """ from __future__ import annotations @@ -41,6 +44,32 @@ def _native_cc() -> bool: return torch.cuda.get_device_capability() >= (8, 9) +def test_constexpr_probe_never_references_a_host_function(): + """triton hashes an @constexpr_function by walking its AST (runtime/jit.py: + cache_key -> record_reference), and a bare reference to a plain python function + raises "Unsupported function referenced: ". Making + e4m3_native_cx defer to the host probe did exactly that, and it took out EVERY + kernel that branches on e4m3 native-ness at once (here: the PLE gather, inside + CUDA graph capture). Modules (``target_info.cuda_capability_geq``) survive the + walk, functions do not -- so unifying the two probes has to go the other way: the + code that owns a buffer follows the buffer (tests/kernels/test_kv_fp8.py).""" + import inspect + + from freetoken.kernel.triton import e4m3_compat + + lines = inspect.getsource(e4m3_compat).splitlines() + start = next(i for i, ln in enumerate(lines) if ln.startswith("def e4m3_native_cx(")) + body = [] + for line in lines[start + 1:]: + if line and not line.startswith((" ", "\t")): + break + body.append(line) + assert "e4m3_native(" not in "\n".join(body), ( + "a constexpr_function may not call a host function: triton's cache-key walk " + "rejects it and every e4m3 kernel stops compiling" + ) + + # ====================================================================================== # 1. Primitives vs the native fp8 unit (needs sm_89+ hardware for the reference). # ====================================================================================== @@ -228,6 +257,60 @@ def _emit_all(path: str) -> None: out["nv_moe_prefill"] = f32(fused_experts_nvfp4( hid4, gup, gus4, gug, dnp, dns4, dng, tw, tids, S)) + # fp8 KV cache: the fused quantize+scatter writer the pools call in store_kv, and the + # QSA sparse reader that turns codes + row scales back into operands. Codes are plain + # bytes on every arch (kv_quant.kv_codes_dtype), so comparing them as BYTES is what + # pins a native run and a forced-EMU run to the very same encoding, and the two + # QSA runs -- one over codes, one over the same real numbers pre-rounded into bf16 -- + # must produce identical bits: the reader casts back to the query dtype before the + # dot, so any decode difference lands on the comparison instead of hiding in a + # tolerance. + from freetoken.kernel.triton.kv_quant import ( + alloc_codes, codes_to_f32, quantize_kv_to_cache, + ) + from freetoken.kernel.triton.qsa import qsa_sparse_paged_attention + + slots, kvh, hd, page = 128, 2, 64, 64 + kv_in = (torch.randn(slots, kvh * hd, device=dev, dtype=torch.float32) * 4.0).to( + torch.bfloat16 + ) + k_codes = alloc_codes((slots, kvh, hd), dev) + v_codes = alloc_codes((slots, kvh, hd), dev) + k_sc = torch.zeros((slots, kvh), dtype=torch.float32, device=dev) + v_sc = torch.zeros_like(k_sc) + quantize_kv_to_cache( + k=kv_in, + v=kv_in.flip(-1).contiguous(), + out_loc=torch.arange(slots, dtype=torch.int32, device=dev), + k_cache=k_codes, + v_cache=v_codes, + k_scale=k_sc, + v_scale=v_sc, + ) + out["kvfp8_codes"] = k_codes.view(torch.uint8).to(torch.int16).cpu() + out["kvfp8_scale"] = k_sc.cpu() + + pages = slots // page + kc4, vc4 = k_codes.view(pages, page, kvh, hd), v_codes.view(pages, page, kvh, hd) + q = torch.randn(2, 2 * kvh, hd, device=dev, dtype=torch.bfloat16) + sel = ( + torch.arange(2 * page, dtype=torch.int32, device=dev)[None, :] + .repeat(2, 1) + .contiguous() + ) + table = torch.arange(pages, dtype=torch.int32, device=dev)[None, :].contiguous() + t2r = torch.zeros(2, dtype=torch.int32, device=dev) + out["kvfp8_qsa"] = f32(qsa_sparse_paged_attention( + q, kc4, vc4, sel, table, t2r, k_scale=k_sc, v_scale=v_sc)) + # The very same numbers, pre-rounded into a bf16 cache. The reader casts its + # dequantized operands to the query dtype before tl.dot, so the two runs must end up + # bit-identical (asserted where launches really execute: tests/kernels/test_qsa_fp8.py + # -- the compile gate below runs warmup-only, where outputs are never written). + out["kvfp8_qsa_bf16"] = f32(qsa_sparse_paged_attention( + q, + (codes_to_f32(kc4) * k_sc.view(pages, page, kvh, 1)).to(torch.bfloat16), + (codes_to_f32(vc4) * v_sc.view(pages, page, kvh, 1)).to(torch.bfloat16), + sel, table, t2r)) torch.save(out, path) diff --git a/tests/kernels/test_kv_fp8.py b/tests/kernels/test_kv_fp8.py new file mode 100644 index 000000000..26deef2ac --- /dev/null +++ b/tests/kernels/test_kv_fp8.py @@ -0,0 +1,262 @@ +"""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 + + +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) + ) + 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" + ) + + +def test_kv_codec_has_no_arch_or_dtype_branch(): + """The two rejected designs had one thing in common: they chose an arm. + + The compile-time fp8-native probe answers a question the allocator already + answered -- and on one box answered wrongly -- while a test against the pointer's + element type is NOT pruned by triton, so the dead arm still gets type-checked. + That is how an int mask fill ended up in front of an fp8 pointer, twice. Both + codecs are straight-line now; pin that, plus the identifiers of the two rejected + designs, so neither creeps back in as a "fast path". + """ + import inspect + + from freetoken.kernel.triton import kv_quant + from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 + + for obj in (kv_load_e4m3_tile_f32, kv_quant._kv_quant_scatter_kernel): + src = inspect.getsource(getattr(obj, "fn", obj)) + for banned in ("e4m3_native", "dtype.element_ty"): + assert banned not in src, f"{banned} is back in {obj.__name__}" + body = src.split('"""')[-1].splitlines() + arms = [ + line.strip() for line in body + if line.strip().startswith(("if ", "elif ", "else")) + ] + assert not arms, f"{obj.__name__} must not branch: {arms}" diff --git a/tests/kernels/test_qsa_fp8.py b/tests/kernels/test_qsa_fp8.py new file mode 100644 index 000000000..1fa6669f8 --- /dev/null +++ b/tests/kernels/test_qsa_fp8.py @@ -0,0 +1,278 @@ +"""QSA sparse attention reading an fp8 KV pool (kernel/triton/qsa/attend.py). + +The oracle is the SAME data in a bf16 cache, not a tolerance. The kernel widens the e4m3 +codes to fp32, multiplies by the row scale, and casts back to the query dtype before +``tl.dot`` -- so a cache holding ``c * s`` and a cache holding ``c`` with scale ``s`` feed +the matmul bit-identical operands, and the two runs must agree bit for bit. Anything wrong +in scale indexing (the page/offset -> slot arithmetic), in the K-vs-V broadcast direction, +or in the masked-fill values shows up as a mismatch instead of slop inside an epsilon. + +Two quantizer variants are covered: + * hand-made codes with power-of-two scales, which also keeps the real values on the e4m3 + grid, so the fp8 buffer, the bf16 buffer and the test agree on the numbers; + * the fused store kernel the pools actually call (``kv_quant.quantize_kv_to_cache``, + amax/448 scales) over ordinary gaussian rows -- whose quantization error against the + ORIGINAL rows is bounded separately, because that part is quality, not exactness. + +Each parametrization also covers a different launcher profile: a small +``rows * kv_heads`` pushes it into split-K (partials + merge kernel), a large one takes +the direct-write path. +""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.kv_quant import ( + FP8, + alloc_codes, + codes_to_f32, + kv_codes_dtype, + quantize_kv_to_cache, +) +from freetoken.kernel.triton.qsa import qsa_sparse_paged_attention + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="Triton attention needs CUDA" +) + +PAGE = 64 # the page size the qsa_sparse backend registers +HEAD_DIM = 64 + + +def _on_grid(shape, device, generator) -> torch.Tensor: + """Values of the form +-m * 2^e with m in [8, 16): four significant bits, i.e. every + one of them is representable in e4m3 AND in bf16, so encoding is lossless and both + caches hold the very same real numbers.""" + mantissa = torch.randint( + 8, 16, shape, device=device, generator=generator, dtype=torch.int32 + ) + exponent = torch.randint( + -6, 5, shape, device=device, generator=generator, dtype=torch.int32 + ) + sign = torch.where( + torch.randint(0, 2, shape, device=device, generator=generator).bool(), 1.0, -1.0 + ) + return sign * mantissa.to(torch.float32) * torch.pow(2.0, exponent.to(torch.float32)) + + +def _code_buffer(values_f32: torch.Tensor) -> torch.Tensor: + """Encode e4m3-exact fp32 values into a buffer of the pool's code dtype.""" + codes = values_f32.to(FP8) + if kv_codes_dtype() is torch.uint8: + codes = codes.view(torch.uint8) + assert codes.dtype == kv_codes_dtype(), codes.dtype + return codes.contiguous() + + +def _layout(rows: int, topk: int, num_req: int): + """(indices, block_table, token_to_req) over three 64-token pages. + + Tokens 32..31+topk straddle pages 0 and 1, and the two requests map their logical + pages to physical ones in OPPOSITE order -- a page-table or page/offset slip cannot + hide behind a symmetric fixture. + """ + device = "cuda" + block_table = torch.tensor( + [[0, 1, 2], [2, 1, 0]], dtype=torch.int32, device=device + )[:num_req].contiguous() + indices = ( + torch.arange(topk, dtype=torch.int32, device=device)[None, :] + 32 + ).repeat(rows, 1).contiguous() + token_to_req = ( + torch.arange(rows, dtype=torch.int32, device=device) % num_req + ).contiguous() + return indices, block_table, token_to_req + + +def _run(codes, scales, q, indices, block_table, token_to_req): + return qsa_sparse_paged_attention( + q, + codes[0], + codes[1], + indices, + block_table, + token_to_req, + k_scale=scales[0], + v_scale=scales[1], + ) + + + +# rows 1 x 1 kv head -> base_programs 1 -> BLOCK_N 16, 4 tiles -> NUM_SPLITS 4 (split-K). +# rows 16 x 2 -> base_programs 32 -> BLOCK_N 64, 1 tile -> NUM_SPLITS 1 (direct). +@pytest.mark.parametrize(("rows", "kv_heads", "topk"), [(1, 1, 64), (16, 2, 64)]) +def test_fp8_codes_match_the_bf16_cache_bit_for_bit(rows, kv_heads, topk): + torch.manual_seed(3) + device = torch.device("cuda") + num_pages, num_query_heads = 3, 2 * kv_heads + slots = num_pages * PAGE + + k_values = _on_grid((num_pages, PAGE, kv_heads, HEAD_DIM), device, None) + v_values = _on_grid((num_pages, PAGE, kv_heads, HEAD_DIM), device, None) + # Alternate the scale exponent by slot -- a scale broadcast along the wrong axis then + # changes the answer instead of cancelling out -- and give K and V opposite parities. + parity = torch.arange(slots, device=device, dtype=torch.float32) % 2 + k_scale = ( + torch.pow(2.0, parity * 5 - 3).unsqueeze(-1).expand(slots, kv_heads).contiguous() + ) + v_scale = ( + torch.pow(2.0, (1 - parity) * 4 - 2).unsqueeze(-1).expand(slots, kv_heads).contiguous() + ) + k_ref = (k_values * k_scale.view(num_pages, PAGE, kv_heads, 1)).to(torch.bfloat16) + v_ref = (v_values * v_scale.view(num_pages, PAGE, kv_heads, 1)).to(torch.bfloat16) + + q = torch.randn(rows, num_query_heads, HEAD_DIM, device=device, dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(rows, topk, num_req=2) + + got = _run( + (_code_buffer(k_values), _code_buffer(v_values)), + (k_scale, v_scale), + q, + indices, + block_table, + token_to_req, + ) + want = qsa_sparse_paged_attention( + q, k_ref, v_ref, indices, block_table, token_to_req + ) + assert torch.equal(got, want), ( + "fp8 QSA attend diverged from the same data in a bf16 cache (max diff " + f"{(got.float() - want.float()).abs().max().item():.3e})" + ) + + +@pytest.mark.parametrize(("rows", "kv_heads", "topk"), [(1, 1, 64), (16, 2, 64)]) +def test_pool_writer_codes_match_the_bf16_cache_bit_for_bit(rows, kv_heads, topk): + """Ordinary rows through the fused quantize+scatter writer store_kv calls: the + codes/scales it produces, decoded on the host with torch's own e4m3 cast (never the + decoder under test), must feed the dot identical operands.""" + torch.manual_seed(5) + device = torch.device("cuda") + num_pages, num_query_heads = 3, 2 * kv_heads + slots = num_pages * PAGE + + k_rows = torch.randn(slots, kv_heads * HEAD_DIM, device=device, dtype=torch.float32) + v_rows = torch.randn(slots, kv_heads * HEAD_DIM, device=device, dtype=torch.float32) + # Wide amplitude spread: every 8th row carries a 64x outlier. That is what a per-row + # amax scale exists to absorb, and what a scale read from the wrong slot blows up on. + outlier = (torch.arange(slots, device=device) % 8 == 0).unsqueeze(-1) + k_rows = k_rows * torch.where(outlier, 64.0, 1.0) + v_rows = v_rows * torch.where(outlier.flip(0), 32.0, 0.5) + + k_flat = alloc_codes((slots, kv_heads, HEAD_DIM), device) + v_flat = alloc_codes((slots, kv_heads, HEAD_DIM), device) + k_scale = torch.zeros((slots, kv_heads), dtype=torch.float32, device=device) + v_scale = torch.zeros_like(k_scale) + quantize_kv_to_cache( + k=k_rows, + v=v_rows, + out_loc=torch.arange(slots, dtype=torch.int32, device=device), + k_cache=k_flat, + v_cache=v_flat, + k_scale=k_scale, + v_scale=v_scale, + ) + torch.cuda.synchronize() + + k_ref = (codes_to_f32(k_flat) * k_scale.unsqueeze(-1)).to(torch.bfloat16) + v_ref = (codes_to_f32(v_flat) * v_scale.unsqueeze(-1)).to(torch.bfloat16) + q = torch.randn(rows, num_query_heads, HEAD_DIM, device=device, dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(rows, topk, num_req=2) + + shape = (num_pages, PAGE, kv_heads, HEAD_DIM) + got = _run( + (k_flat.view(shape), v_flat.view(shape)), + (k_scale, v_scale), + q, + indices, + block_table, + token_to_req, + ) + want = qsa_sparse_paged_attention( + q, k_ref.view(shape), v_ref.view(shape), indices, block_table, token_to_req + ) + assert torch.equal(got, want), ( + "fused-writer codes diverged (max diff " + f"{(got.float() - want.float()).abs().max().item():.3e})" + ) + + # Quantization QUALITY, bounded per row: e4m3's half-ulp is 2^-4 of the binade a + # value lands in, and the writer scales each row so its amax sits at 448. + for source, codes, scales in ((k_rows, k_flat, k_scale), (v_rows, v_flat, v_scale)): + assert torch.isfinite(scales).all() and (scales > 0).all() + decoded = codes_to_f32(codes) * scales.unsqueeze(-1) + original = source.view(slots, kv_heads, HEAD_DIM) + amax = original.abs().amax(dim=-1, keepdim=True) + rel = ((decoded - original).abs() / amax.clamp_min(1e-9)).max() + assert rel.item() <= 0.07, f"e4m3 grid error {rel.item():.4f} above its half-ulp" + + +@pytest.mark.parametrize("which", ["k_only", "v_only", "bf16_cache", "wrong_shape"]) +def test_scale_arguments_are_validated(which): + """A half-supplied or mismatched scale pair must fail loudly: attending over raw + codes while believing they are bf16 is the exact failure mode this feature cannot be + allowed to have, and it produces plausible garbage rather than an error.""" + device = "cuda" + kv_heads, rows = 1, 1 + k = torch.zeros(2, PAGE, kv_heads, HEAD_DIM, device=device, dtype=kv_codes_dtype()) + v = torch.zeros_like(k) + scale = torch.ones(2 * PAGE, kv_heads, dtype=torch.float32, device=device) + q = torch.zeros(rows, 2 * kv_heads, HEAD_DIM, device=device, dtype=torch.bfloat16) + indices = torch.zeros(rows, 8, dtype=torch.int32, device=device) + block_table = torch.zeros(1, 2, dtype=torch.int32, device=device) + token_to_req = torch.zeros(rows, dtype=torch.int32, device=device) + + kwargs = {"k_scale": scale, "v_scale": scale} + if which == "k_only": + kwargs.pop("v_scale") + elif which == "v_only": + kwargs.pop("k_scale") + elif which == "bf16_cache": + k, v = k.to(torch.bfloat16), v.to(torch.bfloat16) + else: + kwargs["k_scale"] = scale[:-1] + + with pytest.raises(ValueError, match="QSA"): + qsa_sparse_paged_attention( + q, k, v, indices, block_table, token_to_req, **kwargs + ) + + +def test_qsa_scoring_refuses_fp8_operands(): + """The indexer's scoring dot is 16-bit only, and the wrapper has to say so. + + In the field an e4m3 ``q_index`` -- produced by a backend that took its scratch dtype + from a pool reporting its STORE dtype -- surfaced as ``CompilationError: Unsupported + rhs dtype fp8e4nv`` inside CUDA-graph capture: a dead scheduler and a stopped API + server, forty seconds after the pool allocated fine. Same mistake, now stopped at + the call with a name on it. + """ + from freetoken.kernel.triton.qsa import qsa_mqa_paged + + device = "cuda" + rows, heads, dim, pages, cmp_page = 2, 2, 64, 2, 16 + good = { + "q": torch.zeros(rows, heads, dim, device=device, dtype=torch.bfloat16), + "k_cache": torch.zeros( + pages, cmp_page, 1, dim, device=device, dtype=torch.bfloat16 + ), + "page_table": torch.zeros(1, pages, dtype=torch.int32, device=device), + "token_to_req": torch.zeros(rows, dtype=torch.int32, device=device), + "query_positions": torch.arange(rows, dtype=torch.int32, device=device), + "sequence_lengths": torch.zeros(1, dtype=torch.int32, device=device), + "compress_ratio": cmp_page, + # Zero columns -> the wrapper returns before launching. This control proves the + # new check still lets the dtype it exists to allow through. + "logits": torch.zeros(rows, 0, dtype=torch.float32, device=device), + "visible_blocks": torch.zeros(rows, dtype=torch.int32, device=device), + } + qsa_mqa_paged(**good) + + for name in ("q", "k_cache"): + bad = dict(good) + bad[name] = alloc_codes(tuple(good[name].shape), device) + with pytest.raises(ValueError, match="16-bit only"): + qsa_mqa_paged(**bad) + diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index 6f4afca9e..ffff4a09b 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -781,3 +781,260 @@ def test_triton_metadata_keeps_full_indices_and_optional_swa_indices(monkeypatch assert metadata.indices.tolist() == [10, 11, 20, 21, 22] assert metadata.swa_indices is not None assert metadata.swa_indices.tolist() == [110, 111, 120, 121, 122] + + +def _fp8_cache(k_rows: torch.Tensor, v_rows: torch.Tensor): + """Quantize ``[slots, heads, dim]`` KV into codes + scales, in the layout the + attention kernels expect.""" + from freetoken.kernel.triton.kv_quant import alloc_codes, quantize_kv_to_cache + + slots, heads, dim = k_rows.shape + k_cache = alloc_codes((slots, heads, dim), k_rows.device) + v_cache = alloc_codes((slots, heads, dim), k_rows.device) + k_scale = torch.zeros((slots, heads), dtype=torch.float32, device=k_rows.device) + v_scale = torch.zeros((slots, heads), dtype=torch.float32, device=k_rows.device) + quantize_kv_to_cache( + k=k_rows.reshape(slots, heads * dim), + v=v_rows.reshape(slots, heads * dim), + out_loc=torch.arange(slots, dtype=torch.int32, device=k_rows.device), + 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 _dequantized(codes: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Host-side decode (torch's e4m3 cast), so the reference never reuses the + software decoder the kernel is being tested against.""" + from freetoken.kernel.triton.kv_quant import codes_to_f32 + + return codes_to_f32(codes) * scale.unsqueeze(-1) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +@pytest.mark.parametrize(("head_dim", "num_kv_heads"), [(64, 4), (128, 2)]) +def test_decode_paged_attention_decodes_fp8_scales(head_dim: int, num_kv_heads: int): + """fp8 decode must equal the SAME data dequantized by hand. + + Not compared against the bf16 cache: the fp8-vs-bf16 gap is quantization error, + already bounded in tests/kernels/test_kv_fp8.py. What this pins is that the + kernel applies the right scale to the right row -- a dropped or mis-indexed + scale is off by a factor of amax/448, which no tolerance hides. + """ + from freetoken.kernel.triton.attention import decode_paged_attention + + torch.manual_seed(11) + device = torch.device("cuda") + batch, num_q_heads, max_kv_splits = 2, 8, 8 + seq_lens = [6, 9] + total_kv = sum(seq_lens) + q = torch.randn(batch, num_q_heads, head_dim, device=device, dtype=torch.bfloat16) + # Row magnitudes far from 1, so a missing scale cannot pass by luck. + k_rows = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 8.0).to( + torch.bfloat16 + ) + v_rows = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 0.05).to( + torch.bfloat16 + ) + indptr = torch.tensor([0, seq_lens[0], total_kv], dtype=torch.int32, device=device) + indices = torch.arange(total_kv, dtype=torch.int32, device=device) + q_positions = torch.tensor( + [seq_lens[0] - 1, seq_lens[1] - 1], dtype=torch.int64, device=device + ) + q_to_req = torch.arange(batch, dtype=torch.int32, device=device) + attn_logits = torch.empty( + batch, num_q_heads, max_kv_splits, head_dim, dtype=torch.float32, device=device + ) + attn_lse = torch.empty(batch, num_q_heads, max_kv_splits, dtype=torch.float32, device=device) + num_kv_splits = torch.full((batch,), max_kv_splits, dtype=torch.int32, device=device) + sm_scale = head_dim**-0.5 + + k_codes, v_codes, k_scale, v_scale = _fp8_cache(k_rows, v_rows) + actual = decode_paged_attention( + q, + k_codes, + v_codes, + indptr, + indices, + q_positions, + attn_logits, + attn_lse, + num_kv_splits, + max_kv_splits, + sm_scale, + k_scale=k_scale, + v_scale=v_scale, + ) + expected = _reference_paged_attention( + q, + _dequantized(k_codes, k_scale), + _dequantized(v_codes, v_scale), + indptr, + indices, + q_to_req, + q_positions, + sm_scale, + None, + ) + # The kernel rounds the decoded operands to the compute dtype before tl.dot; the + # reference stays in fp32, hence the same tolerance the bf16 tests already use. + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +def test_paged_attention_decodes_fp8_scales(): + """The non-tl.dot fallback kernel takes the same scales (it serves head_dim > 256 + with a short prefill, and every decode batch when the grouped path is not used).""" + from freetoken.kernel.triton.attention import paged_attention + + torch.manual_seed(13) + device = torch.device("cuda") + head_dim, num_kv_heads, num_q_heads = 64, 2, 4 + seq_lens = [5, 3] + total_kv = sum(seq_lens) + q = torch.randn(total_kv, num_q_heads, head_dim, device=device, dtype=torch.bfloat16) + k_rows = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 16.0).to( + torch.bfloat16 + ) + v_rows = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 0.02).to( + torch.bfloat16 + ) + indptr = torch.tensor([0, seq_lens[0], total_kv], dtype=torch.int32, device=device) + indices = torch.arange(total_kv, dtype=torch.int32, device=device) + q_to_req = torch.tensor( + [0] * seq_lens[0] + [1] * seq_lens[1], dtype=torch.int32, device=device + ) + q_positions = torch.cat( + [ + torch.arange(seq_lens[0], dtype=torch.int64, device=device), + torch.arange(seq_lens[1], dtype=torch.int64, device=device), + ] + ) + k_codes, v_codes, k_scale, v_scale = _fp8_cache(k_rows, v_rows) + actual = paged_attention( + q, + k_codes, + v_codes, + indptr, + indices, + q_to_req, + q_positions, + head_dim**-0.5, + k_scale=k_scale, + v_scale=v_scale, + ) + expected = _reference_paged_attention( + q, + _dequantized(k_codes, k_scale), + _dequantized(v_codes, v_scale), + indptr, + indices, + q_to_req, + q_positions, + head_dim**-0.5, + None, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=2e-2, rtol=2e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +@pytest.mark.parametrize("use_split_inputs", [False, True]) +def test_extend_paged_attention_decodes_fp8_scales(use_split_inputs: bool): + """Prefill over an fp8 cache. + + With ``use_split_inputs`` the kernel reads a request's own new tokens from + ``k_extend`` (compute dtype, never quantized) and only the cached prefix from the + fp8 codes; without it every row is served from the codes. The reference mirrors + that split, so a scale leaking onto the extend path -- or failing to apply to the + cache path -- cannot pass. + """ + from freetoken.kernel.triton.attention import extend_paged_attention + + torch.manual_seed(14) + device = torch.device("cuda") + head_dim, num_kv_heads, num_q_heads = 64, 2, 8 + cached_lens, extend_lens = [4, 2], [3, 2] + seq_lens = [c + e for c, e in zip(cached_lens, extend_lens)] + total_q, total_kv = sum(extend_lens), sum(seq_lens) + q = torch.randn(total_q, num_q_heads, head_dim, device=device, dtype=torch.bfloat16) + k_extend = (torch.randn(total_q, num_kv_heads, head_dim, device=device) * 6.0).to( + torch.bfloat16 + ) + v_extend = (torch.randn(total_q, num_kv_heads, head_dim, device=device) * 6.0).to( + torch.bfloat16 + ) + # Magnitudes far from 1: a dropped scale is then off by orders of magnitude. + k_cache = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 0.3).to( + torch.bfloat16 + ) + v_cache = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 30.0).to( + torch.bfloat16 + ) + qo_indptr = torch.tensor([0] + extend_lens, dtype=torch.int32, device=device).cumsum_(0) + kv_indptr = torch.tensor([0] + seq_lens, dtype=torch.int32, device=device).cumsum_(0) + indices = torch.arange(total_kv, dtype=torch.int32, device=device) + prefix_lens = torch.tensor(cached_lens, dtype=torch.int32, device=device) + q_to_req = torch.empty(total_q, dtype=torch.int32, device=device) + q_positions = torch.empty(total_q, dtype=torch.int64, device=device) + q_off = kv_off = 0 + for req_idx, (cached_len, extend_len) in enumerate(zip(cached_lens, extend_lens)): + q_to_req[q_off : q_off + extend_len].fill_(req_idx) + q_positions[q_off : q_off + extend_len] = torch.arange( + cached_len, cached_len + extend_len, dtype=torch.int64, device=device + ) + # The step's own tokens are what the engine stores at the tail of the span. + k_cache[kv_off + cached_len : kv_off + cached_len + extend_len] = k_extend[ + q_off : q_off + extend_len + ] + v_cache[kv_off + cached_len : kv_off + cached_len + extend_len] = v_extend[ + q_off : q_off + extend_len + ] + q_off += extend_len + kv_off += cached_len + extend_len + sm_scale = head_dim**-0.5 + + k_codes, v_codes, k_scale, v_scale = _fp8_cache(k_cache, v_cache) + k_ref = _dequantized(k_codes, k_scale).clone() + v_ref = _dequantized(v_codes, v_scale).clone() + if use_split_inputs: + q_off = kv_off = 0 + for cached_len, extend_len in zip(cached_lens, extend_lens): + k_ref[kv_off + cached_len : kv_off + cached_len + extend_len] = k_extend[ + q_off : q_off + extend_len + ] + v_ref[kv_off + cached_len : kv_off + cached_len + extend_len] = v_extend[ + q_off : q_off + extend_len + ] + q_off += extend_len + kv_off += cached_len + extend_len + + actual = extend_paged_attention( + q=q, + k_cache=k_codes, + v_cache=v_codes, + qo_indptr=qo_indptr, + kv_indptr=kv_indptr, + kv_indices=indices, + prefix_lens=prefix_lens, + max_q_len=max(extend_lens), + sm_scale=sm_scale, + k_extend=k_extend if use_split_inputs else None, + v_extend=v_extend if use_split_inputs else None, + k_scale=k_scale, + v_scale=v_scale, + ) + expected = _reference_paged_attention( + q, + k_ref, + v_ref, + kv_indptr, + indices, + q_to_req, + q_positions, + sm_scale, + None, + ) + torch.testing.assert_close(actual.float(), expected.float(), atol=3e-2, rtol=3e-2) diff --git a/tests/kvcache/test_mha_pool_fp8.py b/tests/kvcache/test_mha_pool_fp8.py new file mode 100644 index 000000000..77bd2986e --- /dev/null +++ b/tests/kvcache/test_mha_pool_fp8.py @@ -0,0 +1,271 @@ +"""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): + from freetoken.kvcache.mha_pool import MHAKVCache + + _init_tp() + return MHAKVCache( + num_kv_heads=HEADS, + num_layers=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(): + # Hybrid GDN models back only their full-attention layers; a scale view that + # forgot the remap would hand layer 7's rows to layer 2's attention. + layer_ids = (1, 3) + pool = _pool(layer_ids=layer_ids) + assert pool._kv_buffer.shape[1] == 2 + 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_latent_kv_pool_rejects_fp8(): + """MLA/DSA (and by the same guard BSA/QSA/DSV4) have no scale-read path: asking + for fp8 must fail loudly, not quietly allocate a 16-bit cache the budget priced + as fp8.""" + from freetoken.attention import AttnType + from freetoken.kvcache import create_kvcache_pool + + spec = KVCacheGroupSpec( + name="full", + layer_ids=(0, 1), + num_kv_heads=HEADS, + head_dim=DIM, + sliding_window=None, + mla=True, + attn_type=AttnType.MLA, + ) + mc = SimpleNamespace( + has_swa_attention=False, + has_linear_attention=False, + num_layers=2, + num_kv_heads=HEADS, + head_dim=DIM, + kv_cache_group_specs=lambda: (spec,), + ) + with pytest.raises(ValueError, match="kv-cache-dtype"): + create_kvcache_pool( + model_config=mc, + num_pages=4, + page_size=1, + dtype=torch.bfloat16, + device=DEV, + kv_quant="fp8", + ) + # The same request on the 16-bit path is fine (guards against an over-eager check). + pool = create_kvcache_pool( + model_config=mc, + num_pages=4, + page_size=1, + dtype=torch.bfloat16, + device=DEV, + kv_quant="none", + ) + assert pool.kv_quant == "none" + + +def test_hybrid_swa_pool_also_separates_compute_and_store_dtype(monkeypatch): + """The hybrid-SWA pool is the other family that stores codes, so it must report the + same pair. Any backend sizing scratch off ``dtype`` would take e4m3 home with it + here too -- what that looked like in practice is QSA's indexer (kvcache/base.py).""" + from freetoken.distributed.info import DistributedInfo + from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache + from freetoken.models.config import KVCacheGroupSpec + + monkeypatch.setattr( + "freetoken.kvcache.hybrid_swa_pool.get_tp_info", + lambda: DistributedInfo(rank=0, size=1), + ) + groups = ( + KVCacheGroupSpec( + name="full", layer_ids=(2, 5), num_kv_heads=2, head_dim=DIM, + sliding_window=None, + ), + KVCacheGroupSpec( + name="swa", layer_ids=(0, 1, 3, 4), num_kv_heads=2, head_dim=DIM, + sliding_window=32, + ), + ) + + def build(kv_quant: str): + return HybridSWAKVCache( + groups=groups, + num_layers=6, + num_full_pages=PAGES, + page_size=PAGE_SIZE, + num_swa_tokens=PAGES * PAGE_SIZE, + dtype=torch.bfloat16, + device=DEV, + kv_quant=kv_quant, + ) + + quantized, plain = build("fp8"), build("none") + assert quantized.dtype is torch.bfloat16 + assert plain.dtype is torch.bfloat16 + assert quantized.store_dtype == kv_codes_dtype() + assert plain.store_dtype is torch.bfloat16 + # ...while the buffers really did shrink: the two properties must not be aliases. + assert quantized.k_cache(2).element_size() == 1 + assert plain.k_cache(2).element_size() == 2 + assert quantized.k_scale(2) is not None and plain.k_scale(2) is None diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py new file mode 100644 index 000000000..13fa7f468 --- /dev/null +++ b/tests/kvcache/test_qsa_pool_fp8.py @@ -0,0 +1,209 @@ +"""The QSA pool under ``--kv-cache-dtype fp8``: quantized K/V, 16-bit index tiers. + +Only the paged K/V slabs change. The compressed index slab, the per-request ring and +their scratch rows must stay 16-bit whatever the KV store does -- block selection scores +against them and the score kernel asserts their dtype -- while the byte account the +startup budget, the cache sliders and the runtime rebuild all divide by has to keep +telling the two halves apart. A drift here does not crash: it silently buys the wrong +number of pages. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.attention import AttnType +from freetoken.kvcache.base import spec_kv_bytes_per_token +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 codes_to_f32, kv_codes_dtype +from freetoken.kvcache.qsa_pool import QSAKVCache + +DEV = torch.device("cuda") +PAGE_SIZE = 64 # the page size the qsa_sparse backend registers +LAYER_IDS = (1, 3, 5, 7) +HEADS, DIM, INDEX_DIM, INDEX_LAYERS, RATIO = 2, 64, 32, 4, 4 + + +@pytest.fixture(autouse=True) +def _tp(monkeypatch): + from freetoken.distributed.info import DistributedInfo + + monkeypatch.setattr( + "freetoken.kvcache.mha_pool.get_tp_info", + lambda: DistributedInfo(rank=0, size=1), + ) + + +def _pool(kv_quant="fp8", num_pages=4, num_req_slots=4): + return QSAKVCache( + num_kv_heads=HEADS, + num_layers=8, + head_dim=DIM, + num_pages=num_pages, + page_size=PAGE_SIZE, + dtype=torch.bfloat16, + device=DEV, + index_head_dim=INDEX_DIM, + num_index_layers=INDEX_LAYERS, + index_ratio=RATIO, + num_req_slots=num_req_slots, + layer_ids=LAYER_IDS, + kv_quant=kv_quant, + ) + + +def _spec(): + return KVCacheGroupSpec( + name="full", + layer_ids=LAYER_IDS, + num_kv_heads=HEADS, + head_dim=DIM, + sliding_window=None, + index_head_dim=INDEX_DIM, + num_index_layers=INDEX_LAYERS, + index_ratio=RATIO, + attn_type=AttnType.QSA, + ) + + +def _config(kv_quant, *, page_size=PAGE_SIZE, max_running_req=3): + mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True) + mc.kv_cache_group_specs = lambda: (_spec(),) + return SimpleNamespace( + model_config=mc, + page_size=page_size, + dtype=torch.bfloat16, + tp_info=SimpleNamespace(size=1), + max_running_req=max_running_req, + kv_quant=kv_quant, + ) + + +def test_fp8_replaces_the_kv_slab_and_adds_scale_views(): + pool = _pool() + bf16 = _pool(kv_quant="none") + assert pool.kv_quant == "fp8" + # Same geometry as the 16-bit pool -- only the element type changed, because the + # attend kernels index codes exactly the way they index values. + assert pool._kv_buffer.shape == bf16._kv_buffer.shape + assert pool._kv_buffer.dtype == kv_codes_dtype() + # dtype = compute dtype (what the backend sizes its INDEXER scratch with), + # store_dtype = what the buffer holds. Swapping them is the bug that compiled an + # e4m3 operand into QSA's scoring dot and died at graph capture. + assert pool.dtype is torch.bfloat16 + assert pool.store_dtype == kv_codes_dtype() + assert bf16.dtype is torch.bfloat16 and bf16.store_dtype is torch.bfloat16 + assert pool.k_cache(3).shape == (4, PAGE_SIZE, HEADS, DIM) + assert pool.k_cache(3).element_size() == 1 + slots = 4 * PAGE_SIZE + assert pool.k_scale(3).shape == (slots, HEADS) + assert pool.k_scale(3).dtype is torch.float32 + assert pool.v_scale(3).shape == (slots, HEADS) + # Zero-filled: e4m3 has NaN bit patterns, and the dummy page / unwritten tail rows + # must never read back as one. + assert pool.k_scale(3).abs().sum().item() == 0.0 + assert codes_to_f32(pool.k_cache(3)).abs().sum().item() == 0.0 + # A 16-bit pool keeps answering None, so the backends' k_scale(...) pass-through is + # the only branch that ever differs between the two. + assert bf16.k_scale(3) is None and bf16.v_scale(3) is None + + +@pytest.mark.parametrize("kv_quant", ["none", "fp8"]) +def test_index_tiers_stay_16_bit_whatever_the_kv_store_does(kv_quant): + """Block selection is quantization-agnostic by construction: it reads the compressed + index keys, not the KV rows, so fp8 must not touch these three buffers.""" + pool = _pool(kv_quant=kv_quant) + assert pool.cmp_k_cache(0).dtype is torch.bfloat16 + assert pool.pending_ring(0).dtype is torch.bfloat16 + assert pool.cmp_k_cache(0).shape == (4 * PAGE_SIZE // RATIO + 4, INDEX_DIM) + # ...and the byte account says so too: only the 1-byte KV codes got cheaper. + spec = _spec() + cost = spec_kv_bytes_per_token(spec, _config(kv_quant)) + plain = spec_kv_bytes_per_token(spec, _config("none")) + kv_layers = len(LAYER_IDS) + kv_16bit = 2 * DIM * HEADS * 2 * kv_layers # two slabs, 16-bit codes + scale_term = 2 * kv_layers * HEADS * 4 # one fp32 scale per (slot, head) + index_term = INDEX_DIM * INDEX_LAYERS * 2 // RATIO # the untouched 16-bit slab + assert plain == kv_16bit + index_term + if kv_quant == "fp8": + assert cost == kv_16bit // 2 + scale_term + index_term + else: + assert cost == plain + + +def test_unit_bytes_and_kv_cost_still_agree_when_quantized(): + """The pool's own allocation and the budget model must divide the same way -- the + scale sidecar is priced in base.spec_kv_bytes_per_token, not here.""" + spec, config = _spec(), _config("fp8") + pool = _pool() + kv_bytes, swa_bytes = pool.unit_bytes() + assert swa_bytes == 0 + assert kv_bytes == spec_kv_bytes_per_token(spec, config) + assert kv_bytes * PAGE_SIZE == QSAKVCache.kv_cost(config)[0] + # The 16-bit pool's per-token figure is the reference: codes halve the KV term, the + # scale sidecar and the untouched index slab keep the total above a clean half. + plain_kv = spec_kv_bytes_per_token(spec, _config("none")) + assert plain_kv // 2 < kv_bytes < plain_kv + + +def test_rebuild_resizes_codes_and_scales_together(): + pool = _pool(num_pages=4) + ident = id(pool) + pool.rebuild(16) + assert id(pool) == ident + assert pool.k_cache(1).shape == (16, PAGE_SIZE, HEADS, DIM) + assert pool.k_scale(1).shape == (16 * PAGE_SIZE, HEADS) + assert pool.k_scale(1).abs().sum().item() == 0.0 + assert pool.cmp_k_cache(0).shape == (16 * PAGE_SIZE // RATIO + 4, INDEX_DIM) + + +def test_store_kv_writes_the_slot_the_attend_kernel_will_read(): + """out_loc numbering (page * page_size + offset) is the contract between the fused + writer and the attend kernel's scale slot arithmetic -- this is that round trip.""" + torch.manual_seed(0) + pool = _pool(num_pages=4) + slots = 4 * PAGE_SIZE + rows = (0, 1, 63, 64, 255, 256) # page boundaries included: 63/64 and 255/256 + k = torch.randn(len(rows), HEADS * DIM, device=DEV, dtype=torch.bfloat16) * 3.0 + v = torch.randn(len(rows), HEADS * DIM, device=DEV, dtype=torch.bfloat16) * 0.25 + out_loc = torch.tensor(rows, dtype=torch.int32, device=DEV) + + pool.store_kv(k, v, out_loc, layer_id=3) + torch.cuda.synchronize() + + codes = codes_to_f32(pool.k_cache(3).view(slots, HEADS, DIM)) + decoded_k = codes[out_loc] * pool.k_scale(3)[out_loc].unsqueeze(-1) + decoded_v = codes_to_f32(pool.v_cache(3).view(slots, HEADS, DIM))[out_loc] * pool.v_scale(3)[ + out_loc + ].unsqueeze(-1) + for want, got in ((k, decoded_k), (v, decoded_v)): + w = want.view(len(rows), HEADS, DIM).to(torch.float32) + amax = w.abs().amax(dim=-1, keepdim=True) + assert (((got - w).abs() / amax).max().item()) < 0.07 + # Rows nobody wrote stay zero rather than NaN -- the dummy page depends on it. + untouched = torch.tensor([r for r in range(64) if r not in rows], device=DEV) + assert pool.k_scale(3)[untouched].abs().sum().item() == 0.0 + + +def test_factory_threads_kv_quant_into_the_qsa_pool(): + from freetoken.kvcache import create_kvcache_pool + + mc = SimpleNamespace( + num_layers=8, has_swa_attention=False, has_linear_attention=True, + num_kv_heads=HEADS, head_dim=DIM, dsv4_args=None, + ) + mc.kv_cache_group_specs = lambda: (_spec(),) + pool = create_kvcache_pool( + mc, num_pages=4, page_size=PAGE_SIZE, dtype=torch.bfloat16, device=DEV, + num_req_slots=4, kv_quant="fp8", + ) + assert isinstance(pool, QSAKVCache) and pool.kv_quant == "fp8" + assert pool.k_cache(1).element_size() == 1 and pool.k_scale(1) is not None + diff --git a/tests/models/qwen4_exp/common.py b/tests/models/qwen4_exp/common.py index 1f9c117bd..cfd1e8d35 100644 --- a/tests/models/qwen4_exp/common.py +++ b/tests/models/qwen4_exp/common.py @@ -154,6 +154,7 @@ def __init__( device: str = "cuda", dtype: torch.dtype = torch.bfloat16, page_size: int = 64, + kv_quant: str = "none", ) -> None: from freetoken.attention.qsa_sparse import QSASparseAttnBackend from freetoken.kvcache import create_kvcache_pool @@ -170,6 +171,7 @@ def __init__( dtype=dtype, device=self.device, num_req_slots=self.num_req_slots, + kv_quant=kv_quant, ) self.page_table = torch.zeros( (self.num_req_slots, num_pages * page_size), dtype=torch.int32, device=self.device diff --git a/tests/models/qwen4_exp/test_qsa_backend.py b/tests/models/qwen4_exp/test_qsa_backend.py index 1d3b944ce..742ca00e4 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -5,7 +5,10 @@ exactly the causal prefix and the layer output must match ``TorchDenseQSAReference`` (fp32) and a flashinfer dense run over the same pool; (b) chunked prefill at unaligned cut points equals one-shot prefill (the dual-source compress); -(c) a captured decode replay equals the eager decode step. +(c) a captured decode replay equals the eager decode step; +(d) an fp8 KV pool (``--kv-cache-dtype fp8``) keeps block selection bit-identical to the + 16-bit run -- only the selected K/V rows are read back as e4m3 codes -- and the layer + output stays within quantization error of it. """ from __future__ import annotations @@ -248,3 +251,64 @@ def test_two_qsa_layers_keep_separate_slab_slots(monkeypatch): slab = fixture.pool.cmp_k_cache assert not torch.equal(slab(0), slab(1)) + + +def _prefill_under_kv(monkeypatch, config, kv_quant: str, lengths): + """One prefill of the QSA layer under a given KV store. + + Each call builds its own Fixture on purpose: a Fixture owns the global ctx (pool, + page table, backend), so two KV stores cannot share one scenario. The weight seed + (``Fixture.layer``) and the input seed (``_inputs``) are fixed, so the two runs differ + ONLY in how the K/V rows are stored. + """ + fixture = Fixture(config, num_pages=128, kv_quant=kv_quant) + attn = fixture.layer(QSA_LAYER) + seen = selection_spy(monkeypatch, fixture.backend) + inputs = _inputs(fixture, lengths) + x = torch.cat([row[:n] for row, n in zip(inputs, lengths)]) + reqs = [fixture.req(i, 0, n) for i, n in enumerate(lengths)] + batch = fixture.batch(reqs, "prefill") + out = attn.forward(x, batch) + # the selection lives in a scratch buffer the next forward overwrites + return fixture, out.clone(), seen["indices"].clone(), batch.positions.clone() + + +@requires_cuda +def test_fp8_kv_pool_keeps_selection_and_output(monkeypatch): + """--kv-cache-dtype fp8 through the real layer: e4m3 codes + per-row scales in, same + answer out to within quantization error -- and, because block selection scores 16-bit + compressed index keys that fp8 never touches, the SAME selection bit for bit.""" + config = parsed_config() + lengths = [2051, 1000, 137] # every complete block is selected here + + plain, plain_out, plain_idx, _ = _prefill_under_kv(monkeypatch, config, "none", lengths) + quant, quant_out, quant_idx, positions = _prefill_under_kv( + monkeypatch, config, "fp8", lengths + ) + + # The tripwire for the field failure: the backend sizes its indexer scratch with + # pool.dtype, which must stay the COMPUTE dtype even when store_dtype is e4m3. An + # fp8 q_index compiles into qsa_mqa_paged's dot and dies at graph capture. + assert quant.backend.dtype is torch.bfloat16 + assert plain.backend.dtype is torch.bfloat16 + assert quant.pool.store_dtype != torch.bfloat16 + assert quant.pool.kv_quant == "fp8" and plain.pool.kv_quant == "none" + assert quant.pool.k_cache(QSA_LAYER).element_size() == 1 + assert quant.pool.v_cache(QSA_LAYER).element_size() == 1 + assert plain.pool.k_scale(QSA_LAYER) is None and plain.pool.v_scale(QSA_LAYER) is None + pages, page_size, kv_heads = quant.pool.k_cache(QSA_LAYER).shape[:3] + assert quant.pool.k_scale(QSA_LAYER).shape == (pages * page_size, kv_heads) + assert quant.pool.k_scale(QSA_LAYER).dtype is torch.float32 + + for pool in (plain.pool, quant.pool): + assert pool.cmp_k_cache(0).dtype is torch.bfloat16 + assert torch.equal(quant_idx, plain_idx), ( + "quantizing the KV rows changed which blocks the indexer selected -- the index " + "tier is supposed to be 16-bit in both runs" + ) + _assert_selection_is_causal_prefix(quant_idx, positions) + + # Looser than the 2e-2 the 16-bit run needs against the same reference: e4m3 carries + # four significant bits, so ~1e-2 relative per stored element is the floor here. + torch.testing.assert_close(quant_out.float(), plain_out.float(), rtol=4e-2, atol=4e-2) +