From 03fb0438bc3585eeacb80795517659b451182701 Mon Sep 17 00:00:00 2001 From: ArqAlice Date: Wed, 2 Sep 2026 22:07:25 +0900 Subject: [PATCH 01/12] feat(kvcache): store the KV cache as fp8 e4m3 codes (--kv-cache-dtype fp8) One (token, kv head) row of K and of V becomes head_dim e4m3 codes plus ONE fp32 symmetric scale, in a code buffer with exactly the geometry of the 16-bit KV buffer -- only the element type changes. That halves the bytes per cached token (the scale sidecar costs 4/head_dim of it back, ~3% at head_dim 128), and it is what lets Qwen3.8-Flash-Next serve a 1M-token context on this card. Codes are kept in a plain uint8 buffer on EVERY architecture, and the fp8e4nv type never appears in a kernel signature. Both ways of choosing that per target failed on real hardware and are recorded here so nobody reopens them: the compile-time fp8-native probe (e4m3_compat.e4m3_native_cx) answers the question independently from the host that allocated the buffer and disagreed with it on sm_100, and branching on a pointer's element type is NOT statically pruned -- triton still type-checked the dead arm, whose int mask fill is illegal against an fp8 pointer ("cannot cast int32 to fp8e4nv", raised at CUDA graph capture). What remains is the software encode/decode that already runs wherever the fp8 type is unavailable and is bit-exact per e4m3_compat's header, so the cache holds the same bytes and produces the same numbers on every card (docs/cli.md). - server/args.py, engine/config.py: --kv-cache-dtype {auto,bf16,fp8}, refused at startup for the pools and backends that cannot apply the row scales (attention/__init__.py: BackendInfo.supports_fp8_kv) rather than ignored. - kernel/triton/kv_quant.py: fused quantize+scatter -- one launch under CUDA graph capture, where the slot ids arrive as a device tensor. - kvcache: unit_bytes() counts codes plus the scale sidecar, so ft ctl stats and cache --kv N follow the smaller footprint, and rebuild reallocates the scale buffers alongside the codes (mha, hybrid-SWA and QSA pools). - kvcache/base.py: pool.dtype is the COMPUTE dtype -- what store_kv receives and what a backend sizes its scratch with -- while pool.store_dtype is what the buffer holds. Reporting codes as dtype handed e4m3 to QSA's 16-bit indexer and died compiling qsa_mqa_paged; the contract is now asserted at backend init and in the kernel wrapper. QSA's block-selection keys stay 16-bit: only the selected K/V rows are read back as codes. Tested on: sm_100, 148 SMs, Linux; 524,480 fp8 KV tokens = 6.47 GiB, Qwen3.8-Flash-Next with: ft serve --kv-cache-dtype fp8 -> 1M-token context. Covered by tests/kernels/test_kv_fp8.py, tests/kernels/test_qsa_fp8.py, tests/kernels/test_triton_attention.py, tests/kernels/test_e4m3_compat.py, tests/kvcache/test_mha_pool_fp8.py, tests/kvcache/test_qsa_pool_fp8.py and tests/engine/test_kv_quant_config.py (CUDA-gated; not run on the Windows development box, which has neither triton nor pytest installed). Not included here, on purpose: unifying the two fp8-native probes (triton's cache-key walk rejects a constexpr function that defers to a host one, so warn_if_probes_disagree() reports the disagreement instead), and a hardware decode fast path on sm_89+ (that needs a constexpr flag threaded from the host plus the matching AOT variants, since testing the dtype does not prune). --- docs/cli.md | 24 ++ docs/models.md | 6 + python/freetoken/attention/__init__.py | 9 + python/freetoken/attention/qsa_sparse.py | 14 + python/freetoken/attention/triton.py | 11 + python/freetoken/engine/config.py | 6 + python/freetoken/engine/engine.py | 84 ++++- python/freetoken/kernel/aot_models.py | 26 +- python/freetoken/kernel/triton/attention.py | 317 ++++++++++++++---- python/freetoken/kernel/triton/e4m3_compat.py | 120 ++++++- python/freetoken/kernel/triton/kv_quant.py | 194 +++++++++++ python/freetoken/kernel/triton/qsa/attend.py | 117 +++++-- python/freetoken/kernel/triton/qsa/score.py | 8 + python/freetoken/kvcache/__init__.py | 22 ++ python/freetoken/kvcache/base.py | 78 ++++- python/freetoken/kvcache/hybrid_swa_pool.py | 154 +++++++-- python/freetoken/kvcache/mha_pool.py | 110 +++++- python/freetoken/kvcache/qsa_pool.py | 19 +- python/freetoken/server/args.py | 15 + tests/engine/test_kv_quant_config.py | 184 ++++++++++ tests/kernels/test_e4m3_compat.py | 83 +++++ tests/kernels/test_kv_fp8.py | 262 +++++++++++++++ tests/kernels/test_qsa_fp8.py | 278 +++++++++++++++ tests/kernels/test_triton_attention.py | 257 ++++++++++++++ tests/kvcache/test_mha_pool_fp8.py | 271 +++++++++++++++ tests/kvcache/test_qsa_pool_fp8.py | 209 ++++++++++++ tests/models/qwen4_exp/common.py | 2 + tests/models/qwen4_exp/test_qsa_backend.py | 66 +++- 28 files changed, 2791 insertions(+), 155 deletions(-) create mode 100644 python/freetoken/kernel/triton/kv_quant.py create mode 100644 tests/engine/test_kv_quant_config.py create mode 100644 tests/kernels/test_kv_fp8.py create mode 100644 tests/kernels/test_qsa_fp8.py create mode 100644 tests/kvcache/test_mha_pool_fp8.py create mode 100644 tests/kvcache/test_qsa_pool_fp8.py 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 41b79cd68..03521fd87 100644 --- a/docs/models.md +++ b/docs/models.md @@ -40,3 +40,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 72b01c047..04353a94f 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): @@ -132,6 +138,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 543012f39..6b7348f3c 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -57,6 +57,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 73dc7688d..85ddb0547 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. @@ -1296,6 +1349,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} @@ -1323,7 +1397,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 6268ae338..cbcd5f027 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 @@ -377,8 +380,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]]: @@ -400,9 +404,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..6866316a2 --- /dev/null +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -0,0 +1,194 @@ +"""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, # source row pitch, in elements (the qkv slice is wider than one row) + 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 + + src = t * stride_xs + h * D + d + xk = tl.load(k_src + src, mask=mask, other=0.0).to(tl.float32) + xv = tl.load(v_src + src, 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), + 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 c6c0f1bb9..e2c47b2fa 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -70,6 +70,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 @@ -79,10 +90,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, @@ -111,6 +124,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, ) @@ -122,6 +136,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 @@ -134,6 +149,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] return BSAKVCache( num_kv_heads=spec.num_kv_heads, @@ -203,11 +220,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, MLAKVCache + _reject_unsupported_quant("latent-KV (MLA/DSA)", kv_quant) spec = kv_specs[0] if spec.index_head_dim > 0 and spec.num_index_layers > 0: return DSAKVCache( @@ -238,6 +259,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 6696f65dd..2c92568dc 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) + From 811ccee2b24f86a5e282311976180dbe92acb4be Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Fri, 4 Sep 2026 11:02:14 +0900 Subject: [PATCH 02/12] fix(kernels): give the V tensor its own row pitch in the fp8 KV store `quantize_kv_to_cache` passed `k.stride(0)` as the only source pitch and `_kv_quant_scatter_kernel` used it for both tensors: src = t * stride_xs + h * D + d xk = tl.load(k_src + src, ...) xv = tl.load(v_src + src, ...) The guard above it checks only the inner stride (`k.stride(1) == 1 and v.stride(1) == 1`), never `k.stride(0) == v.stride(0)`, so the kernel carries an undocumented contract: K and V must share one row pitch. When they do not, V is read at K's pitch. In the failing test K is a view of the qkv slice (pitch 1152) while V is materialised by `.clamp()` (pitch 384), so with 8 tokens of 3072 elements: token 0 reads 0 correct by coincidence token 1-2 reads 1152, 2304 in range, WRONG rows token 3-7 reads 3456 .. 8064 past the initialised data 2684 of 3072 codes wrong, all in V, K byte-perfect. Deterministic addressing; only the contents of the uninitialised tail vary with allocator history, which is why the mismatch count drifts (2684 / 2663 / 2676 across runs) while the mismatching positions do not -- the in-range half is exactly 764 every time. Found with `compute-sanitizer --tool initcheck` (TRITON_DISABLE_LINE_INFO=0), which named `kv_quant.py:110`. `memcheck` reports 0 errors because PyTorch's caching allocator rounds allocations up and the bad read stays inside the pooled segment; `racecheck` reports 0 hazards because it is not a race. Fix: pass `v.stride(0)` as its own kernel argument and load each tensor with its own pitch. Verified on RTX 4090 (sm_89): tests/kernels/test_kv_fp8.py 2 failed -> 1 failed, the flip being test_codes_match_the_reference_quantizer_and_reconstruction_is_close; five consecutive standalone runs give K 0/3072 and V 0/3072 with got.sort() == exp.sort(); compute-sanitizer initcheck reports 0 errors on the patched build. Independently confirmed on RTX 5090 D (sm_120) by @Kaempferia: same single flip, same multiset property, 5 runs clean. Note for reviewers: the test's SECOND assertion (dequantised error <= 0.08) passes at 0.035 while V is wrong, so a reconstruction-level check does not catch this class. Only the exact-code assertion does. Co-Authored-By: Claude Opus 5 --- python/freetoken/kernel/triton/kv_quant.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py index 6866316a2..e041e0466 100644 --- a/python/freetoken/kernel/triton/kv_quant.py +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -89,7 +89,9 @@ def _kv_quant_scatter_kernel( k_scale, v_scale, idx_ptr, - stride_xs, # source row pitch, in elements (the qkv slice is wider than one row) + 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) @@ -105,9 +107,9 @@ def _kv_quant_scatter_kernel( d = tl.arange(0, BLOCK_D) mask = d < D - src = t * stride_xs + h * D + d - xk = tl.load(k_src + src, mask=mask, other=0.0).to(tl.float32) - xv = tl.load(v_src + src, mask=mask, other=0.0).to(tl.float32) + 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). @@ -172,6 +174,7 @@ def quantize_kv_to_cache( v_scale, out_loc, k.stride(0), + v.stride(0), k_cache.stride(0), v_cache.stride(0), k_scale.stride(0), From 7f9a05a7f85e2d3d8f31c385a8d3feda0c4a42b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Sat, 5 Sep 2026 22:17:04 +0900 Subject: [PATCH 03/12] test(kvcache): size the fp8 slot round-trip to the rows it indexes `test_store_kv_writes_the_slot_the_attend_kernel_will_read` reaches row 256 -- its comment asks for the 255/256 page boundary -- against a four-page, 256-slot pool, so `codes[out_loc]` gathers one row past the view. Five pages is what that row list needs. The gather raises a device-side assert, and the CUDA context does not recover from it, so in a single-process run everything scheduled afterwards is reported as failing as well -- not because those tests stop working, but because there is no longer a context to run them in. One file per process does not show that. Assisted-by: Claude Opus 5 --- tests/kvcache/test_qsa_pool_fp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py index 13fa7f468..04426b876 100644 --- a/tests/kvcache/test_qsa_pool_fp8.py +++ b/tests/kvcache/test_qsa_pool_fp8.py @@ -168,8 +168,8 @@ 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 + pool = _pool(num_pages=5) + slots = 5 * 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 From dabe93ccd18fdf110d02088a93a754d16adfcd00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Sat, 5 Sep 2026 21:02:36 +0900 Subject: [PATCH 04/12] test(kernels): put the scale-one encoder's V tensor on the device `test_encoder_inverts_the_grid_through_the_scale_one_path` moves `rows` to the device and leaves the V tensor beside it on the host, so the kernel is handed a CPU pointer. Assisted-by: Claude Opus 5 --- tests/kernels/test_kv_fp8.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/kernels/test_kv_fp8.py b/tests/kernels/test_kv_fp8.py index 26deef2ac..817646886 100644 --- a/tests/kernels/test_kv_fp8.py +++ b/tests/kernels/test_kv_fp8.py @@ -179,7 +179,8 @@ def test_encoder_inverts_the_grid_through_the_scale_one_path(): expected[t, j + 1] = code k_cache, _, k_scale, _ = _store( - rows.to(DEV, dtype=torch.bfloat16), torch.zeros(tokens, 1, dim, dtype=torch.bfloat16) + rows.to(DEV, dtype=torch.bfloat16), + torch.zeros(tokens, 1, dim, dtype=torch.bfloat16, device=DEV), ) assert torch.equal(k_scale[:, 0], torch.ones_like(k_scale[:, 0])) got = _canonical_zero(_as_bytes(k_cache)[:, 0, :]) From 5febeeedcffbdda105493461070683c93afa0075 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Sat, 5 Sep 2026 21:08:05 +0900 Subject: [PATCH 05/12] test(kvcache): give the layer-ids remap test a model deep enough for its ids `test_layer_ids_remap_applies_to_scales_too` backs `layer_ids=(1, 3)` on a pool built with `num_layers=LAYERS`, and LAYERS is 3, so id 3 is one past the end and the constructor raises before the assertion it is there to make. The helper now takes the depth. Assisted-by: Claude Opus 5 --- tests/kvcache/test_mha_pool_fp8.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/kvcache/test_mha_pool_fp8.py b/tests/kvcache/test_mha_pool_fp8.py index 77bd2986e..dd3f0542d 100644 --- a/tests/kvcache/test_mha_pool_fp8.py +++ b/tests/kvcache/test_mha_pool_fp8.py @@ -35,13 +35,13 @@ def _tp(): _init_tp() -def _pool(kv_quant="fp8", num_pages=PAGES, layer_ids=None): +def _pool(kv_quant="fp8", num_pages=PAGES, layer_ids=None, num_layers=LAYERS): from freetoken.kvcache.mha_pool import MHAKVCache _init_tp() return MHAKVCache( num_kv_heads=HEADS, - num_layers=LAYERS, + num_layers=num_layers, head_dim=DIM, num_pages=num_pages, page_size=PAGE_SIZE, @@ -84,8 +84,10 @@ def test_fp8_pool_keeps_geometry_and_adds_scale_views(): 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. + # (1, 3) names a subset of a model, so layer id 3 has to be inside it: LAYERS is 3, which + # makes 3 one past the end, so this pool is told the depth the ids imply. layer_ids = (1, 3) - pool = _pool(layer_ids=layer_ids) + pool = _pool(layer_ids=layer_ids, num_layers=4) assert pool._kv_buffer.shape[1] == 2 with pytest.raises(KeyError): pool.k_scale(0) From 0820ff4e9b5a774940f34b74b2a79c1ca2fa8013 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matsumoto=20Takaya=20=28=E6=9D=BE=E6=9C=AC=20=E8=B2=B4?= =?UTF-8?q?=E4=B9=9F=29?= Date: Sat, 5 Sep 2026 21:09:48 +0900 Subject: [PATCH 06/12] test(kernels): give the triton-attention doubles the scale accessors the backend now reads `k_scale` / `v_scale` arrived with the fp8 store, and the backend reads them on every path, including a 16-bit pool -- which answers None. The two hand-rolled `FakeKVCache` classes in this file do not inherit the base pool, so they were left without them and raise `AttributeError` instead. Assisted-by: Claude Opus 5 --- tests/kernels/test_triton_attention.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index ffff4a09b..61ce5360d 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -67,6 +67,14 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + # A 16-bit pool answers None here. The backend reads it on every path since the + # fp8 store landed, so the double answers it too. + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + kv_cache = FakeKVCache() monkeypatch.setattr( "freetoken.attention.triton.get_global_ctx", @@ -636,6 +644,14 @@ def k_cache(self, layer_id): def v_cache(self, layer_id): return self.v + # A 16-bit pool answers None here. The backend reads it on every path since the + # fp8 store landed, so the double answers it too. + def k_scale(self, layer_id): + return None + + def v_scale(self, layer_id): + return None + device = torch.device("cuda") head_dim = 256 page_table = torch.tensor([[0, 1], [2, 3]], dtype=torch.int32, device=device) From 05861fb312781c26027e4b82bacc9bb29033f04f Mon Sep 17 00:00:00 2001 From: Vincent Labreche Date: Sun, 6 Sep 2026 21:01:44 +0000 Subject: [PATCH 07/12] perf(kernel): apply the fp8 KV dequant scale after the dot, not to the tile The per-(token, kv_head) dequant scale is constant down each dot's reduction dim, so it never has to touch K or V: scores[m,n] = (sum_d q[m,d] * k[d,n]) * s_k[n] p @ (diag(s_v) @ v) = (p * s_v[None,:]) @ v Scaling the BLOCK_M x BLOCK_N result instead of the BLOCK_D x BLOCK_N K tile and the BLOCK_N x BLOCK_DV V tile is head_dim/BLOCK_M fewer multiplies. p itself stays unscaled, since l_i accumulates it as the softmax denominator and knows nothing about V's quantization. That also removes the only reason the tile was widened to fp32. kv_load_e4m3_tile_f32 builds an fp16 bit pattern and widens purely so a * 256.0 can put the value back on the true e4m3 scale, and 2^8 is a power of two, so once the scale rides the dot output it folds into that scale exactly. kv_load_e4m3_tile_scaled16 stops before the widen and leaves the fold to the caller, keeping the tile 16-bit through the whole loop. Accuracy improves rather than degrades. The general scale used to multiply before the narrow to the compute dtype, so the product rounded; now the tile reaches the dot exactly (the code's own 3 mantissa bits, |x| <= 1.75) and the scale is applied in fp32 afterwards. Worst-case absolute error in test_extend_paged_attention_decodes_fp8_scales drops 0.281 -> 0.0996 on sm_86. (That test still exceeds its 2e-2 tolerance on this card both before and after -- it fails on 3e5bbdd unpatched too, so it is not introduced here.) The loader's bit placement is also the same number in 4 ops instead of 7: for v = 128s + r, ((v & 0x80) << 8) | ((v & 0x7F) << 7) and (v + (v & 0x80)) << 7 are both (256s + r) << 7. Verified identical on all 256 codes, NaN patterns included, by the new test in tests/kernels/test_e4m3_compat.py. RTX 3070 (sm_86, 8GB, driver 610.57.04), i7-11700KF, Qwen3.6-35B-A3B-NVFP4, --moe-backend hybrid --kv-cache-dtype fp8 --max-seq-len-override 180000 --memory-ratio 0.9 --max-running-requests 1 --max-prefill-length 1024, 2 reps, median, 127 output tokens, unique nonce per request: ctx TTFT 3e5bbdd -> here decode 3e5bbdd -> here 33k 34.13 -> 33.95 s 47.04 -> 48.05 t/s 65k 88.79 -> 85.03 s 38.99 -> 41.61 t/s 100k 168.27 -> 160.42 s 33.59 -> 37.13 t/s Most of the prefill win needs the tile-sizing fix in the next commit; this one is mainly a decode gain on its own. Co-Authored-By: Claude Opus 5 --- python/freetoken/kernel/triton/attention.py | 111 +++++++++++------- python/freetoken/kernel/triton/e4m3_compat.py | 30 +++++ tests/kernels/test_e4m3_compat.py | 43 +++++++ 3 files changed, 140 insertions(+), 44 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index 31ae84847..6166e56fe 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -6,7 +6,31 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import KV_TILE_SCALE from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 as _kv_load_f32 +from freetoken.kernel.triton.e4m3_compat import ( + kv_load_e4m3_tile_scaled16 as _kv_load_s16, +) + + +@triton.jit +def _kv_dequant_scale(scale_ptr, slots, stride, kv_head, mask): + """Per-(token, kv_head) dequant scale for an fp8 KV tile, pre-multiplied by the + 2**8 that :func:`kv_load_e4m3_tile_scaled16` leaves on the tile it returns. + + Applied to a dot's OUTPUT rather than to K or V. The scale is constant down the + reduction dim, so ``scores[m,n] = (sum_d q[m,d]*k[d,n]) * s_k[n]`` and + ``p @ (diag(s_v) @ v) = (p * s_v[None,:]) @ v`` both hold: scaling the + ``BLOCK_M x BLOCK_N`` result costs less than scaling the ``BLOCK_D x BLOCK_N`` K + tile or the ``BLOCK_N x BLOCK_DV`` V tile, by a factor of head_dim / BLOCK_M. + + It is also the more accurate order, which is why the 16-bit tile is safe: every + value the loader returns carries at most the code's own 3 mantissa bits and lies + in +-1.75, so narrowing it to the compute dtype is lossless, whereas multiplying + by a general scale first and narrowing after rounds the product. + """ + s = tl.load(scale_ptr + slots * stride + kv_head, mask=mask, other=0.0) + return s * KV_TILE_SCALE _MAX_KV_SPLITS = 8 @@ -25,10 +49,11 @@ def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[i Larger tiles run materially faster (~2x for head_dim 512 on H100) but their bf16 q/k/v tiles need about ``(BLOCK_M + 2 * BLOCK_N) * BLOCK_D * 2`` bytes of shared - memory, which overflows consumer GPUs (sm_89 ~99KB opt-in) once head_dim >= 256. - Keep the fast tiles where the device's opt-in shared memory fits them (datacenter - A100/H100); shrink only where it does not. ``smem_optin == 0`` (unknown) conservatively - selects the small tiles, i.e. the prior consumer-safe behavior. + memory, which overflows consumer GPUs (sm_89 ~99KB opt-in) once head_dim >= 256. Keep the fast tiles where the device's opt-in shared memory fits + them (datacenter A100/H100); shrink only where it does not. ``smem_optin == 0`` + (unknown) conservatively selects the small tiles, i.e. the prior consumer-safe + behavior. + """ budget = smem_optin * 0.8 # headroom for scores/acc/alignment/triton scratch @@ -38,7 +63,10 @@ def fits(block_m: int, block_n: int) -> bool: if head_dim <= 128: return 128, 64 if head_dim <= 256: - return (128, 64) if fits(128, 64) else (64, 32) + if fits(128, 64): + return 128, 64 + # Reachable by an fp8 cache where a 16-bit one falls through to 64x32. + return (64, 64) if fits(64, 64) else (64, 32) if head_dim <= 384: return (32, 64) if fits(32, 64) else (32, 32) return (32, 64) if fits(32, 64) else (16, 16) @@ -278,14 +306,11 @@ def _decode_grouped_stage1_kernel( slots = tl.load(indices_ptr + kv_start + logical_offs, mask=mask_n, other=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( + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16( k_ptr + slots[None, :] * stride_ks + k_base_offsets, mask_n[None, :] & mask_d[:, None], - ) - k = (k * s_k[None, :]).to(q.dtype) + ).to(q.dtype) else: k = tl.load( k_ptr + slots[None, :] * stride_ks + k_base_offsets, @@ -293,17 +318,16 @@ def _decode_grouped_stage1_kernel( other=0.0, ) scores = tl.dot(q, k) * sm_scale + if HAS_KV_SCALE: + scores = scores * s_k[None, :] scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) 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( + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16( v_ptr + slots[:, None] * stride_vs + v_base_offsets, mask_n[:, None] & mask_dv[None, :], - ) - v = (v * s_v[:, None]).to(q.dtype) + ).to(q.dtype) else: v = tl.load( v_ptr + slots[:, None] * stride_vs + v_base_offsets, @@ -314,7 +338,10 @@ def _decode_grouped_stage1_kernel( m_new = tl.maximum(tl.max(scores, axis=1), m_i) alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + # p itself stays unscaled: l_i is the softmax denominator and knows + # nothing about V's quantization. + pv = (p * s_v[None, :]) if HAS_KV_SCALE else p + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -638,17 +665,14 @@ def _extend_attention_kernel( if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=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( + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16( k_ptr + slots[None, :] * stride_ks + kv_head * stride_kh + offs_d[:, None], mask_n[None, :] & mask_d[:, None], - ) - k = (k * s_k[None, :]).to(q.dtype) + ).to(q.dtype) else: k = tl.load( k_ptr @@ -659,6 +683,8 @@ def _extend_attention_kernel( other=0.0, ) scores = tl.dot(q.to(k.dtype), k) * sm_scale + if HAS_KV_SCALE: + scores = scores * s_k[None, :] scores = tl.where(final_mask, scores, -float("inf")) row_max = tl.max(scores, axis=1) @@ -668,17 +694,14 @@ def _extend_attention_kernel( p = tl.exp(scores - m_new[:, None]) 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( + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16( v_ptr + slots[:, None] * stride_vs + kv_head * stride_vh + offs_dv[None, :], mask_n[:, None] & mask_dv[None, :], - ) - v = (v * s_v[:, None]).to(q.dtype) + ).to(q.dtype) else: v = tl.load( v_ptr @@ -688,7 +711,9 @@ def _extend_attention_kernel( mask=mask_n[:, None] & mask_dv[None, :], other=0.0, ) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + # p stays unscaled for the l_i denominator below. + pv = (p * s_v[None, :]) if HAS_KV_SCALE else p + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new @@ -796,17 +821,14 @@ def _extend_attention_split_kernel( if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=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( + s_k = _kv_dequant_scale(k_scale_ptr, slots, stride_kss, kv_head, mask_n) + k = _kv_load_s16( 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) + ).to(q.dtype) else: k = tl.load( k_cache_ptr @@ -817,6 +839,8 @@ def _extend_attention_split_kernel( other=0.0, ) scores = tl.dot(q.to(k.dtype), k) * sm_scale + if HAS_KV_SCALE: + scores = scores * s_k[None, :] scores = tl.where(final_mask, scores, -float("inf")) row_max = tl.max(scores, axis=1) @@ -826,17 +850,14 @@ def _extend_attention_split_kernel( p = tl.exp(scores - m_new[:, None]) 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( + s_v = _kv_dequant_scale(v_scale_ptr, slots, stride_vss, kv_head, mask_n) + v = _kv_load_s16( 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) + ).to(q.dtype) else: v = tl.load( v_cache_ptr @@ -846,7 +867,9 @@ def _extend_attention_split_kernel( mask=mask_n[:, None] & mask_dv[None, :], other=0.0, ) - acc = acc * alpha[:, None] + tl.dot(p.to(v.dtype), v) + # p stays unscaled for the l_i denominator below. + pv = (p * s_v[None, :]) if HAS_KV_SCALE else p + acc = acc * alpha[:, None] + tl.dot(pv.to(v.dtype), v) l_i = l_i * alpha + tl.sum(p, axis=1) m_i = m_new diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 532c1d7ad..a23b7cde1 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -232,3 +232,33 @@ def kv_load_e4m3_tile_f32(ptrs, mask): QSA sparse one read the same pool, hence this helper lives here. """ return e4m3_u8_to_f32(tl.load(ptrs, mask=mask, other=0)) + + +# What kv_load_e4m3_tile_scaled16 leaves on its tile for the caller to repay. +# tl.constexpr, not a plain float: a @jit function that references a module-level +# python float fails triton's cache-key AST walk with a CompilationError, the same +# way this module's header describes for @constexpr_function and host functions. +# Host-side callers (tests) want ``KV_TILE_SCALE.value``. +KV_TILE_SCALE = tl.constexpr(256.0) + + +@jit +def kv_load_e4m3_tile_scaled16(ptrs, mask): + """Load a tile of KV e4m3 codes as fp16 holding the value times 1/KV_TILE_SCALE. + + :func:`kv_load_e4m3_tile_f32` without its last two steps: the widen to fp32 and + the ``* 256.0`` that puts the tile back on the true e4m3 scale exist only so the + result is directly usable, and they are what force the tile to 32 bits. A caller + that must apply a per-(token, kv_head) dequant scale anyway can fold ``2**8`` into + that scale instead -- exactly, since it is a power of two -- and keep the tile + 16-bit. Callers owe that fold; see :data:`KV_TILE_SCALE`. + + Same straight-line load-with-int-fill as that function, for the same reason (see + its docstring), and the same bit placement, folded: for ``v = 128s + r`` it builds + ``32768s + 128r`` with two masks, two widens, two shifts and an or, while + ``(v + (v & 0x80)) << 7 = (256s + r) << 7`` is the same number in one widen, one + mask, one add and one shift. Verified identical on all 256 codes, NaN patterns + (0x7F/0xFF -> +-480 after the caller's fold) included. + """ + w = tl.load(ptrs, mask=mask, other=0).to(tl.uint16) + return ((w + (w & 0x80)) << 7).to(tl.float16, bitcast=True) diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py index c5ca36553..1a9c57f48 100644 --- a/tests/kernels/test_e4m3_compat.py +++ b/tests/kernels/test_e4m3_compat.py @@ -142,6 +142,49 @@ def k(x_ptr, y_ptr, N, BLOCK: tl.constexpr): assert int(((y != ref) & ~(y.isnan() & ref.isnan())).sum()) == 0 +def test_kv_tile_scaled16_agrees_with_the_f32_loader(): + """kv_load_e4m3_tile_scaled16 is kv_load_e4m3_tile_f32 with its last two steps -- + the widen to fp32 and the ``* 256.0`` -- left for the caller to fold into the + per-(token, kv_head) dequant scale it has to apply anyway. That fold is only legal + if the two loaders agree on every code, so pin it: all 256, NaN patterns included. + + Second assertion pins the property that lets the 16-bit tile be narrowed to a bf16 + compute dtype for free: every value carries at most the code's own 3 mantissa bits + and lies in +-1.75, so bf16 holds it exactly. That is what makes scaling AFTER the + dot strictly more accurate than upstream's scale-then-narrow. + + Third pins masked lanes at zero in both, so a masked tile contributes nothing once + the scale is applied to the dot output instead of to the tile.""" + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import ( + KV_TILE_SCALE, + kv_load_e4m3_tile_f32, + kv_load_e4m3_tile_scaled16, + ) + + @triton.jit + def k(v_ptr, wide_ptr, narrow_ptr, N, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + mask = offs < N + tl.store(wide_ptr + offs, kv_load_e4m3_tile_f32(v_ptr + offs, mask)) + tl.store(narrow_ptr + offs, kv_load_e4m3_tile_scaled16(v_ptr + offs, mask)) + + n = 256 + v = torch.zeros(2 * n, dtype=torch.uint8, device="cuda") + v[:n] = torch.arange(n, dtype=torch.uint8, device="cuda") + wide = torch.empty(2 * n, dtype=torch.float32, device="cuda") + narrow = torch.empty(2 * n, dtype=torch.float16, device="cuda") + k[(1,)](v, wide, narrow, n, BLOCK=2 * n) + + assert torch.equal(narrow.to(torch.float32) * KV_TILE_SCALE.value, wide) + assert torch.equal( + narrow.to(torch.bfloat16).to(torch.float32), narrow.to(torch.float32) + ) + assert not wide[n:].any() and not narrow[n:].any() + + # ====================================================================================== # Shared emit path: every affected wrapper, deterministic inputs. # ====================================================================================== From 73ca76c8260d866621124e9ddd7ac743fd561985 Mon Sep 17 00:00:00 2001 From: Vincent Labreche Date: Sun, 6 Sep 2026 21:02:00 +0000 Subject: [PATCH 08/12] perf(kernel): size the extend tile from the KV cache element size _select_extend_tile budgets shared memory as (BLOCK_M + 2 * BLOCK_N) * BLOCK_D * 2 which charges K and V at 2 bytes/element whatever the cache actually holds. The q tile is always 2 bytes/element, but K and V follow the cache, so a 1-byte fp8 cache is billed for twice the shared memory it uses and falls through to a smaller tile than it has room for. On an RTX 3070 (sm_86, 99KB opt-in) at head_dim 256 that is BLOCK_N 32 where 64 fits. Take the element size as a parameter and bill K/V at it: (BLOCK_M * 2 + 2 * BLOCK_N * kv_bytes) * BLOCK_D kv_bytes=2 is algebraically the previous expression, so every 16-bit cache keeps the tile it had; the existing parametrisation in test_select_extend_tile_is_shared_memory_aware still passes unchanged. The head_dim <= 256 ladder gains a 64x64 rung between 128x64 and 64x32, which only an fp8 cache can reach on a consumer card. The budget stays a conservative proxy rather than an exact model. On this card it correctly rejects both tiles that fail to launch (128x64 and 64x128, which raise OutOfResources: shared memory, Required: 114688, Hardware limit: 101376) and correctly accepts the two the ladder uses. It also rejects 128x32 and 32x128, which do launch -- but those are not on the ladder, and rejecting a tile that would have worked only costs a smaller tile, never a failure. Same setup as the previous commit, measured on top of it: ctx TTFT before -> after decode before -> after 33k 33.95 -> 28.18 s 48.05 -> 47.94 t/s 65k 85.03 -> 64.36 s 41.61 -> 40.92 t/s 100k 160.42 -> 112.78 s 37.13 -> 36.26 t/s Decode is untouched by this commit (it only moves the extend/prefill tile); the small differences there are run-to-run noise at 2 reps. Co-Authored-By: Claude Opus 5 --- python/freetoken/kernel/triton/attention.py | 17 +++++++----- tests/kernels/test_triton_attention.py | 30 +++++++++++++++++++++ 2 files changed, 41 insertions(+), 6 deletions(-) diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index 6166e56fe..134e5b504 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -44,21 +44,26 @@ def _optin_smem_bytes(device_index: int) -> int: return int(getattr(props, "shared_memory_per_block_optin", 0)) -def _select_extend_tile(head_dim: int, block_d: int, smem_optin: int) -> tuple[int, int]: +def _select_extend_tile( + head_dim: int, block_d: int, smem_optin: int, kv_bytes: int = 2 +) -> tuple[int, int]: """Pick ``(BLOCK_M, BLOCK_N)`` for the extend/prefill kernel, shared-memory aware. - Larger tiles run materially faster (~2x for head_dim 512 on H100) but their bf16 - q/k/v tiles need about ``(BLOCK_M + 2 * BLOCK_N) * BLOCK_D * 2`` bytes of shared - memory, which overflows consumer GPUs (sm_89 ~99KB opt-in) once head_dim >= 256. Keep the fast tiles where the device's opt-in shared memory fits + Larger tiles run materially faster (~2x for head_dim 512 on H100) but their q/k/v + tiles need shared memory, which overflows consumer GPUs (sm_89 ~99KB opt-in) once + head_dim >= 256. Keep the fast tiles where the device's opt-in shared memory fits them (datacenter A100/H100); shrink only where it does not. ``smem_optin == 0`` (unknown) conservatively selects the small tiles, i.e. the prior consumer-safe behavior. + ``kv_bytes`` is the KV cache's element size: q is always 2 bytes/element but K and + V follow the cache, so a 1-byte fp8 cache fits a tile a 16-bit one cannot. Passing + 2 reproduces the previous budget exactly, so the 16-bit ladder is unchanged. """ budget = smem_optin * 0.8 # headroom for scores/acc/alignment/triton scratch def fits(block_m: int, block_n: int) -> bool: - return (block_m + 2 * block_n) * block_d * 2 <= budget + return (block_m * 2 + 2 * block_n * kv_bytes) * block_d <= budget if head_dim <= 128: return 128, 64 @@ -988,7 +993,7 @@ def extend_paged_attention( # shared memory fits them, shrink on consumer GPUs (sm_89 ~99KB) where the default # 128x64 overflows once head_dim >= 256 (e.g. gemma4: SWA 256, full-attention 512). block_m, block_n = _select_extend_tile( - head_dim, block_d, _optin_smem_bytes(q.device.index) + head_dim, block_d, _optin_smem_bytes(q.device.index), k_cache.element_size() ) grid = (qo_indptr.numel() - 1, num_q_heads, triton.cdiv(max_q_len, block_m)) if k_extend is not None or v_extend is not None: diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index ffff4a09b..7731d8c42 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -614,6 +614,36 @@ def test_select_extend_tile_is_shared_memory_aware(head_dim, smem_optin, expecte assert _select_extend_tile(head_dim, block_d, smem_optin) == expected +@pytest.mark.parametrize( + ("head_dim", "smem_optin", "expected_16bit", "expected_fp8"), + [ + # consumer opt-in smem (sm_86/sm_89 ~99KB): a 1-byte cache buys BLOCK_N 32 -> 64 + (256, 101376, (64, 32), (64, 64)), + # where the fast tile already fits, the cache dtype changes nothing + (256, 232448, (128, 64), (128, 64)), + # unknown budget stays conservative for both + (256, 0, (64, 32), (64, 32)), + ], +) +def test_select_extend_tile_uses_kv_cache_element_size( + head_dim, smem_optin, expected_16bit, expected_fp8 +): + """The q tile is always 2 bytes/element but K and V follow the cache, so charging + K/V at 2 bytes regardless makes an fp8 cache run a smaller tile than it has shared + memory for. On an RTX 3070 (sm_86, 99KB opt-in, head_dim 256) that was BLOCK_N 32 + where 64 fits, worth ~9% of prefill time at 99k context. + + The 16-bit column is the no-regression half: ``kv_bytes=2`` reproduces the previous + budget identically, since (M + 2N) * D * 2 == (M*2 + 2N*2) * D.""" + import triton + + from freetoken.kernel.triton.attention import _select_extend_tile + + block_d = triton.next_power_of_2(head_dim) + assert _select_extend_tile(head_dim, block_d, smem_optin, 2) == expected_16bit + assert _select_extend_tile(head_dim, block_d, smem_optin, 1) == expected_fp8 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") def test_triton_backend_stores_kv_and_matches_reference(monkeypatch): from freetoken.attention import AttentionSpec From 2f554c9a7c78e876a4c427d277065ecbf3abc0f9 Mon Sep 17 00:00:00 2001 From: ArqAlice Date: Mon, 7 Sep 2026 18:42:31 +0900 Subject: [PATCH 09/12] feat(kvcache): add nvfp4 kv quantization --- benchmarks/README.md | 9 + benchmarks/bench_kv_quant.py | 77 ++++++ docs/cli.md | 27 +- python/freetoken/attention/__init__.py | 3 + python/freetoken/attention/qsa_sparse.py | 7 +- python/freetoken/attention/triton.py | 21 +- python/freetoken/engine/engine.py | 17 +- python/freetoken/kernel/triton/attention.py | 137 +++++++++- python/freetoken/kernel/triton/kv_nvfp4.py | 124 +++++++++ python/freetoken/kernel/triton/qsa/attend.py | 71 +++++- python/freetoken/kvcache/__init__.py | 13 + python/freetoken/kvcache/base.py | 28 +- python/freetoken/kvcache/hybrid_swa_pool.py | 70 ++++- python/freetoken/kvcache/mha_pool.py | 50 +++- python/freetoken/kvcache/qsa_pool.py | 3 + python/freetoken/server/args.py | 4 +- tests/engine/test_kv_quant_config.py | 74 ++++++ tests/kernels/test_kv_nvfp4.py | 255 +++++++++++++++++++ tests/kernels/test_qsa_nvfp4.py | 141 ++++++++++ tests/kvcache/test_qsa_pool_fp8.py | 14 +- 20 files changed, 1097 insertions(+), 48 deletions(-) create mode 100644 benchmarks/bench_kv_quant.py create mode 100644 python/freetoken/kernel/triton/kv_nvfp4.py create mode 100644 tests/kernels/test_kv_nvfp4.py create mode 100644 tests/kernels/test_qsa_nvfp4.py diff --git a/benchmarks/README.md b/benchmarks/README.md index 6218903f2..5cbb0b34e 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -29,3 +29,12 @@ python benchmarks/bench_offload_cache_copy.py For host RAM vs PCIe bandwidth and the offload/hybrid backend pick, use `ft bench bw` instead — it writes the JSON profile the engine reads. + +**`bench_kv_quant.py`** compares BF16, FP8 and NVFP4 KV storage bytes, one-step +scatter latency and paged decode latency on synthetic inputs. No checkpoint is +required. Keep the GPU idle and use identical arguments for A/B comparisons; +this does not measure model quality or end-to-end serving throughput. + +```bash +PYTHONPATH=python:. uv run python benchmarks/bench_kv_quant.py --lengths 1024,8192,32768 +``` diff --git a/benchmarks/bench_kv_quant.py b/benchmarks/bench_kv_quant.py new file mode 100644 index 000000000..18707eb07 --- /dev/null +++ b/benchmarks/bench_kv_quant.py @@ -0,0 +1,77 @@ +"""Paged KV storage/scatter/decode microbenchmark, independent of model weights. + +Run with PYTHONPATH=python:. uv run python benchmarks/bench_kv_quant.py. +Compare identical arguments on the baseline and candidate; this does not measure +end-to-end model quality, TTFT, or serving throughput. +""" + +import argparse +import json + +import torch +import triton.testing + +from freetoken.distributed import set_tp_info +from freetoken.kernel.triton.attention import decode_paged_attention +from freetoken.kvcache.mha_pool import MHAKVCache + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--formats", default="none,fp8,nvfp4") + parser.add_argument("--lengths", default="1024,8192,32768") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--heads", type=int, default=4) + parser.add_argument("--group", type=int, default=4) + parser.add_argument("--dim", type=int, default=128) + args = parser.parse_args() + set_tp_info(rank=0, size=1) + torch.manual_seed(42) + batch, heads, dim = args.batch, args.heads, args.dim + qheads = heads * args.group + device = torch.device("cuda") + results = [] + for length in map(int, args.lengths.split(",")): + slots = batch * length + k, v = [torch.randn(slots, heads * dim, device=device, dtype=torch.bfloat16) + for _ in range(2)] + loc = torch.arange(slots, device=device, dtype=torch.int32) + q = torch.randn(batch, qheads, dim, device=device, dtype=torch.bfloat16) + indptr = torch.arange(batch + 1, device=device, dtype=torch.int32) * length + pos = torch.full((batch,), length - 1, device=device, dtype=torch.int32) + scratch = torch.empty(batch, qheads, 8, dim, device=device) + lse = torch.empty(batch, qheads, 8, device=device) + splits = torch.full((batch,), 8, device=device, dtype=torch.int32) + out = torch.empty_like(q) + for quant in args.formats.split(","): + pool = MHAKVCache(heads, 1, dim, slots, 1, q.dtype, device, kv_quant=quant) + pool.store_kv(k, v, loc, 0) + extra = {} + if quant == "nvfp4": + extra = dict(kv_quant=quant, k_block_scale=pool.k_block_scale(0), + v_block_scale=pool.v_block_scale(0)) + kc, vc = [getattr(pool, name)(0).flatten(0, 1) for name in ("k_cache", "v_cache")] + + def decode(): + return decode_paged_attention(q, kc, vc, indptr, loc, pos, + scratch, lse, splits, 8, dim ** -.5, out=out, + k_scale=pool.k_scale(0), v_scale=pool.v_scale(0), **extra) + + decode() + decode_ms = triton.testing.do_bench(decode, warmup=100, rep=300) + store_ms = triton.testing.do_bench( + lambda: pool.store_kv(k[-batch:], v[-batch:], loc[-batch:], 0), + warmup=100, rep=300) + record = dict(format=quant, length=length, batch=batch, + kv_bytes=pool.unit_bytes()[0] * slots, + decode_ms=decode_ms, store_ms=store_ms) + results.append(record) + print(json.dumps(record), flush=True) + del pool, kc, vc + del k, v + print(json.dumps(dict(gpu=torch.cuda.get_device_name(), torch=torch.__version__, + heads=heads, group=args.group, dim=dim, results=results))) + + +if __name__ == "__main__": + main() diff --git a/docs/cli.md b/docs/cli.md index 32a6b6f3e..552499a5b 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -71,7 +71,7 @@ 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) | +| `--kv-cache-dtype` | bf16 | `bf16`, `fp8`, or `nvfp4` (see [NVFP4 KV cache](#nvfp4-kv-cache)): FP8 stores 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 @@ -197,3 +197,28 @@ profile that `ft serve --moe-backend auto` and `--moe-hybrid-max-fetch -1` then - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe by that factor. + +### NVFP4 KV cache + +`ft serve --model --kv-cache-dtype nvfp4 --attention-backend triton` +opts into packed E2M1 KV storage. The initial implementation supports plain paged +FULL attention (MHA/GQA), hybrid-SWA, and the full-attention portion of hybrid-linear +models, and QSA. Head dimensions must be divisible by 16. MLA/DSA, DSV4, and BSA +pools are rejected at startup. `auto` selects Triton or QSA sparse attention for +supported models. + +Each K or V row stores `head_dim / 2` packed bytes, `head_dim / 16` E4M3 block-scale +bytes, and one FP32 row scale. At head_dim 128 this is 76 bytes, versus 256 for +BF16 and 132 for the existing FP8 format. Pool management, recurrent states, +attention workspace and model weights consume additional memory. + +The second-level scale is dynamic per token/head, so appending a token never +rescales an existing prefix. This is a FreeToken KV layout, not an external +NVFP4 checkpoint or attention-library ABI. K/V are restored inside attention; +Q and attention arithmetic retain their compute precision. The MoE weight option +`--nvfp4-backend` is independent. Prefill uses fresh compute-dtype K/V while +cached prefixes are restored, as in the FP8 path. + +NVFP4 is opt-in: assess quality on your checkpoint and workload before using it +for long-context inference. Capacity savings do not guarantee faster decode; +packing, reconstruction, and the selected attention backend affect throughput. diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 012759ad3..f344f2f53 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -38,6 +38,7 @@ class BackendInfo: # 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 + supports_nvfp4_kv: bool = False SUPPORTED_ATTENTION_BACKENDS = Registry[BackendCreator]("Attention Backend") @@ -90,6 +91,7 @@ def create_fa_backend(config: ModelConfig): supported_types=frozenset({AttnType.FULL, AttnType.SWA}), consumes_attn_spec=True, supports_fp8_kv=True, + supports_nvfp4_kv=True, ), ) def create_triton_backend(config: ModelConfig): @@ -146,6 +148,7 @@ def create_m3_sparse_backend(config: ModelConfig): # 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, + supports_nvfp4_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 68153b1d4..8423afbcb 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -292,8 +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. + # K/V scale tensors are independent of the BF16 index tier, so selection is + # quantization-agnostic; only sparse K/V attention reconstructs the codes. return qsa_sparse_paged_attention( q, self.kvcache.k_cache(layer_id), @@ -304,6 +304,9 @@ def qsa_forward( torch.empty_like(q), k_scale=self.kvcache.k_scale(layer_id), v_scale=self.kvcache.v_scale(layer_id), + kv_quant=self.kvcache.kv_quant, + k_block_scale=self.kvcache.k_block_scale(layer_id), + v_block_scale=self.kvcache.v_block_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 c731c2636..58aa6d7c8 100644 --- a/python/freetoken/attention/triton.py +++ b/python/freetoken/attention/triton.py @@ -151,10 +151,14 @@ def forward( k_raw = self.kvcache.k_cache(layer_id) v_raw = self.kvcache.v_cache(layer_id) - kv_heads, head_dim = k_raw.shape[-2], k_raw.shape[-1] - 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) + kv_heads, stored_dim = k_raw.shape[-2], k_raw.shape[-1] + head_dim = q.shape[-1] + kv_quant = getattr(self.kvcache, "kv_quant", "none") + assert stored_dim == (head_dim // 2 if kv_quant == "nvfp4" else head_dim) + k_cache = k_raw.view(-1, kv_heads, stored_dim) + v_cache = v_raw.view(-1, kv_heads, stored_dim) + k_block_scale = self.kvcache.k_block_scale(layer_id) if kv_quant == "nvfp4" else None + v_block_scale = self.kvcache.v_block_scale(layer_id) if kv_quant == "nvfp4" else None # 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) @@ -188,6 +192,9 @@ def forward( sinks=spec.sinks, k_scale=k_scale, v_scale=v_scale, + kv_quant=kv_quant, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, ) if ( (not metadata.is_decode) @@ -210,6 +217,9 @@ def forward( v_extend=v.view(q.shape[0], kv_heads, head_dim), k_scale=k_scale, v_scale=v_scale, + kv_quant=kv_quant, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, ) return paged_attention( q=q, @@ -224,6 +234,9 @@ def forward( sinks=spec.sinks, k_scale=k_scale, v_scale=v_scale, + kv_quant=kv_quant, + k_block_scale=k_block_scale, + v_block_scale=v_block_scale, ) def prepare_metadata(self, batch: Batch) -> None: diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cc87d95bf..045f60c65 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -115,7 +115,7 @@ def _backend_requirements_met(name: str) -> bool: # --kv-cache-dtype spellings -> the stored EngineConfig.kv_quant value. -KV_QUANT_ALIASES = {"auto": "none", "bf16": "none", "none": "none", "fp8": "fp8"} +KV_QUANT_ALIASES = {"auto": "none", "bf16": "none", "none": "none", "fp8": "fp8", "nvfp4": "nvfp4"} def _resolve_kv_quant(value: str | None) -> str: @@ -134,8 +134,11 @@ def _backend_supports_kv_quant(name: str, kv_quant: str) -> bool: KV pool (an unquantized pool needs nothing from the backend).""" if kv_quant == "none": return True + if kv_quant not in ("fp8", "nvfp4"): + return False return all( - attention_backend_info(part.strip()).supports_fp8_kv for part in name.split(",") + getattr(attention_backend_info(part.strip()), f"supports_{kv_quant}_kv") + for part in name.split(",") ) @@ -237,7 +240,7 @@ def _validate_attention_backend_choice(config, override, required: frozenset[Att name for name in ("trtllm", "fi", "fa", "triton") if required <= attention_backend_info(name).supported_types - and attention_backend_info(name).supports_fp8_kv + and _backend_supports_kv_quant(name, kv_quant) ] raise ValueError( f"--kv-cache-dtype {kv_quant} needs an attention backend that decodes the KV " @@ -1358,6 +1361,14 @@ def override(attr: str, value: Any): # this is dangerous, use with caution # 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 == "nvfp4": + if required_attn_types - {AttnType.FULL, AttnType.SWA, AttnType.QSA}: + raise ValueError( + "--kv-cache-dtype nvfp4 requires a plain paged FULL, hybrid-SWA, or QSA KV pool" + ) + for spec in model_config.kv_cache_group_specs(): + if spec.head_dim % 16: + raise ValueError("--kv-cache-dtype nvfp4 requires head_dim divisible by 16") 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 diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index 31ae84847..f6a9a5419 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -7,6 +7,7 @@ import triton.language as tl from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 as _kv_load_f32 +from freetoken.kernel.triton.kv_nvfp4 import load_nvfp4 _MAX_KV_SPLITS = 8 @@ -51,6 +52,8 @@ def _paged_attention_kernel( v_ptr, k_scale_ptr, v_scale_ptr, + k_block_ptr, + v_block_ptr, o_ptr, indptr_ptr, indices_ptr, @@ -75,6 +78,7 @@ def _paged_attention_kernel( SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ): q_tok = tl.program_id(0) q_head = tl.program_id(1) @@ -114,7 +118,13 @@ 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) - if HAS_KV_SCALE: + if KV_NVFP4: + k = load_nvfp4( + k_ptr, k_block_ptr, k_scale_ptr, slots[:, None], kv_head, + offs_d[None, :], + offs_n[:, None] < kv_len, stride_ks, stride_kh, stride_kss, D, + ) + elif 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( @@ -147,7 +157,13 @@ def _paged_attention_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new) - if HAS_KV_SCALE: + if KV_NVFP4: + v = load_nvfp4( + v_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, + offs_d[None, :], + offs_n[:, None] < kv_len, stride_vs, stride_vh, stride_vss, D, + ) + elif HAS_KV_SCALE: s_v = tl.load( v_scale_ptr + slots * stride_vss + kv_head, mask=offs_n < kv_len, @@ -188,6 +204,8 @@ def _decode_grouped_stage1_kernel( v_ptr, k_scale_ptr, v_scale_ptr, + k_block_ptr, + v_block_ptr, sm_scale, indptr_ptr, indices_ptr, @@ -221,6 +239,7 @@ def _decode_grouped_stage1_kernel( DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ): batch_id = tl.program_id(0) head_block_id = tl.program_id(1) @@ -277,7 +296,13 @@ 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) - if HAS_KV_SCALE: + if KV_NVFP4: + k = load_nvfp4( + k_ptr, k_block_ptr, k_scale_ptr, slots[None, :], kv_head, + offs_d[:, None], + mask_n[None, :], stride_ks, stride_kh, stride_kss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: s_k = tl.load( k_scale_ptr + slots * stride_kss + kv_head, mask=mask_n, other=0.0 ) @@ -295,7 +320,13 @@ def _decode_grouped_stage1_kernel( scores = tl.dot(q, k) * sm_scale scores = tl.where(mask_h[:, None] & mask_n[None, :], scores, -float("inf")) - if HAS_KV_SCALE: + if KV_NVFP4: + v = load_nvfp4( + v_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, + offs_dv[None, :], + mask_n[:, None], stride_vs, stride_vh, stride_vss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: s_v = tl.load( v_scale_ptr + slots * stride_vss + kv_head, mask=mask_n, other=0.0 ) @@ -414,6 +445,27 @@ def _decode_stage2_kernel( ) +def _validate_kv_format(q, k, v, kr, vr, quant, kb, vb): + quant = quant if quant is not None else ("fp8" if kr is not None else "none") + if quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unknown kv_quant {quant!r}") + d = q.shape[-1] + assert k.shape == v.shape + assert k.stride(-1) == v.stride(-1) == 1 + assert (kr is not None) == (vr is not None) == (quant != "none") + if quant == "nvfp4": + assert d % 16 == 0 and k.shape[-1] == d // 2 + for codes, row, block in ((k, kr, kb), (v, vr, vb)): + assert codes.dtype == torch.uint8 + assert row.dtype == torch.float32 and row.is_contiguous() + assert block is not None and block.dtype == torch.uint8 + assert block.shape == (*codes.shape[:2], d // 16) and block.is_contiguous() + assert codes.device == row.device == block.device == q.device + else: + assert k.shape[-1] == d and kb is None and vb is None + return quant == "nvfp4" + + def decode_paged_attention( q: torch.Tensor, k_cache: torch.Tensor, @@ -431,9 +483,15 @@ def decode_paged_attention( out: torch.Tensor | None = None, k_scale: torch.Tensor | None = None, v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: """SGLang-style split-k grouped decode attention for one query per request. + ``kv_quant="nvfp4"`` requires packed half-width codes, FP32 row scales and + uint8 E4M3 block scales. Omitted ``kv_quant`` retains FP8 scale inference. + ``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. @@ -446,7 +504,8 @@ def decode_paged_attention( num_kv_heads = k_cache.shape[1] assert batch == indptr.numel() - 1 assert v_cache.shape[1] == num_kv_heads - assert k_cache.shape[-1] == head_dim and v_cache.shape[-1] == head_dim + nvfp4 = _validate_kv_format(q, k_cache, v_cache, k_scale, v_scale, + kv_quant, k_block_scale, v_block_scale) assert num_q_heads % num_kv_heads == 0 assert attn_logits.shape[0] >= batch assert attn_logits.shape[1] >= num_q_heads @@ -468,6 +527,8 @@ def decode_paged_attention( 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 + k_block_arg = k_block_scale if nvfp4 else k_cache + v_block_arg = v_block_scale if nvfp4 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), @@ -489,6 +550,8 @@ def decode_paged_attention( v_cache, k_scale_arg, v_scale_arg, + k_block_arg, + v_block_arg, sm_scale, indptr, indices, @@ -522,6 +585,7 @@ def decode_paged_attention( DV=head_dim, SLIDING_WINDOW=sliding_window or 0, HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, num_warps=4, num_stages=2, ) @@ -560,6 +624,8 @@ def _extend_attention_kernel( v_ptr, k_scale_ptr, v_scale_ptr, + k_block_ptr, + v_block_ptr, o_ptr, qo_indptr_ptr, kv_indptr_ptr, @@ -586,6 +652,7 @@ def _extend_attention_kernel( SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -637,7 +704,13 @@ 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) - if HAS_KV_SCALE: + if KV_NVFP4: + k = load_nvfp4( + k_ptr, k_block_ptr, k_scale_ptr, slots[None, :], kv_head, + offs_d[:, None], + mask_n[None, :], stride_ks, stride_kh, stride_kss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: s_k = tl.load( k_scale_ptr + slots * stride_kss + kv_head, mask=mask_n, other=0.0 ) @@ -667,7 +740,13 @@ def _extend_attention_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - if HAS_KV_SCALE: + if KV_NVFP4: + v = load_nvfp4( + v_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, + offs_dv[None, :], + mask_n[:, None], stride_vs, stride_vh, stride_vss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: s_v = tl.load( v_scale_ptr + slots * stride_vss + kv_head, mask=mask_n, other=0.0 ) @@ -712,6 +791,8 @@ def _extend_attention_split_kernel( v_cache_ptr, k_scale_ptr, v_scale_ptr, + k_block_ptr, + v_block_ptr, o_ptr, qo_indptr_ptr, kv_indptr_ptr, @@ -742,6 +823,7 @@ def _extend_attention_split_kernel( SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ): seq_id = tl.program_id(0) q_head = tl.program_id(1) @@ -795,7 +877,13 @@ def _extend_attention_split_kernel( if not skip_tile: slots = tl.load(kv_indices_ptr + kv_start + kv_offsets, mask=mask_n, other=0) - if HAS_KV_SCALE: + if KV_NVFP4: + k = load_nvfp4( + k_cache_ptr, k_block_ptr, k_scale_ptr, slots[None, :], kv_head, + offs_d[:, None], + mask_n[None, :], stride_kcs, stride_kch, stride_kss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: s_k = tl.load( k_scale_ptr + slots * stride_kss + kv_head, mask=mask_n, other=0.0 ) @@ -825,7 +913,13 @@ def _extend_attention_split_kernel( alpha = tl.exp(m_i - m_new) p = tl.exp(scores - m_new[:, None]) - if HAS_KV_SCALE: + if KV_NVFP4: + v = load_nvfp4( + v_cache_ptr, v_block_ptr, v_scale_ptr, slots[:, None], kv_head, + offs_dv[None, :], + mask_n[:, None], stride_vcs, stride_vch, stride_vss, D, + ).to(q.dtype) + elif HAS_KV_SCALE: s_v = tl.load( v_scale_ptr + slots * stride_vss + kv_head, mask=mask_n, other=0.0 ) @@ -924,6 +1018,9 @@ def extend_paged_attention( v_extend: torch.Tensor | None = None, k_scale: torch.Tensor | None = None, v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Block-tiled causal prefill/extend attention over paged KV cache. @@ -940,7 +1037,8 @@ def extend_paged_attention( assert qo_indptr.numel() == kv_indptr.numel() assert prefix_lens.numel() == qo_indptr.numel() - 1 assert v_cache.shape[1] == num_kv_heads - assert k_cache.shape[-1] == head_dim and v_cache.shape[-1] == head_dim + nvfp4 = _validate_kv_format(q, k_cache, v_cache, k_scale, v_scale, + kv_quant, k_block_scale, v_block_scale) assert num_q_heads % num_kv_heads == 0 if sinks is not None: assert sinks.is_cuda @@ -954,6 +1052,8 @@ def extend_paged_attention( # 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 + k_block_arg = k_block_scale if nvfp4 else k_cache + v_block_arg = v_block_scale if nvfp4 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), @@ -983,6 +1083,8 @@ def extend_paged_attention( v_cache, k_scale_arg, v_scale_arg, + k_block_arg, + v_block_arg, o, qo_indptr, kv_indptr, @@ -1013,6 +1115,7 @@ def extend_paged_attention( SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, num_warps=8, num_stages=1, ) @@ -1024,6 +1127,8 @@ def extend_paged_attention( v_cache, k_scale_arg, v_scale_arg, + k_block_arg, + v_block_arg, o, qo_indptr, kv_indptr, @@ -1050,6 +1155,7 @@ def extend_paged_attention( SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, num_warps=8, num_stages=1, ) @@ -1071,6 +1177,9 @@ def paged_attention( block_n: int = 32, k_scale: torch.Tensor | None = None, v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Paged causal attention for one layer. @@ -1086,7 +1195,8 @@ def paged_attention( num_tokens, num_q_heads, head_dim = q.shape num_kv_heads = k_cache.shape[1] assert v_cache.shape[1] == num_kv_heads - assert k_cache.shape[-1] == head_dim and v_cache.shape[-1] == head_dim + nvfp4 = _validate_kv_format(q, k_cache, v_cache, k_scale, v_scale, + kv_quant, k_block_scale, v_block_scale) assert num_q_heads % num_kv_heads == 0 if sinks is not None: assert sinks.is_cuda @@ -1099,6 +1209,8 @@ def paged_attention( 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 + k_block_arg = k_block_scale if nvfp4 else k_cache + v_block_arg = v_block_scale if nvfp4 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), @@ -1112,6 +1224,8 @@ def paged_attention( v_cache, k_scale_arg, v_scale_arg, + k_block_arg, + v_block_arg, o, indptr, indices, @@ -1136,6 +1250,7 @@ def paged_attention( SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, HAS_KV_SCALE=has_kv_scale, + KV_NVFP4=nvfp4, num_warps=8 if head_dim >= 256 else 4, num_stages=2, ) diff --git a/python/freetoken/kernel/triton/kv_nvfp4.py b/python/freetoken/kernel/triton/kv_nvfp4.py new file mode 100644 index 000000000..91bfc602b --- /dev/null +++ b/python/freetoken/kernel/triton/kv_nvfp4.py @@ -0,0 +1,124 @@ +"""Row-scaled NVFP4 KV storage, with low nibble first and 16-wide blocks. + +Each (token, head) row has an FP32 scale; each block has an E4M3 scale +stored as uint8. Reconstruction is E2M1 * block_scale * row_scale. +Row-local second-level scales keep appends independent of the cached prefix. +""" + +import torch +import triton +import triton.language as tl + +from .e4m3_compat import e4m3_f32_to_u8, e4m3_u8_to_f32, round_e4m3 + + +@triton.jit +def _encode_e2m1(x): + a = tl.abs(x) + code = tl.full(x.shape, 0, tl.int32) + code = tl.where(a > 0.25, 1, code) + code = tl.where(a >= 0.75, 2, code) + code = tl.where(a > 1.25, 3, code) + code = tl.where(a >= 1.75, 4, code) + code = tl.where(a > 2.5, 5, code) + code = tl.where(a >= 3.5, 6, code) + code = tl.where(a > 5.0, 7, code) + return code | tl.where(x < 0, 8, 0) + + +@triton.jit +def _decode_e2m1(code): + # Place E2M1 bits in FP16, then compensate the exponent-bias difference (15 - 1). + bits = ((code & 8).to(tl.uint16) << 12) | ((code & 7).to(tl.uint16) << 9) + return bits.to(tl.float16, bitcast=True).to(tl.float32) * 16384.0 + + +@triton.jit +def load_nvfp4(ptr, block_ptr, row_ptr, slots, head, dims, slot_mask, + stride_slot, stride_head, stride_row, D: tl.constexpr): + TRANSPOSE: tl.constexpr = dims.shape[0] != 1 + WIDTH: tl.constexpr = dims.shape[0] if TRANSPOSE else dims.shape[1] + TOKENS: tl.constexpr = slots.shape[0] * slots.shape[1] + slot = slots.reshape(TOKENS).to(tl.int64) + valid = slot_mask.reshape(TOKENS) + dim = tl.arange(0, WIDTH) + packed = tl.load( + ptr + slot[:, None] * stride_slot + head * stride_head + (dim[None, :] // 2), + valid[:, None] & (dim[None, :] < D), other=0, + ).to(tl.int32) + codes = tl.where((dim[None, :] & 1) == 0, packed & 15, packed >> 4) + block_dim = dim // 16 + block = tl.load( + block_ptr + (slot[:, None] * stride_row + head) * (D // 16) + block_dim[None, :], + valid[:, None] & (dim[None, :] < D), other=0, + ) + row = tl.load(row_ptr + slot * stride_row + head, valid, other=0) + value = _decode_e2m1(codes) * e4m3_u8_to_f32(block) * row[:, None] + if TRANSPOSE: + return tl.trans(value) + else: + return value + + +@triton.jit +def _quantize_row(src, dst, block_ptr, row_ptr, t, h, slot, stride_src, + HEADS: tl.constexpr, D: tl.constexpr, BLOCKS: tl.constexpr): + blocks = tl.arange(0, BLOCKS) + dims = blocks[:, None] * 16 + tl.arange(0, 16)[None, :] + x = tl.load(src + t * stride_src + h * D + dims, dims < D, other=0).to(tl.float32) + amax = tl.max(tl.abs(x), 1) + row_scale = tl.maximum(tl.max(amax, 0), 1e-10) / (6.0 * 448.0) + block_scale = round_e4m3(tl.minimum(tl.div_rn(amax, 6.0 * row_scale), 448.0)) + # Quantize against the scale actually stored, including E4M3 rounding/underflow. + denom = tl.where(block_scale > 0, block_scale * row_scale, 1.0) + normalized = tl.where(block_scale[:, None] > 0, tl.div_rn(x, denom[:, None]), 0.0) + codes = _encode_e2m1(normalized).reshape(BLOCKS, 8, 2) + lo, hi = tl.split(codes) + packed = lo | (hi << 4) + byte_dims = blocks[:, None] * 8 + tl.arange(0, 8)[None, :] + tl.store(dst + (slot * HEADS + h) * (D // 2) + byte_dims, packed, byte_dims < D // 2) + tl.store(block_ptr + (slot * HEADS + h) * (D // 16) + blocks, + e4m3_f32_to_u8(block_scale), blocks < D // 16) + tl.store(row_ptr + slot * HEADS + h, row_scale) + + +@triton.jit +def _scatter(k, v, kc, vc, kb, vb, kr, vr, indices, stride_k, stride_v, + HEADS: tl.constexpr, D: tl.constexpr, BLOCKS: tl.constexpr): + t, h = tl.program_id(0), tl.program_id(1) + slot = tl.load(indices + t).to(tl.int64) + _quantize_row(k, kc, kb, kr, t, h, slot, stride_k, HEADS, D, BLOCKS) + _quantize_row(v, vc, vb, vr, t, h, slot, stride_v, HEADS, D, BLOCKS) + + +def quantize_nvfp4_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, + k_block_scale: torch.Tensor, + v_block_scale: torch.Tensor, +) -> None: + tokens, width = k.shape + slots, heads, packed_dim = k_cache.shape + dim = packed_dim * 2 + assert dim % 16 == 0 and width == heads * dim + assert v.shape == k.shape and k.stride(1) == v.stride(1) == 1 + assert k.dtype in (torch.float16, torch.bfloat16, torch.float32) and v.dtype == k.dtype + assert out_loc.shape == (tokens,) and out_loc.dtype in (torch.int32, torch.int64) + assert out_loc.is_contiguous() + for codes, row, block in ((k_cache, k_scale, k_block_scale), (v_cache, v_scale, v_block_scale)): + assert codes.shape == (slots, heads, packed_dim) and codes.dtype == torch.uint8 + assert row.shape == (slots, heads) and row.dtype == torch.float32 + assert block.shape == (slots, heads, dim // 16) and block.dtype == torch.uint8 + assert codes.is_contiguous() and row.is_contiguous() and block.is_contiguous() + assert codes.device == row.device == block.device == k.device + assert k.is_cuda and v.device == out_loc.device == k.device + if tokens: + _scatter[(tokens, heads)](k, v, k_cache, v_cache, k_block_scale, v_block_scale, + k_scale, v_scale, out_loc, k.stride(0), v.stride(0), + HEADS=heads, D=dim, BLOCKS=triton.next_power_of_2(dim // 16), + num_warps=4, enable_fp_fusion=False) diff --git a/python/freetoken/kernel/triton/qsa/attend.py b/python/freetoken/kernel/triton/qsa/attend.py index 7e5815487..f6e223385 100644 --- a/python/freetoken/kernel/triton/qsa/attend.py +++ b/python/freetoken/kernel/triton/qsa/attend.py @@ -10,6 +10,7 @@ import triton.language as tl from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 +from freetoken.kernel.triton.kv_nvfp4 import load_nvfp4 @triton.jit @@ -19,6 +20,8 @@ def _qsa_sparse_paged_gqa_splitk_kernel( v_cache_ptr, k_scale_ptr, v_scale_ptr, + k_block_scale_ptr, + v_block_scale_ptr, indices_ptr, block_table_ptr, token_to_req_ptr, @@ -56,6 +59,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( # 16-bit values. The bf16 branch below stays exactly as it was, instruction for # instruction, for the unquantized default. HAS_KV_SCALE: tl.constexpr, + KV_NVFP4: tl.constexpr, ) -> None: # row * stride can overflow int32 for large row counts. row = tl.program_id(0).to(tl.int64) @@ -111,7 +115,35 @@ 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) - if HAS_KV_SCALE: + if KV_NVFP4: + scale_slot = safe_page * PAGE_SIZE + page_offset + keys = load_nvfp4( + k_cache_ptr, + k_block_scale_ptr, + k_scale_ptr, + scale_slot[None, :], + kv_head, + dim_offsets[:, None], + valid[None, :], + stride_k_token, + stride_k_head, + stride_kss, + HEAD_DIM, + ).to(query.dtype) + values = load_nvfp4( + v_cache_ptr, + v_block_scale_ptr, + v_scale_ptr, + scale_slot[:, None], + kv_head, + dim_offsets[None, :], + valid[:, None], + stride_v_token, + stride_v_head, + stride_vss, + HEAD_DIM, + ).to(query.dtype) + elif 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. @@ -286,8 +318,11 @@ def qsa_sparse_paged_attention( out: torch.Tensor | None = None, k_scale: torch.Tensor | None = None, v_scale: torch.Tensor | None = None, + kv_quant: str | None = None, + k_block_scale: torch.Tensor | None = None, + v_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: - """Run sparse GQA directly over paged K/V caches (bf16, or e4m3 + row scales).""" + """Run sparse GQA over bf16, FP8, or packed NVFP4 K/V caches.""" 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") @@ -297,12 +332,25 @@ def qsa_sparse_paged_attention( raise ValueError("QSA sparse attention metadata has invalid shapes") if logical_indices.shape[1] <= 0: raise ValueError("QSA sparse attention requires a positive selection width") - if q.shape[2] != k_cache.shape[3] or q.shape[1] % k_cache.shape[2]: + kv_quant = kv_quant if kv_quant is not None else ("fp8" if k_scale is not None else "none") + if kv_quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unknown QSA kv_quant {kv_quant!r}") + stored_dim = q.shape[2] // 2 if kv_quant == "nvfp4" else q.shape[2] + if ( + (kv_quant == "nvfp4" and q.shape[2] % 16) + or k_cache.shape[3] != stored_dim + or q.shape[1] % k_cache.shape[2] + ): 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 - if (k_scale is None) != (v_scale is None): + if (k_scale is None) != (v_scale is None) or (k_block_scale is None) != (v_block_scale is None): raise ValueError("QSA sparse attention requires both KV scale tensors") + nvfp4 = kv_quant == "nvfp4" + if nvfp4 and (k_scale is None or k_block_scale is None): + raise ValueError("QSA NVFP4 requires row and block scale tensors") + if not nvfp4 and (k_block_scale is not None or v_block_scale is not None): + raise ValueError("QSA block scales require kv_quant='nvfp4'") 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 @@ -316,6 +364,18 @@ def qsa_sparse_paged_attention( 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 + if nvfp4: + block_want = (*want, head_dim // 16) + if ( + k_cache.dtype is not torch.uint8 + or v_cache.dtype is not torch.uint8 + or k_block_scale.dtype is not torch.uint8 + or v_block_scale.dtype is not torch.uint8 + or k_block_scale.shape != block_want + or v_block_scale.shape != block_want + ): + raise ValueError(f"QSA NVFP4 block scales must have shape {block_want}") + assert k_block_scale.stride(2) == v_block_scale.stride(2) == 1 else: assert q.dtype == k_cache.dtype == v_cache.dtype assert logical_indices.dtype == block_table.dtype == torch.int32 @@ -377,6 +437,8 @@ def qsa_sparse_paged_attention( # 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, + k_cache if k_block_scale is None else k_block_scale, + v_cache if v_block_scale is None else v_block_scale, logical_indices, block_table, token_to_req, @@ -411,6 +473,7 @@ def qsa_sparse_paged_attention( BLOCK_M=block_m, BLOCK_N=block_n, HAS_KV_SCALE=k_scale is not None, + KV_NVFP4=nvfp4, num_warps=partial_warps, num_stages=2, ) diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index a7510d69f..4cd576d8c 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -144,6 +144,19 @@ def create_kvcache_pool( num_req_slots: int | None = None, kv_quant: str = "none", ) -> BaseKVCachePool: + if kv_quant == "nvfp4": + from freetoken.attention import AttnType + + if any( + spec.attn_type not in (AttnType.FULL, AttnType.SWA, AttnType.QSA) + or spec.mla + or spec.head_dim % 16 + for spec in model_config.kv_cache_group_specs() + ): + raise ValueError( + "--kv-cache-dtype nvfp4 requires paged FULL, hybrid-SWA, or QSA groups " + "with head_dim divisible by 16" + ) if model_config.has_swa_attention: from .hybrid_swa_pool import HybridSWAKVCache diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index abcb916c1..276f53524 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -35,15 +35,18 @@ def kv_storage_bytes_per_elem(config) -> int: 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. + """Sidecar bytes per token: FP32 row scales plus NVFP4 E4M3 block scales. + The unquantized pool has no scales. 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": + if getattr(config, "kv_quant", "none") not in ("fp8", "nvfp4"): 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 + scale_bytes = FP8_KV_SCALE_BYTES + if getattr(config, "kv_quant", "none") == "nvfp4": + scale_bytes += spec.head_dim // 16 + return (1 if spec.mla else 2) * spec.num_layers * heads * scale_bytes def spec_kv_bytes_per_token(spec, config) -> int: @@ -56,11 +59,16 @@ def spec_kv_bytes_per_token(spec, config) -> int: ``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.""" + if getattr(config, "kv_quant", "none") == "nvfp4": + if spec.head_dim % 16: + raise ValueError("NVFP4 KV requires head_dim divisible by 16") + row_bytes = spec.head_dim // 2 + else: + row_bytes = spec.head_dim * kv_storage_bytes_per_elem(config) per_token = ( (1 if spec.mla else 2) # MLA latent groups store one slab (V aliases K) - * spec.head_dim + * row_bytes * div_even(spec.num_kv_heads, config.tp_info.size, allow_replicate=True) - * kv_storage_bytes_per_elem(config) * spec.num_layers ) return ( @@ -204,6 +212,14 @@ 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 + def k_block_scale(self, index: int) -> torch.Tensor | None: + """NVFP4 E4M3 scales for ``k_cache(index)``, or ``None`` for other layouts.""" + return None + + def v_block_scale(self, index: int) -> torch.Tensor | None: + """NVFP4 E4M3 scales for ``v_cache(index)``, or ``None`` for other layouts.""" + return None + @property @abstractmethod def device(self) -> torch.device: ... diff --git a/python/freetoken/kvcache/hybrid_swa_pool.py b/python/freetoken/kvcache/hybrid_swa_pool.py index 1813128b2..bd69b2884 100644 --- a/python/freetoken/kvcache/hybrid_swa_pool.py +++ b/python/freetoken/kvcache/hybrid_swa_pool.py @@ -23,8 +23,11 @@ class _KVGroupStorage: k_buffer: torch.Tensor v_buffer: torch.Tensor storage_shape: tuple[int, int, int] + logical_head_dim: int # (2, num_layers, num_slots, local_kv_heads) fp32, or None for an unquantized group. scale_buffer: torch.Tensor | None = None + # (2, num_layers, num_slots, local_kv_heads, head_dim // 16) uint8, or None. + block_scale_buffer: torch.Tensor | None = None def _alloc_group_storage( @@ -36,7 +39,7 @@ def _alloc_group_storage( store_dtype: torch.dtype, outer_size: int, inner_size: int, - quantized: bool, + kv_quant: str, ) -> _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. @@ -46,8 +49,9 @@ def _alloc_group_storage( 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: + stored_dim = head_dim // 2 if kv_quant == "nvfp4" else head_dim + shape = (2, num_layers, outer_size, inner_size, local_kv_heads, stored_dim) + if kv_quant != "none": from freetoken.kernel.triton.kv_quant import alloc_codes buffer = alloc_codes(shape, device) @@ -59,12 +63,21 @@ def _alloc_group_storage( else: buffer = torch.empty(shape, device=device, dtype=store_dtype) scale = None + block_scale = None + if kv_quant == "nvfp4": + block_scale = torch.zeros( + (2, num_layers, outer_size * inner_size, local_kv_heads, head_dim // 16), + device=device, + dtype=torch.uint8, + ) return _KVGroupStorage( buffer=buffer, k_buffer=buffer[0], v_buffer=buffer[1], - storage_shape=(outer_size * inner_size, local_kv_heads, head_dim), + storage_shape=(outer_size * inner_size, local_kv_heads, stored_dim), + logical_head_dim=head_dim, scale_buffer=scale, + block_scale_buffer=block_scale, ) @@ -82,9 +95,13 @@ def __init__( num_swa_tokens: int | None = None, kv_quant: str = "none", ) -> None: + if kv_quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unsupported hybrid-SWA kv_quant {kv_quant!r}") 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)}") + if kv_quant == "nvfp4" and any(spec.head_dim % 16 for spec in specs.values()): + raise ValueError("NVFP4 KV requires head_dim divisible by 16") from .mha_pool import _kv_store_dtype @@ -154,7 +171,7 @@ def _allocate_group( store_dtype=_kv_store_dtype(dtype, kv_quant), outer_size=outer_size, inner_size=inner_size, - quantized=kv_quant != "none", + kv_quant=kv_quant, ) @staticmethod @@ -265,6 +282,16 @@ def v_scale(self, index: int) -> torch.Tensor | None: scale = self._storages[ref.group].scale_buffer return None if scale is None else scale[1][ref.index] + def k_block_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scale = self._storages[ref.group].block_scale_buffer + return None if scale is None else scale[0][ref.index] + + def v_block_scale(self, index: int) -> torch.Tensor | None: + ref = self.layers_mapping[index] + scale = self._storages[ref.group].block_scale_buffer + return None if scale is None else scale[1][ref.index] + def store_kv( self, k: torch.Tensor, @@ -277,6 +304,21 @@ def store_kv( indices = out_loc if ref.group == "swa": indices = self.translate_loc_from_full_to_swa(out_loc) + if self.kv_quant == "nvfp4": + from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_to_cache + + quantize_nvfp4_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], + k_block_scale=storage.block_scale_buffer[0][ref.index], + v_block_scale=storage.block_scale_buffer[1][ref.index], + ) + return if self.kv_quant == "fp8": from freetoken.kernel.triton.kv_quant import quantize_kv_to_cache @@ -328,22 +370,22 @@ 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, storage dtype, and whether codes are fp8. + # kv heads, logical head_dim, device, storage dtype, and quantized sidecars. # (Plain ints + device/dtype handles, no tensor.) - _, num_layers, _old_outer, _old_inner, local_kv_heads, head_dim = group.buffer.shape + _, num_layers, _old_outer, _old_inner, local_kv_heads, _stored_dim = group.buffer.shape return ( num_layers, local_kv_heads, - head_dim, + group.logical_head_dim, group.buffer.device, group.buffer.dtype, - group.scale_buffer is not None, + "nvfp4" if group.block_scale_buffer is not None else ("fp8" if group.scale_buffer is not None else "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, store_dtype, quantized = geom + num_layers, local_kv_heads, head_dim, device, store_dtype, kv_quant = geom return _alloc_group_storage( num_layers=num_layers, local_kv_heads=local_kv_heads, @@ -352,7 +394,7 @@ def _alloc_group(geom: tuple, outer_size: int, inner_size: int) -> _KVGroupStora store_dtype=store_dtype, outer_size=outer_size, inner_size=inner_size, - quantized=quantized, + kv_quant=kv_quant, ) def rebuild(self, num_full_pages: int, num_swa_tokens: int | None = None) -> None: @@ -445,6 +487,12 @@ def unit_bytes(self) -> tuple[int, int]: 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 + if self.full_kv_pool.block_scale_buffer is not None: + fs = self.full_kv_pool.block_scale_buffer + kv += int(fs.numel() * fs.element_size()) // full_tokens + if self.swa_kv_pool.block_scale_buffer is not None: + ss = self.swa_kv_pool.block_scale_buffer + swa_b += int(ss.numel() * ss.element_size()) // self._swa_num_tokens return kv, swa_b diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index c58fab366..f198887a7 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/python/freetoken/kvcache/mha_pool.py @@ -13,6 +13,8 @@ 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 == "nvfp4": + return torch.uint8 if kv_quant == "fp8": from freetoken.kernel.triton.kv_quant import kv_codes_dtype @@ -37,6 +39,10 @@ class MHAKVCache(BaseKVCachePool): :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. + + ``kv_quant="nvfp4"`` packs two E2M1 values per byte and adds one E4M3 scale + per 16 values, alongside the FP32 row scale. Logical head_dim is retained + separately so rebuild never mistakes the packed width for the model width. """ def __init__( @@ -53,6 +59,10 @@ def __init__( ) -> None: tp_info = get_tp_info() local_kv_heads = div_even(num_kv_heads, tp_info.size, allow_replicate=True) + _kv_store_dtype(dtype, kv_quant) + if kv_quant == "nvfp4" and head_dim % 16: + raise ValueError("NVFP4 KV requires head_dim divisible by 16") + self._head_dim = head_dim self._num_layers = num_layers self.kv_quant = kv_quant self._compute_dtype = dtype @@ -85,8 +95,10 @@ def _alloc( 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": + stored_dim = head_dim // 2 if self.kv_quant == "nvfp4" else head_dim + shape = (2, num_storage_layers, num_pages, page_size, local_kv_heads, stored_dim) + self._block_scale_buffer = None + if self.kv_quant in ("fp8", "nvfp4"): from freetoken.kernel.triton.kv_quant import alloc_codes self._kv_buffer = alloc_codes(shape, self._device) @@ -100,9 +112,14 @@ def _alloc( shape, device=self._device, dtype=self._compute_dtype ) self._scale_buffer = None + if self.kv_quant == "nvfp4": + self._block_scale_buffer = torch.zeros( + (2, num_storage_layers, num_pages * page_size, local_kv_heads, head_dim // 16), + device=self._device, dtype=torch.uint8, + ) 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._storage_shape = (num_pages * page_size, local_kv_heads, stored_dim) def rebuild(self, num_pages: int) -> None: """Reallocate the KV buffer for ``num_pages`` pages IN PLACE. @@ -111,8 +128,10 @@ def rebuild(self, num_pages: int) -> None: existing buffer; only the page count changes. Views and ``_storage_shape`` are 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 + _, num_storage_layers, _old_pages, page_size, local_kv_heads, _stored_dim = self._kv_buffer.shape + head_dim = self._head_dim device = self._device + self._block_scale_buffer = None self._k_buffer = None self._v_buffer = None self._kv_buffer = None @@ -145,6 +164,8 @@ def unit_bytes(self) -> tuple[int, int]: if self._scale_buffer is not None: sc = self._scale_buffer kv += int(sc.numel() * sc.element_size()) // tokens + if self._block_scale_buffer is not None: + kv += self._block_scale_buffer.numel() // tokens return kv, 0 def _dense(self, layer_id: int) -> int: @@ -171,6 +192,16 @@ def v_scale(self, index: int) -> torch.Tensor | None: return None return self._scale_buffer[1][self._dense(index)] + def k_block_scale(self, index: int) -> torch.Tensor | None: + if self._block_scale_buffer is None: + return None + return self._block_scale_buffer[0][self._dense(index)] + + def v_block_scale(self, index: int) -> torch.Tensor | None: + if self._block_scale_buffer is None: + return None + return self._block_scale_buffer[1][self._dense(index)] + def store_kv( self, k: torch.Tensor, @@ -179,6 +210,17 @@ def store_kv( layer_id: int, ) -> None: dense = self._dense(layer_id) + if self.kv_quant == "nvfp4": + from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_to_cache + + quantize_nvfp4_to_cache( + k, v, out_loc, + self._k_buffer[dense].view(self._storage_shape), + self._v_buffer[dense].view(self._storage_shape), + self.k_scale(layer_id), self.v_scale(layer_id), + self.k_block_scale(layer_id), self.v_block_scale(layer_id), + ) + return if self.kv_quant == "fp8": from freetoken.kernel.triton.kv_quant import quantize_kv_to_cache diff --git a/python/freetoken/kvcache/qsa_pool.py b/python/freetoken/kvcache/qsa_pool.py index e79e59009..15d81203a 100644 --- a/python/freetoken/kvcache/qsa_pool.py +++ b/python/freetoken/kvcache/qsa_pool.py @@ -65,6 +65,8 @@ def __init__( layer_ids: Sequence[int] | None = None, kv_quant: str = "none", ) -> None: + if kv_quant not in ("none", "fp8", "nvfp4"): + raise ValueError(f"unsupported QSA kv_quant {kv_quant!r}") if index_ratio < 1 or page_size % index_ratio != 0: # slot // index_ratio only names one group when a group never straddles a page. raise ValueError( @@ -160,6 +162,7 @@ def rebuild(self, num_pages: int) -> 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 + self._block_scale_buffer = None raise @classmethod diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 36e7e231d..54244693a 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -373,13 +373,15 @@ def _infer_reasoning_parser(model_path: str) -> str | None: dest="kv_quant", type=str, default=ServerArgs.kv_quant, - choices=["auto", "bf16", "fp8"], + choices=["auto", "bf16", "fp8", "nvfp4"], 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)." + " 'nvfp4' stores packed E2M1 with block/row scales; supports only" + " paged FULL, hybrid-SWA, or QSA sparse attention with head_dim divisible by 16." ), ) diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py index 9b870273c..34eea2e50 100644 --- a/tests/engine/test_kv_quant_config.py +++ b/tests/engine/test_kv_quant_config.py @@ -94,6 +94,7 @@ def test_kv_quant_spellings(): assert _resolve_kv_quant("auto") == "none" assert _resolve_kv_quant("bf16") == "none" assert _resolve_kv_quant("FP8") == "fp8" + assert _resolve_kv_quant("NVFP4") == "nvfp4" assert _resolve_kv_quant(None) == "none" with pytest.raises(ValueError, match="kv-cache-dtype"): _resolve_kv_quant("q8") @@ -114,6 +115,13 @@ def test_only_the_backends_that_read_scales_declare_fp8_support(): # producing plausible garbage instead of an error. assert fp8 == {"triton", "qsa_sparse"} + nvfp4 = { + name + for name in SUPPORTED_ATTENTION_BACKENDS.supported_names() + if attention_backend_info(name).supports_nvfp4_kv + } + assert nvfp4 == {"triton", "qsa_sparse"} + def test_auto_avoids_the_fast_backends_for_fp8(monkeypatch): from freetoken.engine.engine import _adjust_config @@ -152,6 +160,72 @@ def test_explicit_triton_is_accepted(monkeypatch): assert config.kv_quant == "fp8" +def test_nvfp4_auto_selects_triton(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("full", attention_backend="auto", kv_quant="nvfp4") + _adjust_config(config) + assert config.attention_backend == "triton" + + +@pytest.mark.parametrize("kind", ["mla", "dsa", "dsv4", "bsa"]) +def test_nvfp4_rejects_unsupported_pools_before_allocation(monkeypatch, kind): + from freetoken.engine.engine import _adjust_config + from freetoken.kvcache import create_kvcache_pool + + _patch_fast_machine(monkeypatch) + config = _config(kind, attention_backend="auto", kv_quant="nvfp4") + with pytest.raises(ValueError, match="nvfp4"): + _adjust_config(config) + with pytest.raises(ValueError, match="nvfp4"): + create_kvcache_pool(config.model_config, num_pages=4, page_size=1, + dtype=torch.bfloat16, device=torch.device("cpu"), kv_quant="nvfp4") + + +@pytest.mark.parametrize("backend", ["fi", "fa", "trtllm", "fi,triton", "triton,fi"]) +def test_nvfp4_rejects_backends_without_its_layout(monkeypatch, backend): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + monkeypatch.setattr("freetoken.engine.engine.is_sm100_family", lambda: True) + config = _config("full", attention_backend=backend, kv_quant="nvfp4") + with pytest.raises(ValueError, match="nvfp4"): + _adjust_config(config) + + +def test_nvfp4_rejects_partial_blocks(monkeypatch): + from dataclasses import replace + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("full", attention_backend="auto", kv_quant="nvfp4") + spec = replace(config.model_config.kv_cache_group_specs()[0], head_dim=72) + config.model_config.kv_cache_group_specs = lambda: (spec,) + with pytest.raises(ValueError, match="divisible by 16"): + _adjust_config(config) + + +def test_nvfp4_accepts_hybrid_swa(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("swa", attention_backend="auto", kv_quant="nvfp4") + _adjust_config(config) + assert config.kv_quant == "nvfp4" + assert config.attention_backend == "triton" + + +def test_nvfp4_accepts_qsa(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config("qsa", attention_backend="auto", kv_quant="nvfp4") + _adjust_config(config) + assert config.kv_quant == "nvfp4" + assert config.attention_backend == "qsa_sparse" + + @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 diff --git a/tests/kernels/test_kv_nvfp4.py b/tests/kernels/test_kv_nvfp4.py new file mode 100644 index 000000000..49b4ea0fa --- /dev/null +++ b/tests/kernels/test_kv_nvfp4.py @@ -0,0 +1,255 @@ +"""NVFP4 KV against independent torch rounding and dense attention references.""" + +from types import SimpleNamespace + +import pytest +import torch + +from freetoken.distributed import set_tp_info, try_get_tp_info +from freetoken.kvcache.mha_pool import MHAKVCache + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _pool(dim=128, heads=2, slots=96, layer_ids=None): + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + return MHAKVCache(heads, 4, dim, slots // 4, 4, torch.bfloat16, + torch.device("cuda"), layer_ids=layer_ids, kv_quant="nvfp4") + + +def _reference(x): + shape = x.shape + x = x.float().reshape(*shape[:-1], -1, 16) + row = x.abs().flatten(-2).amax(-1).clamp_min(1e-10) / 2688.0 + block = (x.abs().amax(-1) / (6 * row[..., None])).clamp_max(448).to(torch.float8_e4m3fn) + denom = block.float() * row[..., None] + normalized = torch.where(denom[..., None] > 0, x / denom[..., None], 0) + grid = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6], device=x.device) + distance = (normalized.abs()[..., None] - grid).abs() + nearest = distance == distance.amin(-1, keepdim=True) + codes = torch.arange(8, device=x.device).expand_as(distance) + # Prefer even codes on ties, independently of the kernel's threshold encoding. + rank = torch.where(nearest, codes % 2 * 8 + codes, 32) + code = rank.argmin(-1) | ((normalized < 0).long() * 8) + code = code.reshape(shape) + packed = (code[..., ::2] | (code[..., 1::2] << 4)).to(torch.uint8) + return packed, block.view(torch.uint8), row + + +def _decode(pool, which, layer=1): + codes = getattr(pool, f"{which}_cache")(layer).flatten(0, 1) + block = getattr(pool, f"{which}_block_scale")(layer).view(torch.float8_e4m3fn).float() + row = getattr(pool, f"{which}_scale")(layer) + code = torch.stack((codes & 15, codes >> 4), -1).flatten(-2).long() + grid = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6, + 0, -.5, -1, -1.5, -2, -3, -4, -6], device=codes.device) + return grid[code] * block.repeat_interleave(16, -1) * row[..., None] + + +@pytest.mark.parametrize("dim", [16, 48, 64, 128, 256, 512]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_scatter_matches_independent_reference(dim, dtype): + torch.manual_seed(41) + pool = _pool(dim) + # K is a projection slice; V is independently contiguous. + k = torch.randn(5, 4 * dim, device="cuda", dtype=dtype)[:, :2 * dim] + v = torch.randn(5, 2 * dim, device="cuda", dtype=dtype) + k[0].zero_() + k[1, :16] *= 100 + v[2] *= 1e-5 + loc = torch.tensor([7, 1, 95, 31, 12], device="cuda", dtype=torch.int64) + pool.store_kv(k, v, loc, 1) + for which, source in (("k", k), ("v", v)): + packed, block, row = _reference(source.reshape(5, 2, dim)) + torch.testing.assert_close(getattr(pool, f"{which}_cache")(1).flatten(0, 1)[loc], packed) + torch.testing.assert_close(getattr(pool, f"{which}_block_scale")(1)[loc], block) + torch.testing.assert_close(getattr(pool, f"{which}_scale")(1)[loc], row) + assert torch.isfinite(_decode(pool, which)).all() + assert torch.count_nonzero(_decode(pool, which)[0]) == 0 + + +def test_e2m1_grid_and_round_to_even_boundaries(): + positive = torch.tensor([0, .25, .5, .75, 1, 1.25, 1.5, 1.75, + 2, 2.5, 3, 3.5, 4, 5, 6], device="cuda") + values = torch.cat((positive, -positive)) + values = torch.cat((values, values.nextafter(torch.full_like(values, float("inf"))), + values.nextafter(torch.full_like(values, -float("inf"))))) + pool = _pool(dim=32, slots=192) + rows = torch.zeros(values.numel(), 2, 32, device="cuda") + rows[:, :, 0] = values[:, None] + rows[:, :, 15] = 6 + rows[:, :, 31] = 2688 # Forces row_scale=1 and the first block_scale=1. + loc = torch.arange(values.numel(), device="cuda", dtype=torch.int32) + pool.store_kv(rows.flatten(1), rows.flatten(1), loc, 1) + packed, block, row = _reference(rows) + torch.testing.assert_close(pool.k_cache(1).flatten(0, 1)[loc], packed) + torch.testing.assert_close(pool.k_block_scale(1)[loc], block) + torch.testing.assert_close(pool.k_scale(1)[loc], row) + + +@pytest.mark.parametrize("dim", [64, 128, 256]) +@pytest.mark.parametrize("mode", ["paged", "decode", "extend", "split"]) +def test_attention_reads_packed_cache(dim, mode): + from freetoken.kernel.triton.attention import ( + decode_paged_attention, extend_paged_attention, paged_attention, + ) + + torch.manual_seed(42) + pool = _pool(dim) + n, prefix = 67, 35 + k, v = [torch.randn(n, 2 * dim, device="cuda", dtype=torch.bfloat16) for _ in range(2)] + loc = torch.randperm(95, device="cuda")[:n] + 1 + pool.store_kv(k, v, loc, 1) + tokens = 1 if mode in ("paged", "decode") else n - prefix + q = torch.randn(tokens, 6, dim, device="cuda", dtype=torch.bfloat16) + indptr = torch.tensor([0, n], device="cuda", dtype=torch.int32) + pos = torch.arange(n - tokens, n, device="cuda") + args = dict(q=q, k_cache=pool.k_cache(1).flatten(0, 1), + v_cache=pool.v_cache(1).flatten(0, 1), k_scale=pool.k_scale(1), + v_scale=pool.v_scale(1), k_block_scale=pool.k_block_scale(1), + v_block_scale=pool.v_block_scale(1), kv_quant="nvfp4", sm_scale=dim ** -.5) + kd, vd = _decode(pool, "k")[loc], _decode(pool, "v")[loc] + if mode == "split": + kd[prefix:] = k[prefix:].view(-1, 2, dim).float() + vd[prefix:] = v[prefix:].view(-1, 2, dim).float() + # Tensor-core paths round restored K/V to the query dtype before dot products. + if mode != "paged": + kd, vd = kd.to(q.dtype).float(), vd.to(q.dtype).float() + kd, vd = [x.repeat_interleave(3, 1).transpose(0, 1) for x in (kd, vd)] + score = torch.einsum("thd,hnd->thn", q.float(), kd) * dim ** -.5 + score.masked_fill_(torch.arange(n, device="cuda")[None, None, :] > pos[:, None, None], -float("inf")) + ref = torch.einsum("thn,hnd->thd", score.softmax(-1), vd).to(q.dtype) + if mode == "paged": + actual = paged_attention(**args, indptr=indptr, indices=loc, + q_to_req=torch.zeros(tokens, device="cuda", dtype=torch.int32), q_positions=pos) + elif mode == "decode": + actual = decode_paged_attention(**args, indptr=indptr, indices=loc, q_positions=pos, + attn_logits=torch.empty(1, 6, 8, dim, device="cuda"), + attn_lse=torch.empty(1, 6, 8, device="cuda"), + num_kv_splits=torch.tensor([8], device="cuda", dtype=torch.int32), max_kv_splits=8) + else: + extra = {} if mode == "extend" else dict(k_extend=k[prefix:].view(-1, 2, dim), v_extend=v[prefix:].view(-1, 2, dim)) + actual = extend_paged_attention(**args, **extra, + qo_indptr=torch.tensor([0, tokens], device="cuda", dtype=torch.int32), + kv_indptr=indptr, kv_indices=loc, + prefix_lens=torch.tensor([prefix], device="cuda", dtype=torch.int32), max_q_len=tokens) + torch.testing.assert_close(actual, ref, atol=0.008, rtol=0.025) + + +def test_pool_budget_rebuild_and_layer_mapping(): + from freetoken.kvcache.base import spec_kv_bytes_per_token + from freetoken.models.config import KVCacheGroupSpec + + pool = _pool(layer_ids=(1, 3)) + spec = KVCacheGroupSpec(name="full", layer_ids=(1, 3), num_kv_heads=2, head_dim=128, sliding_window=None) + cfg = SimpleNamespace(kv_quant="nvfp4", dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1)) + assert spec_kv_bytes_per_token(spec, cfg) == pool.unit_bytes()[0] == 2 * 2 * 2 * 76 + with pytest.raises(KeyError): + pool.k_block_scale(0) + for pages in (32, 8): + pool.rebuild(pages) + assert pool.k_cache(3).shape == (pages, 4, 2, 64) + assert pool.k_block_scale(3).shape == (pages * 4, 2, 8) + assert pool.unit_bytes()[0] == 608 + assert torch.count_nonzero(_decode(pool, "k", 3)) == 0 + + +def test_store_cuda_graph_replay_changes_slots(): + pool = _pool() + k = torch.randn(2, 256, device="cuda", dtype=torch.bfloat16) + loc = torch.tensor([1, 2], device="cuda", dtype=torch.int32) + pool.store_kv(k, k, loc, 1) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + pool.store_kv(k, k, loc, 1) + k.mul_(2) + loc.copy_(torch.tensor([4, 9], device="cuda", dtype=torch.int32)) + graph.replay() + expected, _, _ = _reference(k.view(2, 2, 128)) + torch.testing.assert_close(pool.k_cache(1).flatten(0, 1)[loc], expected) + + +def test_hybrid_swa_pool_packs_both_groups_and_rebuilds(monkeypatch): + 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("full", (1,), 2, 64, None), + KVCacheGroupSpec("swa", (0,), 2, 128, 32), + ) + pool = HybridSWAKVCache( + groups=groups, + num_layers=2, + num_full_pages=12, + page_size=1, + num_swa_tokens=18, + dtype=torch.bfloat16, + device=torch.device("cuda"), + kv_quant="nvfp4", + ) + assert pool.k_cache(1).shape == (12, 1, 2, 32) + assert pool.k_cache(0).shape == (18, 1, 2, 64) + assert pool.k_block_scale(1).shape == (12, 2, 4) + assert pool.k_block_scale(0).shape == (18, 2, 8) + assert pool.unit_bytes() == (160, 304) + + full_loc = torch.tensor([4, 9], device="cuda", dtype=torch.int32) + swa_loc = torch.tensor([2, 7], device="cuda", dtype=torch.int32) + pool.alloc_swa(swa_loc) + for layer, dim, loc, cache_loc in ((1, 64, full_loc, full_loc), (0, 128, swa_loc, None)): + rows = torch.randn(2, 2 * dim, device="cuda", dtype=torch.bfloat16) + pool.store_kv(rows, rows, loc, layer) + if cache_loc is None: + cache_loc = pool.translate_loc_from_full_to_swa(loc) + expected, block, row = _reference(rows.view(2, 2, dim)) + codes = pool.k_cache(layer).flatten(0, 1)[cache_loc] + torch.testing.assert_close(codes, expected) + torch.testing.assert_close(pool.k_block_scale(layer)[cache_loc], block) + torch.testing.assert_close(pool.k_scale(layer)[cache_loc], row) + + pool.rebuild(num_full_pages=6, num_swa_tokens=10) + assert pool.k_cache(1).shape == (6, 1, 2, 32) + assert pool.k_cache(0).shape == (10, 1, 2, 64) + assert pool.unit_bytes() == (160, 304) + assert pool.swa_available_size() == 9 + + +def test_backend_decode_graph_replays_new_kv_and_page_tables(monkeypatch): + from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata + from freetoken.kernel.triton.attention import paged_attention + + pool = _pool() + monkeypatch.setattr("freetoken.attention.triton.get_global_ctx", + lambda: SimpleNamespace(kv_cache=pool)) + backend = TritonAttentionBackend(SimpleNamespace(num_qo_heads=6, head_dim=128)) + q = torch.randn(2, 6, 128, device="cuda", dtype=torch.bfloat16) + k, v = [torch.randn(2, 256, device="cuda", dtype=q.dtype) for _ in range(2)] + loc = torch.tensor([7, 11], device="cuda", dtype=torch.int32) + indptr = torch.tensor([0, 1, 2], device="cuda", dtype=torch.int32) + positions = torch.zeros(2, device="cuda", dtype=torch.int32) + indices = loc.clone() + req = torch.arange(2, device="cuda", dtype=torch.int32) + metadata = TritonMetadata(cu_seqlens_q_gpu=indptr, indptr=indptr, indices=indices, + q_to_req=req, q_positions=positions, is_decode=True, prefix_lens=positions, + max_q_len=1) + batch = SimpleNamespace(out_loc=loc, attn_metadata=metadata) + backend.forward(q, k, v, 1, batch) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = backend.forward(q, k, v, 1, batch) + for new_slots in ([31, 44], [7, 11]): + k.normal_() + v.normal_() + q.normal_() + loc.copy_(torch.tensor(new_slots, device="cuda", dtype=torch.int32)) + indices.copy_(loc) + graph.replay() + ref = paged_attention(q, _decode(pool, "k").to(q.dtype), + _decode(pool, "v").to(q.dtype), indptr, indices, req, positions, 128 ** -.5) + torch.testing.assert_close(actual, ref, atol=0.008, rtol=0.025) diff --git a/tests/kernels/test_qsa_nvfp4.py b/tests/kernels/test_qsa_nvfp4.py new file mode 100644 index 000000000..591c6f1d4 --- /dev/null +++ b/tests/kernels/test_qsa_nvfp4.py @@ -0,0 +1,141 @@ +"""QSA sparse attention over the packed NVFP4 KV tier.""" + +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_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 +HEAD_DIM = 64 + + +def _layout(rows: int, topk: int): + block_table = torch.tensor( + [[0, 1, 2], [2, 1, 0]], dtype=torch.int32, device="cuda" + ).contiguous() + indices = ( + torch.arange(topk, dtype=torch.int32, device="cuda")[None, :] + 32 + ).repeat(rows, 1).contiguous() + token_to_req = torch.arange(rows, dtype=torch.int32, device="cuda") % 2 + return indices, block_table, token_to_req.contiguous() + + +def _decode(codes: torch.Tensor, row_scale: torch.Tensor, block_scale: torch.Tensor): + """Decode with torch, independently of the Triton read path under test.""" + lo = (codes & 15).to(torch.long) + hi = (codes >> 4).to(torch.long) + code = torch.stack((lo, hi), dim=-1).reshape(*codes.shape[:-1], HEAD_DIM) + magnitude = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device=codes.device + )[code & 7] + value = torch.where(code < 8, magnitude, -magnitude) + block = block_scale.view(torch.float8_e4m3fn).to(torch.float32) + return (value * block.repeat_interleave(16, dim=-1) * row_scale.unsqueeze(-1)).to( + torch.bfloat16 + ) + + +@pytest.mark.parametrize(("rows", "kv_heads"), [(1, 1), (16, 2)]) +def test_qsa_nvfp4_matches_its_bf16_decode(rows, kv_heads): + """Page-table indirection must address the same packed slot and both scale tiers.""" + torch.manual_seed(19) + pages, query_heads, slots = 3, 2 * kv_heads, 3 * PAGE + k = torch.randn(slots, kv_heads * HEAD_DIM, device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) * 0.25 + out_loc = torch.arange(slots, dtype=torch.int32, device="cuda") + code_shape = (slots, kv_heads, HEAD_DIM // 2) + block_shape = (slots, kv_heads, HEAD_DIM // 16) + k_codes = torch.zeros(code_shape, dtype=torch.uint8, device="cuda") + v_codes = torch.zeros_like(k_codes) + k_scale = torch.zeros((slots, kv_heads), dtype=torch.float32, device="cuda") + v_scale = torch.zeros_like(k_scale) + k_block = torch.zeros(block_shape, dtype=torch.uint8, device="cuda") + v_block = torch.zeros_like(k_block) + quantize_nvfp4_to_cache( + k, v, out_loc, k_codes, v_codes, k_scale, v_scale, k_block, v_block + ) + torch.cuda.synchronize() + + q = torch.randn(rows, query_heads, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(rows, topk=64) + packed_shape = (pages, PAGE, kv_heads, HEAD_DIM // 2) + dense_shape = (pages, PAGE, kv_heads, HEAD_DIM) + got = qsa_sparse_paged_attention( + q, + k_codes.view(packed_shape), + v_codes.view(packed_shape), + indices, + block_table, + token_to_req, + k_scale=k_scale, + v_scale=v_scale, + kv_quant="nvfp4", + k_block_scale=k_block, + v_block_scale=v_block, + ) + want = qsa_sparse_paged_attention( + q, + _decode(k_codes, k_scale, k_block).view(dense_shape), + _decode(v_codes, v_scale, v_block).view(dense_shape), + indices, + block_table, + token_to_req, + ) + assert torch.equal(got, want), ( + "NVFP4 QSA attend diverged from its bf16 decode (max diff " + f"{(got.float() - want.float()).abs().max().item():.3e})" + ) + + +def test_qsa_nvfp4_requires_both_block_scale_tensors(): + q = torch.zeros(1, 2, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + codes = torch.zeros(2, PAGE, 1, HEAD_DIM // 2, device="cuda", dtype=torch.uint8) + scales = torch.ones(2 * PAGE, 1, device="cuda", dtype=torch.float32) + indices = torch.zeros(1, 8, device="cuda", dtype=torch.int32) + block_table = torch.zeros(1, 2, device="cuda", dtype=torch.int32) + token_to_req = torch.zeros(1, device="cuda", dtype=torch.int32) + + with pytest.raises(ValueError, match="NVFP4"): + qsa_sparse_paged_attention( + q, + codes, + codes, + indices, + block_table, + token_to_req, + k_scale=scales, + v_scale=scales, + kv_quant="nvfp4", + ) + + +def test_qsa_nvfp4_splitk_reads_value_rows(): + """A zero query makes the split-K result the selected V-row average.""" + torch.manual_seed(23) + slots = 3 * PAGE + v = torch.randn(slots, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + codes = torch.zeros((slots, 1, HEAD_DIM // 2), dtype=torch.uint8, device="cuda") + row = torch.zeros((slots, 1), dtype=torch.float32, device="cuda") + block = torch.zeros((slots, 1, HEAD_DIM // 16), dtype=torch.uint8, device="cuda") + quantize_nvfp4_to_cache( + v, v, torch.arange(slots, dtype=torch.int32, device="cuda"), codes, codes, + row, row, block, block, + ) + q = torch.zeros(1, 2, HEAD_DIM, device="cuda", dtype=torch.bfloat16) + indices, block_table, token_to_req = _layout(1, topk=64) + args = (indices, block_table, token_to_req) + got = qsa_sparse_paged_attention( + q, codes.view(3, PAGE, 1, HEAD_DIM // 2), codes.view(3, PAGE, 1, HEAD_DIM // 2), + *args, k_scale=row, v_scale=row, kv_quant="nvfp4", k_block_scale=block, + v_block_scale=block, + ) + ref = _decode(codes, row, block).view(3, PAGE, 1, HEAD_DIM) + want = qsa_sparse_paged_attention(q, ref, ref, *args) + assert torch.equal(got, want) diff --git a/tests/kvcache/test_qsa_pool_fp8.py b/tests/kvcache/test_qsa_pool_fp8.py index 04426b876..c0a91f862 100644 --- a/tests/kvcache/test_qsa_pool_fp8.py +++ b/tests/kvcache/test_qsa_pool_fp8.py @@ -115,7 +115,7 @@ def test_fp8_replaces_the_kv_slab_and_adds_scale_views(): assert bf16.k_scale(3) is None and bf16.v_scale(3) is None -@pytest.mark.parametrize("kv_quant", ["none", "fp8"]) +@pytest.mark.parametrize("kv_quant", ["none", "fp8", "nvfp4"]) 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.""" @@ -134,6 +134,9 @@ def test_index_tiers_stay_16_bit_whatever_the_kv_store_does(kv_quant): assert plain == kv_16bit + index_term if kv_quant == "fp8": assert cost == kv_16bit // 2 + scale_term + index_term + elif kv_quant == "nvfp4": + block_term = 2 * kv_layers * HEADS * (DIM // 16) + assert cost == kv_16bit // 4 + scale_term + block_term + index_term else: assert cost == plain @@ -207,3 +210,12 @@ def test_factory_threads_kv_quant_into_the_qsa_pool(): 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 + +def test_nvfp4_replaces_only_qsa_kv_tiers(): + pool = _pool(kv_quant="nvfp4") + slots = 4 * PAGE_SIZE + assert pool.k_cache(3).shape == (4, PAGE_SIZE, HEADS, DIM // 2) + assert pool.k_block_scale(3).shape == (slots, HEADS, DIM // 16) + assert pool.v_block_scale(3).dtype is torch.uint8 + assert pool.cmp_k_cache(0).dtype is torch.bfloat16 + From ca3675ecde8d53385ddb32cf4d611c7230d0b897 Mon Sep 17 00:00:00 2001 From: ArqAlice Date: Mon, 7 Sep 2026 18:57:18 +0900 Subject: [PATCH 10/12] test(kernels): stabilize fp8 extend attention regression --- tests/kernels/test_triton_attention.py | 33 +++++++++++++++++--------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index 61ce5360d..8241bcdc9 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -976,22 +976,27 @@ def test_extend_paged_attention_decodes_fp8_scales(use_split_inputs: bool): 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( + k_extend = (torch.randn(total_q, num_kv_heads, head_dim, device=device) * 3.0).to( torch.bfloat16 ) - v_extend = (torch.randn(total_q, num_kv_heads, head_dim, device=device) * 6.0).to( + v_extend = (torch.randn(total_q, num_kv_heads, head_dim, device=device) * 3.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( + v_cache = (torch.randn(total_kv, num_kv_heads, head_dim, device=device) * 3.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) + # KV positions are logical, but FP8 codes and their scales are addressed by the + # physical slots from the page table. A contiguous table masks a regression that + # looks scales up with the logical position instead of the slot. + indices = torch.tensor( + [10, 3, 8, 1, 9, 0, 7, 2, 6, 4, 5], 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) @@ -1012,18 +1017,24 @@ def test_extend_paged_attention_decodes_fp8_scales(use_split_inputs: bool): 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) + num_slots = total_kv + 1 + k_slots = torch.zeros( + num_slots, num_kv_heads, head_dim, dtype=k_cache.dtype, device=device + ) + v_slots = torch.zeros_like(k_slots) + k_slots[indices.to(torch.long)] = k_cache + v_slots[indices.to(torch.long)] = v_cache + k_codes, v_codes, k_scale, v_scale = _fp8_cache(k_slots, v_slots) 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 - ] + current_slots = indices[ + kv_off + cached_len : kv_off + cached_len + extend_len + ].to(torch.long) + k_ref[current_slots] = k_extend[q_off : q_off + extend_len].float() + v_ref[current_slots] = v_extend[q_off : q_off + extend_len].float() q_off += extend_len kv_off += cached_len + extend_len From 9b103b04f9c8a1544dbe857013dd129170defbd7 Mon Sep 17 00:00:00 2001 From: ArqAlice Date: Tue, 8 Sep 2026 08:48:57 +0900 Subject: [PATCH 11/12] feat(kvcache): add fp8 support for dsa kv cache --- python/freetoken/attention/__init__.py | 5 +- python/freetoken/attention/dsa.py | 1 + python/freetoken/engine/engine.py | 13 ++-- .../freetoken/kernel/triton/glm_dsa_sparse.py | 65 +++++++++++++++---- python/freetoken/kernel/triton/kv_quant.py | 52 +++++++++++++++ python/freetoken/kvcache/__init__.py | 3 +- python/freetoken/kvcache/dsa_pool.py | 48 ++++++++++++-- tests/attention/test_dsa_kpool.py | 16 +++++ tests/engine/test_kv_quant_config.py | 22 +++++-- tests/kvcache/test_dsa_pool.py | 40 +++++++++++- tests/kvcache/test_kv_cache_rebuild.py | 7 ++ tests/kvcache/test_mha_pool_fp8.py | 25 ++++--- 12 files changed, 247 insertions(+), 50 deletions(-) diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 012759ad3..31de4ba96 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -110,7 +110,10 @@ def create_dsv4_sparse_backend(config: ModelConfig): @SUPPORTED_ATTENTION_BACKENDS.register( "dsa", - BackendInfo(supported_types=frozenset({AttnType.MLA, AttnType.DSA})), + BackendInfo( + supported_types=frozenset({AttnType.MLA, AttnType.DSA}), + supports_fp8_kv=True, + ), ) def create_dsa_backend(config: ModelConfig): # MLA with a grouped index (index_ratio > 1) is the kpool indexer layout. diff --git a/python/freetoken/attention/dsa.py b/python/freetoken/attention/dsa.py index 398d17002..9aedb3ed7 100644 --- a/python/freetoken/attention/dsa.py +++ b/python/freetoken/attention/dsa.py @@ -201,6 +201,7 @@ def _attend( return glm_dsa_sparse_attn( q_cat, self.kvcache.latent_rows(layer_id), sel, self.sm_scale, counts=cnt, d_v=self.kv_lora_rank, + pool_scale=self.kvcache.latent_scale(layer_id), ) def mla_forward( diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index cc87d95bf..f0e6bfbb4 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1360,18 +1360,17 @@ def override(attr: str, value: Any): # this is dangerous, use with caution 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. + # kernel: plain paged and hybrid-SWA, QSA sparse, and DSA/MLA. QSA's index + # tier and DSA's index-key/tail tiers stay bf16; the DSA kernel dequantizes + # selected latent rows with their per-token scale. Other sparse families have + # no scale-read path and remain rejected before weights load. quant_unsupported = required_attn_types - { - AttnType.FULL, AttnType.SWA, AttnType.QSA, + AttnType.FULL, AttnType.SWA, AttnType.QSA, AttnType.MLA, AttnType.DSA, } 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 " + "hybrid-SWA, QSA sparse and DSA/MLA KV pools; this model also needs " f"{', '.join(sorted(t.value for t in quant_unsupported))} attention " "(use --kv-cache-dtype bf16)." ) diff --git a/python/freetoken/kernel/triton/glm_dsa_sparse.py b/python/freetoken/kernel/triton/glm_dsa_sparse.py index 8aaeeb9ba..4415a68e9 100644 --- a/python/freetoken/kernel/triton/glm_dsa_sparse.py +++ b/python/freetoken/kernel/triton/glm_dsa_sparse.py @@ -28,6 +28,8 @@ import triton import triton.language as tl +from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 + BLOCK_H = 16 BLOCK_T = 32 MAX_SPLITS = 32 @@ -36,11 +38,12 @@ @triton.jit def _glm_dsa_sparse_kernel( - q_ptr, pool_ptr, o_ptr, idx_ptr, cnt_ptr, + q_ptr, pool_ptr, pool_scale_ptr, o_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, stride_qb, stride_qm, stride_qh, stride_qd, stride_pn, stride_pd, + stride_ps, stride_ob, stride_om, stride_oh, stride_od, stride_ib, stride_im, stride_it, stride_nb, stride_nm, @@ -50,6 +53,7 @@ def _glm_dsa_sparse_kernel( BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, HAS_ROPE: tl.constexpr, + HAS_FP8: tl.constexpr, ): pid_m = tl.program_id(0) pid_b = tl.program_id(1) @@ -82,11 +86,22 @@ def _glm_dsa_sparse_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 kv_base = pool_ptr + idxs[:, None] * stride_pn - kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_FP8: + row_scale = tl.load(pool_scale_ptr + idxs * stride_ps, mask=valid, other=0.0) + kv_v = kv_load_e4m3_tile_f32( + kv_base + offs_v[None, :] * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores = tl.dot(q_v, tl.trans(kv_v)) if HAS_ROPE: - kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_FP8: + kv_r = kv_load_e4m3_tile_f32( + kv_base + (D_V + offs_r[None, :]) * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores += tl.dot(q_r, tl.trans(kv_r)) scores = scores * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -200,11 +215,12 @@ def glm_dsa_decode_logits( @triton.jit def _glm_dsa_splitk_kernel( - q_ptr, pool_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, + q_ptr, pool_ptr, pool_scale_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, stride_qb, stride_qm, stride_qh, stride_qd, stride_pn, stride_pd, + stride_ps, stride_mb, stride_mm, stride_mh, stride_ms, stride_md, stride_lb, stride_lm, stride_lh, stride_ls, stride_ib, stride_im, stride_it, @@ -215,6 +231,7 @@ def _glm_dsa_splitk_kernel( BLOCK_T: tl.constexpr, HAS_COUNTS: tl.constexpr, HAS_ROPE: tl.constexpr, + HAS_FP8: tl.constexpr, NUM_SPLITS: tl.constexpr, ): """Stage 1 (decode flash-decoding): each program reduces one BLOCK_T-aligned slice of @@ -256,11 +273,22 @@ def _glm_dsa_splitk_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 kv_base = pool_ptr + idxs[:, None] * stride_pn - kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_FP8: + row_scale = tl.load(pool_scale_ptr + idxs * stride_ps, mask=valid, other=0.0) + kv_v = kv_load_e4m3_tile_f32( + kv_base + offs_v[None, :] * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_v = tl.load(kv_base + offs_v[None, :] * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores = tl.dot(q_v, tl.trans(kv_v)) if HAS_ROPE: - kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) + if HAS_FP8: + kv_r = kv_load_e4m3_tile_f32( + kv_base + (D_V + offs_r[None, :]) * stride_pd, valid[:, None] + ) * row_scale[:, None] + else: + kv_r = tl.load(kv_base + (D_V + offs_r[None, :]) * stride_pd, mask=valid[:, None], other=0.0).to(tl.float32) scores += tl.dot(q_r, tl.trans(kv_r)) scores = scores * scale scores = tl.where(valid[None, :], scores, -float("inf")) @@ -351,6 +379,7 @@ def glm_dsa_sparse_attn( softmax_scale: float, counts: torch.Tensor | None = None, # [b, m] int32 live columns per query (device-read) d_v: int = 512, + pool_scale: torch.Tensor | None = None, # [rows] fp32, one scale per quantized latent row force_splits: int | None = None, # tests only: 0 = single-program, N = split-k N ) -> torch.Tensor: """Sparse MLA attention over gathered latent rows; returns ``[b, m, h, d_v]``. @@ -364,6 +393,14 @@ def glm_dsa_sparse_attn( d_r = d - d_v topk = topk_idxs.shape[-1] assert pool.shape[-1] == d, (pool.shape, d) + has_fp8 = pool_scale is not None + if has_fp8: + assert pool.dtype == torch.uint8, pool.dtype + assert pool_scale.shape == (pool.shape[0],), (pool_scale.shape, pool.shape) + assert pool_scale.dtype is torch.float32, pool_scale.dtype + scale_pool = pool_scale.contiguous() + else: + scale_pool = pool q = q.contiguous() pool_2d = pool.reshape(-1, d) assert pool_2d.stride(-1) == 1 @@ -385,19 +422,22 @@ def glm_dsa_sparse_attn( mid_lse = q.new_empty(b, m, h, n_splits, dtype=torch.float32) grid1 = (m * n_splits, b, triton.cdiv(h, BLOCK_H)) _glm_dsa_splitk_kernel[grid1]( - q, pool_2d, mid_o, mid_lse, idx, cnt, + q, pool_2d, scale_pool, mid_o, mid_lse, idx, cnt, float(softmax_scale), h, topk, q.stride(0), q.stride(1), q.stride(2), q.stride(3), pool_2d.stride(0), pool_2d.stride(1), + scale_pool.stride(0) if has_fp8 else 0, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4), mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3), idx.stride(0), 0 if broadcast_m else idx.stride(1), idx.stride(2), stride_nb, stride_nm, D_V=d_v, D_R=d_r, BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, - HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, NUM_SPLITS=n_splits, - num_warps=4, num_stages=2, + HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, NUM_SPLITS=n_splits, + # The 512-wide latent accumulator exceeds the 99 KiB shared-memory + # limit of consumer Blackwell GPUs with a two-stage pipeline. + num_warps=4, num_stages=1, ) grid2 = (m, b, h) _glm_dsa_merge_kernel[grid2]( @@ -412,18 +452,19 @@ def glm_dsa_sparse_attn( grid = (m, b, triton.cdiv(h, BLOCK_H)) _glm_dsa_sparse_kernel[grid]( - q, pool_2d, o, idx, cnt, + q, pool_2d, scale_pool, o, idx, cnt, float(softmax_scale), h, topk, q.stride(0), q.stride(1), q.stride(2), q.stride(3), pool_2d.stride(0), pool_2d.stride(1), + scale_pool.stride(0) if has_fp8 else 0, o.stride(0), o.stride(1), o.stride(2), o.stride(3), idx.stride(0), 0 if broadcast_m else idx.stride(1), idx.stride(2), stride_nb, stride_nm, D_V=d_v, D_R=d_r, BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, - HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, - num_warps=4, num_stages=2, + HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, + num_warps=4, num_stages=1, ) return o diff --git a/python/freetoken/kernel/triton/kv_quant.py b/python/freetoken/kernel/triton/kv_quant.py index e041e0466..c10ad9c4b 100644 --- a/python/freetoken/kernel/triton/kv_quant.py +++ b/python/freetoken/kernel/triton/kv_quant.py @@ -185,6 +185,57 @@ def quantize_kv_to_cache( ) +@triton.jit +def _kv_quant_rows_scatter_kernel( + src, dst, scale_ptr, idx_ptr, + stride_src, stride_dst, stride_scale, + D: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """Quantize one whole latent row per program and scatter it by ``out_loc``.""" + t = tl.program_id(0) + pos = tl.load(idx_ptr + t).to(tl.int64) + d = tl.arange(0, BLOCK_D) + mask = d < D + x = tl.load(src + t * stride_src + d, mask=mask, other=0.0).to(tl.float32) + s = tl.maximum(tl.max(tl.abs(x), axis=0), 1e-10) / 448.0 + q = tl.clamp(x / s, -448.0, 448.0) + tl.store(dst + pos * stride_dst + d, e4m3_f32_to_u8(round_e4m3(q)), mask=mask) + tl.store(scale_ptr + pos * stride_scale, s) + + +def quantize_rows_to_cache( + rows: torch.Tensor, + out_loc: torch.Tensor, + cache: torch.Tensor, + scales: torch.Tensor, +) -> None: + """Quantize and scatter one scale-bearing FP8 row per token. + + MLA stores its absorbed K/V state as one latent row, rather than separate K and + V heads. Its scale granularity is therefore one FP32 value per ``(token, + layer)`` row, which is the single-head MLA case already priced by + ``kv_scale_bytes_per_token``. + """ + tokens = rows.shape[0] + if tokens == 0: + return + assert rows.dim() == 2 and rows.stride(1) == 1, tuple(rows.shape) + assert cache.dim() == 2 and cache.shape[1] == rows.shape[1], ( + tuple(cache.shape), tuple(rows.shape), + ) + assert scales.shape == (cache.shape[0],), (tuple(scales.shape), tuple(cache.shape)) + assert cache.dtype == kv_codes_dtype(), cache.dtype + assert scales.dtype == KV_SCALE_DTYPE, scales.dtype + dim = rows.shape[1] + _kv_quant_rows_scatter_kernel[(tokens,)]( + rows, cache, scales, out_loc, + rows.stride(0), cache.stride(0), scales.stride(0), + D=dim, BLOCK_D=triton.next_power_of_2(dim), + num_warps=4 if dim > 256 else 1, + ) + + __all__ = [ "FP8", "KV_QUANT_FP8", @@ -193,5 +244,6 @@ def quantize_kv_to_cache( "codes_to_f32", "kv_codes_dtype", "quantize_kv_to_cache", + "quantize_rows_to_cache", ] diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index a7510d69f..94b71cdc0 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -229,7 +229,6 @@ def create_kvcache_pool( 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). @@ -245,6 +244,7 @@ def create_kvcache_pool( index_head_dim=spec.index_head_dim, num_index_layers=spec.num_index_layers, layer_ids=layer_ids, + kv_quant=kv_quant, ) if spec.index_ratio > 1: # kpool tail rings are keyed by Req.table_idx; + 1 covers the dummy request row. @@ -264,6 +264,7 @@ def create_kvcache_pool( dtype=dtype, device=device, layer_ids=layer_ids, + kv_quant=kv_quant, ) spec = kv_specs[0] if len(kv_specs) == 1 else None diff --git a/python/freetoken/kvcache/dsa_pool.py b/python/freetoken/kvcache/dsa_pool.py index 13225cec0..311b27fb4 100644 --- a/python/freetoken/kvcache/dsa_pool.py +++ b/python/freetoken/kvcache/dsa_pool.py @@ -43,6 +43,7 @@ def __init__( dtype: torch.dtype, device: torch.device, layer_ids: "tuple[int, ...] | None" = None, + kv_quant: str = "none", ) -> None: self._latent_dim = latent_dim if layer_ids is None: @@ -54,6 +55,7 @@ def __init__( self._page_size = page_size self._dtype = dtype self._device = device + self.kv_quant = kv_quant self._alloc(num_pages) def _local_layer(self, layer_id: int) -> int: @@ -61,11 +63,21 @@ def _local_layer(self, layer_id: int) -> int: def _alloc(self, num_pages: int) -> None: self._num_pages = num_pages - self._kv_buffer = torch.empty( - (1, self._num_layers, num_pages, self._page_size, 1, self._latent_dim), - device=self._device, - dtype=self._dtype, - ) + shape = (1, self._num_layers, num_pages, self._page_size, 1, self._latent_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( + (self._num_layers, num_pages * self._page_size), + device=self._device, + dtype=torch.float32, + ) + elif self.kv_quant == "none": + self._kv_buffer = torch.empty(shape, device=self._device, dtype=self._dtype) + self._scale_buffer = None + else: + raise ValueError(f"unknown kv_quant {self.kv_quant!r}") # -- views (addressed by GLOBAL layer id; remapped when layer_ids was given) -- def k_cache(self, layer_id: int) -> torch.Tensor: @@ -82,6 +94,12 @@ def latent_rows(self, layer_id: int) -> torch.Tensor: """Row-flat latent view ``[num_pages * page_size, latent_dim]``.""" return self._kv_buffer[0, self._local_layer(layer_id)].view(-1, self._latent_dim) + def latent_scale(self, layer_id: int) -> torch.Tensor | None: + """One FP32 scale per latent row, or ``None`` for the compute-dtype pool.""" + if self._scale_buffer is None: + return None + return self._scale_buffer[self._local_layer(layer_id)] + # -- writes ----------------------------------------------------------------- def store_kv( self, @@ -98,6 +116,12 @@ def store_kv( """ rows = self.latent_rows(layer_id) split = rows.shape[1] - k_rope.shape[-1] + if self.kv_quant == "fp8": + from freetoken.kernel.triton.kv_quant import quantize_rows_to_cache + + latent = torch.cat((c_kv, k_rope), dim=-1).contiguous() + quantize_rows_to_cache(latent, out_loc, rows, self.latent_scale(layer_id)) + return rows[out_loc, :split] = c_kv rows[out_loc, split:] = k_rope @@ -105,6 +129,7 @@ def rebuild(self, num_pages: int) -> None: """In-place resize (frees the old slab first; object identity preserved -- callers re-derive views per forward, same contract as MHAKVCache.rebuild).""" self._kv_buffer = None + self._scale_buffer = None if self._device.type == "cuda": torch.cuda.synchronize(self._device) torch.cuda.empty_cache() @@ -127,7 +152,11 @@ def rebuild_from_config( def unit_bytes(self) -> tuple[int, int]: buf = self._kv_buffer - return int(buf.numel() * buf.element_size()) // (self._num_pages * self._page_size), 0 + tokens = self._num_pages * self._page_size + kv = int(buf.numel() * buf.element_size()) // tokens + if self._scale_buffer is not None: + kv += int(self._scale_buffer.numel() * self._scale_buffer.element_size()) // tokens + return kv, 0 # -- pool properties ---------------------------------------------------------- @property @@ -138,6 +167,10 @@ def device(self) -> torch.device: def dtype(self) -> torch.dtype: return self._dtype + @property + def store_dtype(self) -> torch.dtype: + return self._kv_buffer.dtype + @property def num_layers(self) -> int: return self._num_layers @@ -158,12 +191,13 @@ def __init__( index_head_dim: int, num_index_layers: int, layer_ids: "tuple[int, ...] | None" = None, + kv_quant: str = "none", ) -> None: self._index_head_dim = index_head_dim self._num_index_layers = num_index_layers super().__init__( latent_dim, num_layers, num_pages, page_size, dtype, device, - layer_ids=layer_ids, + layer_ids=layer_ids, kv_quant=kv_quant, ) def _index_rows(self, num_pages: int) -> int: diff --git a/tests/attention/test_dsa_kpool.py b/tests/attention/test_dsa_kpool.py index fb856e175..963df83aa 100644 --- a/tests/attention/test_dsa_kpool.py +++ b/tests/attention/test_dsa_kpool.py @@ -76,6 +76,22 @@ def harness(monkeypatch): return backend, pool, ape +def test_fp8_latent_cache_keeps_kpool_index_tiers_bf16(): + from freetoken.kvcache.dsa_pool import KpoolDSAKVCache + + pool = KpoolDSAKVCache( + latent_dim=LATENT, num_layers=1, num_pages=8, page_size=64, + dtype=torch.bfloat16, device=torch.device(DEV), + index_head_dim=DI, num_index_layers=1, + index_ratio=KPOOL, num_req_slots=4, kv_quant="fp8", + ) + assert pool.latent_rows(0).dtype is torch.uint8 + assert pool.latent_scale(0).dtype is torch.float32 + assert pool.index_k_cache(0).dtype is torch.bfloat16 + assert pool.tail_k(0).dtype is torch.bfloat16 + assert pool.tail_gate(0).dtype is torch.bfloat16 + + def _req(device_len, cached_len=0): return SimpleNamespace( table_idx=0, device_len=device_len, extend_len=device_len - cached_len, diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py index 9b870273c..2c8e5545c 100644 --- a/tests/engine/test_kv_quant_config.py +++ b/tests/engine/test_kv_quant_config.py @@ -107,12 +107,10 @@ def test_only_the_backends_that_read_scales_declare_fp8_support(): 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"} + # triton serves plain paged / hybrid-SWA pools, qsa_sparse dequantizes selected + # rows, and dsa dequantizes MLA latent rows. External backends and sparse + # families without a scale path must stay out of this set. + assert fp8 == {"triton", "qsa_sparse", "dsa"} def test_auto_avoids_the_fast_backends_for_fp8(monkeypatch): @@ -152,7 +150,17 @@ def test_explicit_triton_is_accepted(monkeypatch): assert config.kv_quant == "fp8" -@pytest.mark.parametrize("kind", ["mla", "dsa", "dsv4", "bsa"]) +@pytest.mark.parametrize("kind", ["mla", "dsa"]) +def test_mla_and_dsa_select_the_scale_reading_backend(monkeypatch, kind): + from freetoken.engine.engine import _adjust_config + + _patch_fast_machine(monkeypatch) + config = _config(kind, attention_backend="auto", kv_quant="fp8") + _adjust_config(config) + assert config.attention_backend == "dsa" + + +@pytest.mark.parametrize("kind", ["dsv4", "bsa"]) def test_pool_families_without_a_scale_read_path_are_rejected(monkeypatch, kind): from freetoken.engine.engine import _adjust_config diff --git a/tests/kvcache/test_dsa_pool.py b/tests/kvcache/test_dsa_pool.py index 7f4b1c709..28b130497 100644 --- a/tests/kvcache/test_dsa_pool.py +++ b/tests/kvcache/test_dsa_pool.py @@ -16,13 +16,14 @@ LATENT, IDX_DIM = 80, 32 # latent 64+16: kernel spans need pow-2 dims (512+64 in the real model) -def _pool(num_pages: int): +def _pool(num_pages: int, kv_quant: str = "none"): from freetoken.kvcache.dsa_pool import DSAKVCache return DSAKVCache( latent_dim=LATENT, num_layers=2, num_pages=num_pages, page_size=1, dtype=torch.bfloat16, device=torch.device("cuda"), index_head_dim=IDX_DIM, num_index_layers=1, + kv_quant=kv_quant, ) @@ -91,6 +92,39 @@ def test_sparse_decode_reads_grown_pool_through_backend_kernels(): assert (o[0, 0].float() - ref_o).abs().max().item() < 2e-2 +def test_fp8_latent_rows_decode_with_their_scales(): + """The DSA kernel must read codes and scales from matching physical rows.""" + from freetoken.kernel.triton.glm_dsa_sparse import glm_dsa_sparse_attn + from freetoken.kernel.triton.kv_quant import codes_to_f32 + + torch.manual_seed(7) + pool = _pool(128, kv_quant="fp8") + rows = torch.randperm(128, device="cuda")[:96].to(torch.int32) + c_kv = torch.randn(96, LATENT - 16, device="cuda", dtype=torch.bfloat16) + k_rope = torch.randn(96, 16, device="cuda", dtype=torch.bfloat16) + pool.store_kv(c_kv, k_rope, rows, layer_id=0) + assert pool.latent_rows(0).dtype is torch.uint8 + assert pool.latent_scale(0).dtype is torch.float32 + + q = torch.randn(1, 1, 4, LATENT, device="cuda", dtype=torch.bfloat16) + sel = rows[:64].view(1, 1, -1) + cnt = torch.tensor([[64]], device="cuda", dtype=torch.int32) + out = glm_dsa_sparse_attn( + q, pool.latent_rows(0), sel, 0.1, counts=cnt, d_v=LATENT - 16, + pool_scale=pool.latent_scale(0), + ) + out_split = glm_dsa_sparse_attn( + q, pool.latent_rows(0), sel, 0.1, counts=cnt, d_v=LATENT - 16, + pool_scale=pool.latent_scale(0), force_splits=4, + ) + decoded = codes_to_f32(pool.latent_rows(0)) * pool.latent_scale(0).unsqueeze(-1) + picked = decoded[sel.view(-1).long()] + logits = (q[0, 0].float() @ picked.T) * 0.1 + ref = logits.softmax(-1) @ picked[:, : LATENT - 16] + assert (out[0, 0].float() - ref).abs().max().item() < 2e-2 + assert (out_split.float() - ref).abs().max().item() < 2e-2 + + def test_mla_pool_selected_by_group_spec(): """The factory keys MLA/DSA pools off the attention-group spec, never the model payload -- and zeroed index dims (the dense ablation) fall back to MLAKVCache.""" @@ -120,6 +154,10 @@ def cfg(index_dim, n_idx): dsa = create_kvcache_pool(model_config=cfg(IDX_DIM, 1), num_pages=8, page_size=1, device=torch.device("cuda"), dtype=torch.bfloat16) assert isinstance(dsa, DSAKVCache) + fp8_dsa = create_kvcache_pool(model_config=cfg(IDX_DIM, 1), num_pages=8, page_size=1, + device=torch.device("cuda"), dtype=torch.bfloat16, + kv_quant="fp8") + assert isinstance(fp8_dsa, DSAKVCache) and fp8_dsa.store_dtype is torch.uint8 mla = create_kvcache_pool(model_config=cfg(0, 0), num_pages=8, page_size=1, device=torch.device("cuda"), dtype=torch.bfloat16) assert isinstance(mla, MLAKVCache) and not isinstance(mla, DSAKVCache) diff --git a/tests/kvcache/test_kv_cache_rebuild.py b/tests/kvcache/test_kv_cache_rebuild.py index a5ec94a08..ec5259277 100644 --- a/tests/kvcache/test_kv_cache_rebuild.py +++ b/tests/kvcache/test_kv_cache_rebuild.py @@ -64,6 +64,13 @@ def test_mla_and_dsa_rebuild_from_config_and_unit_bytes(): # the index slab's per-token bytes ride on top of the latent slab's, each floored on its own assert dsa.unit_bytes() == (layers * latent * 2 + n_idx * idx_dim * 2, 0) + fp8 = MLAKVCache(latent_dim=latent, num_layers=layers, num_pages=8, page_size=1, + dtype=torch.bfloat16, device=torch.device("cpu"), kv_quant="fp8") + fp8.rebuild_from_config(config=None, num_pages=20) + assert fp8.latent_rows(0).dtype is torch.uint8 + # One byte per latent element plus one FP32 scale for each latent layer. + assert fp8.unit_bytes() == (layers * latent + layers * 4, 0) + def _hybrid_groups(): from freetoken.models.config import KVCacheGroupSpec diff --git a/tests/kvcache/test_mha_pool_fp8.py b/tests/kvcache/test_mha_pool_fp8.py index dd3f0542d..8d74b0edd 100644 --- a/tests/kvcache/test_mha_pool_fp8.py +++ b/tests/kvcache/test_mha_pool_fp8.py @@ -182,10 +182,8 @@ def test_fp8_lands_just_above_half_the_bytes(): 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.""" +def test_latent_kv_pool_accepts_fp8(): + """MLA carries one latent row per token, so it uses one FP32 scale per row.""" from freetoken.attention import AttnType from freetoken.kvcache import create_kvcache_pool @@ -206,16 +204,15 @@ def test_latent_kv_pool_rejects_fp8(): 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). + fp8 = create_kvcache_pool( + model_config=mc, + num_pages=4, + page_size=1, + dtype=torch.bfloat16, + device=DEV, + kv_quant="fp8", + ) + assert fp8.kv_quant == "fp8" and fp8.store_dtype is torch.uint8 pool = create_kvcache_pool( model_config=mc, num_pages=4, From 04d462149e211dfb503540d58d92907fd331a5c6 Mon Sep 17 00:00:00 2001 From: ArqAlice Date: Tue, 8 Sep 2026 22:02:03 +0900 Subject: [PATCH 12/12] feat(kvcache): support nvfp4 latent dsa kv --- docs/cli.md | 15 ++- python/freetoken/attention/__init__.py | 1 + python/freetoken/attention/dsa.py | 2 + python/freetoken/engine/engine.py | 6 +- .../freetoken/kernel/triton/glm_dsa_sparse.py | 103 ++++++++++++++---- python/freetoken/kernel/triton/kv_nvfp4.py | 34 +++++- python/freetoken/kvcache/__init__.py | 5 +- python/freetoken/kvcache/dsa_pool.py | 60 ++++++++-- tests/attention/test_dsa_kpool.py | 75 ++++++++++++- tests/engine/test_kv_quant_config.py | 25 ++--- tests/kernels/test_kv_nvfp4.py | 88 +++++++++++++++ tests/kvcache/test_dsa_pool.py | 54 +++++++++ tests/models/test_glm5_next_config.py | 22 ++++ tests/models/test_glm_dsa.py | 25 +++-- 14 files changed, 442 insertions(+), 73 deletions(-) diff --git a/docs/cli.md b/docs/cli.md index 552499a5b..4fa6cbc3e 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -203,9 +203,13 @@ profile that `ft serve --moe-backend auto` and `--moe-hybrid-max-fetch -1` then `ft serve --model --kv-cache-dtype nvfp4 --attention-backend triton` opts into packed E2M1 KV storage. The initial implementation supports plain paged FULL attention (MHA/GQA), hybrid-SWA, and the full-attention portion of hybrid-linear -models, and QSA. Head dimensions must be divisible by 16. MLA/DSA, DSV4, and BSA -pools are rejected at startup. `auto` selects Triton or QSA sparse attention for -supported models. +models, QSA, and MLA/DSA (including GLM-5.3-Flash). Head dimensions must be +divisible by 16. DSV4 and BSA pools are rejected at startup. `auto` selects +Triton, QSA sparse, or DSA attention for supported models. For MLA/DSA use +`--attention-backend auto` or `--attention-backend dsa`. Only the latent slab is +quantized; indexer keys, kpool tails/gates, and recurrent states retain their +existing precision. A 512-element latent row occupies 292 bytes instead of 1024 +bytes in BF16, excluding those other tiers. Each K or V row stores `head_dim / 2` packed bytes, `head_dim / 16` E4M3 block-scale bytes, and one FP32 row scale. At head_dim 128 this is 76 bytes, versus 256 for @@ -216,8 +220,9 @@ The second-level scale is dynamic per token/head, so appending a token never rescales an existing prefix. This is a FreeToken KV layout, not an external NVFP4 checkpoint or attention-library ABI. K/V are restored inside attention; Q and attention arithmetic retain their compute precision. The MoE weight option -`--nvfp4-backend` is independent. Prefill uses fresh compute-dtype K/V while -cached prefixes are restored, as in the FP8 path. +`--nvfp4-backend` is independent. Paged MHA prefill uses fresh compute-dtype K/V +while cached prefixes are restored, as in the FP8 path. MLA/DSA stores fresh +latent rows first and reads the quantized cache in both prefill and decode. NVFP4 is opt-in: assess quality on your checkpoint and workload before using it for long-context inference. Capacity savings do not guarantee faster decode; diff --git a/python/freetoken/attention/__init__.py b/python/freetoken/attention/__init__.py index 1483b1faf..dbcc1f2fa 100644 --- a/python/freetoken/attention/__init__.py +++ b/python/freetoken/attention/__init__.py @@ -115,6 +115,7 @@ def create_dsv4_sparse_backend(config: ModelConfig): BackendInfo( supported_types=frozenset({AttnType.MLA, AttnType.DSA}), supports_fp8_kv=True, + supports_nvfp4_kv=True, ), ) def create_dsa_backend(config: ModelConfig): diff --git a/python/freetoken/attention/dsa.py b/python/freetoken/attention/dsa.py index 9aedb3ed7..2d6781da5 100644 --- a/python/freetoken/attention/dsa.py +++ b/python/freetoken/attention/dsa.py @@ -202,6 +202,8 @@ def _attend( q_cat, self.kvcache.latent_rows(layer_id), sel, self.sm_scale, counts=cnt, d_v=self.kv_lora_rank, pool_scale=self.kvcache.latent_scale(layer_id), + kv_quant=self.kvcache.kv_quant, + pool_block_scale=self.kvcache.latent_block_scale(layer_id), ) def mla_forward( diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 8fb4e3db4..6ce170289 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1362,15 +1362,15 @@ def override(attr: str, value: Any): # this is dangerous, use with caution kv_quant = _resolve_kv_quant(getattr(config, "kv_quant", "none")) override("kv_quant", kv_quant) if kv_quant == "nvfp4": - if required_attn_types - {AttnType.FULL, AttnType.SWA, AttnType.QSA}: + if required_attn_types - {AttnType.FULL, AttnType.SWA, AttnType.QSA, AttnType.MLA, AttnType.DSA}: raise ValueError( - "--kv-cache-dtype nvfp4 requires a plain paged FULL, hybrid-SWA, or QSA KV pool" + "--kv-cache-dtype nvfp4 requires a paged FULL, hybrid-SWA, QSA, or MLA/DSA KV pool" ) for spec in model_config.kv_cache_group_specs(): if spec.head_dim % 16: raise ValueError("--kv-cache-dtype nvfp4 requires head_dim divisible by 16") if kv_quant != "none": - # fp8 codes are wired through the pools that hand their rows to a Triton + # Quantized codes are wired through the pools that hand their rows to a Triton # kernel: plain paged and hybrid-SWA, QSA sparse, and DSA/MLA. QSA's index # tier and DSA's index-key/tail tiers stay bf16; the DSA kernel dequantizes # selected latent rows with their per-token scale. Other sparse families have diff --git a/python/freetoken/kernel/triton/glm_dsa_sparse.py b/python/freetoken/kernel/triton/glm_dsa_sparse.py index 4415a68e9..d73ece8cf 100644 --- a/python/freetoken/kernel/triton/glm_dsa_sparse.py +++ b/python/freetoken/kernel/triton/glm_dsa_sparse.py @@ -29,6 +29,7 @@ import triton.language as tl from freetoken.kernel.triton.e4m3_compat import kv_load_e4m3_tile_f32 +from freetoken.kernel.triton.kv_nvfp4 import load_nvfp4 BLOCK_H = 16 BLOCK_T = 32 @@ -38,7 +39,7 @@ @triton.jit def _glm_dsa_sparse_kernel( - q_ptr, pool_ptr, pool_scale_ptr, o_ptr, idx_ptr, cnt_ptr, + q_ptr, pool_ptr, pool_scale_ptr, pool_block_ptr, o_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, stride_qb, stride_qm, stride_qh, stride_qd, @@ -54,6 +55,7 @@ def _glm_dsa_sparse_kernel( HAS_COUNTS: tl.constexpr, HAS_ROPE: tl.constexpr, HAS_FP8: tl.constexpr, + HAS_NVFP4: tl.constexpr, ): pid_m = tl.program_id(0) pid_b = tl.program_id(1) @@ -65,11 +67,15 @@ def _glm_dsa_sparse_kernel( q_base = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh q_v = tl.load(q_base + offs_v[None, :] * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_v = q_v.to(q_ptr.dtype.element_ty) if HAS_ROPE: # NoPE checkpoints (glm5_next) have D_R == 0: tl.arange needs a non-empty # span, so the whole rope half is compiled out on the constexpr. offs_r = tl.arange(0, D_R) q_r = tl.load(q_base + (D_V + offs_r[None, :]) * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_r = q_r.to(q_ptr.dtype.element_ty) m_i = tl.full((BLOCK_H,), -float("inf"), dtype=tl.float32) l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) @@ -86,7 +92,13 @@ def _glm_dsa_sparse_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 kv_base = pool_ptr + idxs[:, None] * stride_pn - if HAS_FP8: + if HAS_NVFP4: + kv_v = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_v[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: row_scale = tl.load(pool_scale_ptr + idxs * stride_ps, mask=valid, other=0.0) kv_v = kv_load_e4m3_tile_f32( kv_base + offs_v[None, :] * stride_pd, valid[:, None] @@ -96,7 +108,13 @@ def _glm_dsa_sparse_kernel( scores = tl.dot(q_v, tl.trans(kv_v)) if HAS_ROPE: - if HAS_FP8: + if HAS_NVFP4: + kv_r = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_r[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, DIM_OFFSET=D_V, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: kv_r = kv_load_e4m3_tile_f32( kv_base + (D_V + offs_r[None, :]) * stride_pd, valid[:, None] ) * row_scale[:, None] @@ -114,7 +132,7 @@ def _glm_dsa_sparse_kernel( acc = acc * alpha[:, None] + tl.dot(p.to(kv_v.dtype), kv_v) m_i = m_new - o = acc / l_i[:, None] + o = tl.where(l_i[:, None] > 0, acc / l_i[:, None], 0.0) o_ptrs = o_ptr + pid_b * stride_ob + pid_m * stride_om + offs_h[:, None] * stride_oh + offs_v[None, :] * stride_od tl.store(o_ptrs, o.to(o_ptr.dtype.element_ty), mask=h_mask[:, None]) @@ -215,7 +233,7 @@ def glm_dsa_decode_logits( @triton.jit def _glm_dsa_splitk_kernel( - q_ptr, pool_ptr, pool_scale_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, + q_ptr, pool_ptr, pool_scale_ptr, pool_block_ptr, mid_o_ptr, mid_lse_ptr, idx_ptr, cnt_ptr, scale, H, TOPK, stride_qb, stride_qm, stride_qh, stride_qd, @@ -232,6 +250,7 @@ def _glm_dsa_splitk_kernel( HAS_COUNTS: tl.constexpr, HAS_ROPE: tl.constexpr, HAS_FP8: tl.constexpr, + HAS_NVFP4: tl.constexpr, NUM_SPLITS: tl.constexpr, ): """Stage 1 (decode flash-decoding): each program reduces one BLOCK_T-aligned slice of @@ -262,9 +281,13 @@ def _glm_dsa_splitk_kernel( if split_end > split_start: q_base = q_ptr + pid_b * stride_qb + pid_m * stride_qm + offs_h[:, None] * stride_qh q_v = tl.load(q_base + offs_v[None, :] * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_v = q_v.to(q_ptr.dtype.element_ty) if HAS_ROPE: offs_r = tl.arange(0, D_R) q_r = tl.load(q_base + (D_V + offs_r[None, :]) * stride_qd, mask=h_mask[:, None], other=0.0).to(tl.float32) + if HAS_NVFP4: + q_r = q_r.to(q_ptr.dtype.element_ty) idx_base = idx_ptr + pid_b * stride_ib + pid_m * stride_im for start in range(split_start, split_end, BLOCK_T): @@ -273,7 +296,13 @@ def _glm_dsa_splitk_kernel( idxs = tl.load(idx_base + offs_t * stride_it, mask=t_mask, other=-1) valid = idxs >= 0 kv_base = pool_ptr + idxs[:, None] * stride_pn - if HAS_FP8: + if HAS_NVFP4: + kv_v = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_v[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: row_scale = tl.load(pool_scale_ptr + idxs * stride_ps, mask=valid, other=0.0) kv_v = kv_load_e4m3_tile_f32( kv_base + offs_v[None, :] * stride_pd, valid[:, None] @@ -283,7 +312,13 @@ def _glm_dsa_splitk_kernel( scores = tl.dot(q_v, tl.trans(kv_v)) if HAS_ROPE: - if HAS_FP8: + if HAS_NVFP4: + kv_r = load_nvfp4( + pool_ptr, pool_block_ptr, pool_scale_ptr, + idxs[:, None], 0, offs_r[None, :], valid[:, None], + stride_pn, 0, stride_ps, D_V + D_R, DIM_OFFSET=D_V, + ).to(q_ptr.dtype.element_ty) + elif HAS_FP8: kv_r = kv_load_e4m3_tile_f32( kv_base + (D_V + offs_r[None, :]) * stride_pd, valid[:, None] ) * row_scale[:, None] @@ -351,7 +386,7 @@ def _glm_dsa_merge_kernel( l_i = l_i * alpha + beta m_i = m_new - o = acc / l_i + o = tl.where(l_i > 0, acc / l_i, 0.0) o_ptrs = ( o_ptr + pid_b * stride_ob + pid_m * stride_om + pid_h * stride_oh + offs_v * stride_od @@ -381,6 +416,8 @@ def glm_dsa_sparse_attn( d_v: int = 512, pool_scale: torch.Tensor | None = None, # [rows] fp32, one scale per quantized latent row force_splits: int | None = None, # tests only: 0 = single-program, N = split-k N + kv_quant: str | None = None, + pool_block_scale: torch.Tensor | None = None, ) -> torch.Tensor: """Sparse MLA attention over gathered latent rows; returns ``[b, m, h, d_v]``. @@ -388,21 +425,45 @@ def glm_dsa_sparse_attn( list across all queries with STRIDE 0 -- the identity-selection dense path: every query reads the same position-ordered row list, causally bounded by its own ``counts[q] = position + 1``, with zero per-query index materialization. + + NVFP4 restores tiles to Q's compute dtype before the dot products, with FP32 + accumulators. FP32 dot operands exceed consumer GPU shared memory at width 576. """ b, m, h, d = q.shape d_r = d - d_v topk = topk_idxs.shape[-1] - assert pool.shape[-1] == d, (pool.shape, d) - has_fp8 = pool_scale is not None - if has_fp8: + if kv_quant is None: + kv_quant = "fp8" if pool_scale is not None else "none" + assert kv_quant in ("none", "fp8", "nvfp4"), kv_quant + has_fp8 = kv_quant == "fp8" + has_nvfp4 = kv_quant == "nvfp4" + stored_dim = d // 2 if has_nvfp4 else d + assert pool.ndim == 2 and pool.shape[-1] == stored_dim, (pool.shape, d) + if has_fp8 or has_nvfp4: + assert pool_scale is not None assert pool.dtype == torch.uint8, pool.dtype assert pool_scale.shape == (pool.shape[0],), (pool_scale.shape, pool.shape) assert pool_scale.dtype is torch.float32, pool_scale.dtype scale_pool = pool_scale.contiguous() else: + assert pool_scale is None + assert pool.is_floating_point() scale_pool = pool + if has_nvfp4: + assert d % 16 == 0 + assert pool_block_scale is not None + assert pool_block_scale.shape == (pool.shape[0], d // 16) + assert pool_block_scale.dtype == torch.uint8 + assert pool_block_scale.device == pool.device + block_pool = pool_block_scale.contiguous() + else: + assert pool_block_scale is None + block_pool = pool + assert pool.device == q.device + if has_fp8 or has_nvfp4: + assert pool_scale.device == pool.device q = q.contiguous() - pool_2d = pool.reshape(-1, d) + pool_2d = pool.reshape(-1, stored_dim) assert pool_2d.stride(-1) == 1 idx = topk_idxs.contiguous().to(torch.int32) broadcast_m = idx.shape[1] == 1 and m > 1 @@ -416,25 +477,27 @@ def glm_dsa_sparse_attn( else: cnt, stride_nb, stride_nm = idx, 0, 0 + # Packed gathers need additional layout conversions at the 512-wide latent size. + block_t = 16 if has_nvfp4 else BLOCK_T n_splits = _split_count(b, m, h, topk, q.device) if force_splits is None else force_splits if n_splits: mid_o = q.new_empty(b, m, h, n_splits, d_v, dtype=torch.float32) mid_lse = q.new_empty(b, m, h, n_splits, dtype=torch.float32) grid1 = (m * n_splits, b, triton.cdiv(h, BLOCK_H)) _glm_dsa_splitk_kernel[grid1]( - q, pool_2d, scale_pool, mid_o, mid_lse, idx, cnt, + q, pool_2d, scale_pool, block_pool, mid_o, mid_lse, idx, cnt, float(softmax_scale), h, topk, q.stride(0), q.stride(1), q.stride(2), q.stride(3), pool_2d.stride(0), pool_2d.stride(1), - scale_pool.stride(0) if has_fp8 else 0, + scale_pool.stride(0) if has_fp8 or has_nvfp4 else 0, mid_o.stride(0), mid_o.stride(1), mid_o.stride(2), mid_o.stride(3), mid_o.stride(4), mid_lse.stride(0), mid_lse.stride(1), mid_lse.stride(2), mid_lse.stride(3), idx.stride(0), 0 if broadcast_m else idx.stride(1), idx.stride(2), stride_nb, stride_nm, D_V=d_v, D_R=d_r, - BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, - HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, NUM_SPLITS=n_splits, + BLOCK_H=BLOCK_H, BLOCK_T=block_t, + HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, NUM_SPLITS=n_splits, # The 512-wide latent accumulator exceeds the 99 KiB shared-memory # limit of consumer Blackwell GPUs with a two-stage pipeline. num_warps=4, num_stages=1, @@ -452,18 +515,18 @@ def glm_dsa_sparse_attn( grid = (m, b, triton.cdiv(h, BLOCK_H)) _glm_dsa_sparse_kernel[grid]( - q, pool_2d, scale_pool, o, idx, cnt, + q, pool_2d, scale_pool, block_pool, o, idx, cnt, float(softmax_scale), h, topk, q.stride(0), q.stride(1), q.stride(2), q.stride(3), pool_2d.stride(0), pool_2d.stride(1), - scale_pool.stride(0) if has_fp8 else 0, + scale_pool.stride(0) if has_fp8 or has_nvfp4 else 0, o.stride(0), o.stride(1), o.stride(2), o.stride(3), idx.stride(0), 0 if broadcast_m else idx.stride(1), idx.stride(2), stride_nb, stride_nm, D_V=d_v, D_R=d_r, - BLOCK_H=BLOCK_H, BLOCK_T=BLOCK_T, - HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, + BLOCK_H=BLOCK_H, BLOCK_T=block_t, + HAS_COUNTS=has_counts, HAS_ROPE=d_r > 0, HAS_FP8=has_fp8, HAS_NVFP4=has_nvfp4, num_warps=4, num_stages=1, ) return o diff --git a/python/freetoken/kernel/triton/kv_nvfp4.py b/python/freetoken/kernel/triton/kv_nvfp4.py index 91bfc602b..04727deec 100644 --- a/python/freetoken/kernel/triton/kv_nvfp4.py +++ b/python/freetoken/kernel/triton/kv_nvfp4.py @@ -35,13 +35,14 @@ def _decode_e2m1(code): @triton.jit def load_nvfp4(ptr, block_ptr, row_ptr, slots, head, dims, slot_mask, - stride_slot, stride_head, stride_row, D: tl.constexpr): + stride_slot, stride_head, stride_row, D: tl.constexpr, + DIM_OFFSET: tl.constexpr = 0): TRANSPOSE: tl.constexpr = dims.shape[0] != 1 WIDTH: tl.constexpr = dims.shape[0] if TRANSPOSE else dims.shape[1] TOKENS: tl.constexpr = slots.shape[0] * slots.shape[1] slot = slots.reshape(TOKENS).to(tl.int64) valid = slot_mask.reshape(TOKENS) - dim = tl.arange(0, WIDTH) + dim = DIM_OFFSET + tl.arange(0, WIDTH) packed = tl.load( ptr + slot[:, None] * stride_slot + head * stride_head + (dim[None, :] // 2), valid[:, None] & (dim[None, :] < D), other=0, @@ -82,6 +83,35 @@ def _quantize_row(src, dst, block_ptr, row_ptr, t, h, slot, stride_src, tl.store(row_ptr + slot * HEADS + h, row_scale) +@triton.jit +def _scatter_rows(src, dst, block, row, indices, stride_src, + D: tl.constexpr, BLOCKS: tl.constexpr): + t = tl.program_id(0) + slot = tl.load(indices + t).to(tl.int64) + _quantize_row(src, dst, block, row, t, 0, slot, stride_src, 1, D, BLOCKS) + + +def quantize_nvfp4_rows_to_cache(rows, out_loc, cache, scales, block_scales) -> None: + """Quantize a single MLA latent slab, with one second-level scale per token.""" + tokens, dim = rows.shape + slots = cache.shape[0] + assert dim % 16 == 0 and rows.stride(1) == 1 + assert rows.dtype in (torch.float16, torch.bfloat16, torch.float32) + assert cache.shape == (slots, dim // 2) and cache.dtype == torch.uint8 + assert scales.shape == (slots,) and scales.dtype == torch.float32 + assert block_scales.shape == (slots, dim // 16) and block_scales.dtype == torch.uint8 + assert out_loc.shape == (tokens,) and out_loc.dtype in (torch.int32, torch.int64) + for tensor in (cache, scales, block_scales, out_loc): + assert tensor.is_contiguous() and tensor.device == rows.device + assert rows.is_cuda + if tokens: + _scatter_rows[(tokens,)]( + rows, cache, block_scales, scales, out_loc, rows.stride(0), + D=dim, BLOCKS=triton.next_power_of_2(dim // 16), + num_warps=4, enable_fp_fusion=False, + ) + + @triton.jit def _scatter(k, v, kc, vc, kb, vb, kr, vr, indices, stride_k, stride_v, HEADS: tl.constexpr, D: tl.constexpr, BLOCKS: tl.constexpr): diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index 555de38d6..24d2ec790 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -148,13 +148,12 @@ def create_kvcache_pool( from freetoken.attention import AttnType if any( - spec.attn_type not in (AttnType.FULL, AttnType.SWA, AttnType.QSA) - or spec.mla + spec.attn_type not in (AttnType.FULL, AttnType.SWA, AttnType.QSA, AttnType.MLA, AttnType.DSA) or spec.head_dim % 16 for spec in model_config.kv_cache_group_specs() ): raise ValueError( - "--kv-cache-dtype nvfp4 requires paged FULL, hybrid-SWA, or QSA groups " + "--kv-cache-dtype nvfp4 requires paged FULL, hybrid-SWA, QSA, or MLA/DSA groups " "with head_dim divisible by 16" ) if model_config.has_swa_attention: diff --git a/python/freetoken/kvcache/dsa_pool.py b/python/freetoken/kvcache/dsa_pool.py index 311b27fb4..6d9e2d4f3 100644 --- a/python/freetoken/kvcache/dsa_pool.py +++ b/python/freetoken/kvcache/dsa_pool.py @@ -46,6 +46,9 @@ def __init__( kv_quant: str = "none", ) -> None: self._latent_dim = latent_dim + if kv_quant == "nvfp4" and latent_dim % 16: + raise ValueError("NVFP4 KV requires latent_dim divisible by 16") + self._stored_dim = latent_dim // 2 if kv_quant == "nvfp4" else latent_dim if layer_ids is None: self._num_layers = num_layers self._layer_index: dict[int, int] | None = None @@ -63,16 +66,20 @@ def _local_layer(self, layer_id: int) -> int: def _alloc(self, num_pages: int) -> None: self._num_pages = num_pages - shape = (1, self._num_layers, num_pages, self._page_size, 1, self._latent_dim) - if self.kv_quant == "fp8": - from freetoken.kernel.triton.kv_quant import alloc_codes - - self._kv_buffer = alloc_codes(shape, self._device) + shape = (1, self._num_layers, num_pages, self._page_size, 1, self._stored_dim) + self._block_scale_buffer = None + if self.kv_quant in ("fp8", "nvfp4"): + self._kv_buffer = torch.zeros(shape, device=self._device, dtype=torch.uint8) self._scale_buffer = torch.zeros( (self._num_layers, num_pages * self._page_size), device=self._device, dtype=torch.float32, ) + if self.kv_quant == "nvfp4": + self._block_scale_buffer = torch.zeros( + (self._num_layers, num_pages * self._page_size, self._latent_dim // 16), + device=self._device, dtype=torch.uint8, + ) elif self.kv_quant == "none": self._kv_buffer = torch.empty(shape, device=self._device, dtype=self._dtype) self._scale_buffer = None @@ -81,7 +88,7 @@ def _alloc(self, num_pages: int) -> None: # -- views (addressed by GLOBAL layer id; remapped when layer_ids was given) -- def k_cache(self, layer_id: int) -> torch.Tensor: - """Paged latent view ``[num_pages, page_size, latent_dim]``.""" + """Paged latent view; NVFP4 stores ``latent_dim // 2`` bytes per row.""" return self._kv_buffer[0, self._local_layer(layer_id)].view( self._num_pages, self._page_size, -1 ) @@ -91,8 +98,8 @@ def v_cache(self, layer_id: int) -> torch.Tensor: return self.k_cache(layer_id) def latent_rows(self, layer_id: int) -> torch.Tensor: - """Row-flat latent view ``[num_pages * page_size, latent_dim]``.""" - return self._kv_buffer[0, self._local_layer(layer_id)].view(-1, self._latent_dim) + """Row-flat latent view ``[num_pages * page_size, stored_dim]``.""" + return self._kv_buffer[0, self._local_layer(layer_id)].view(-1, self._stored_dim) def latent_scale(self, layer_id: int) -> torch.Tensor | None: """One FP32 scale per latent row, or ``None`` for the compute-dtype pool.""" @@ -100,6 +107,12 @@ def latent_scale(self, layer_id: int) -> torch.Tensor | None: return None return self._scale_buffer[self._local_layer(layer_id)] + def latent_block_scale(self, layer_id: int) -> torch.Tensor | None: + """E4M3 bytes per 16 latent elements, present only for NVFP4.""" + if self._block_scale_buffer is None: + return None + return self._block_scale_buffer[self._local_layer(layer_id)] + # -- writes ----------------------------------------------------------------- def store_kv( self, @@ -115,7 +128,18 @@ def store_kv( store.cu (two-width store). """ rows = self.latent_rows(layer_id) - split = rows.shape[1] - k_rope.shape[-1] + split = self._latent_dim - k_rope.shape[-1] + assert c_kv.shape == (out_loc.numel(), split) + assert k_rope.shape[0] == c_kv.shape[0] + if self.kv_quant == "nvfp4": + from freetoken.kernel.triton.kv_nvfp4 import quantize_nvfp4_rows_to_cache + + latent = torch.cat((c_kv, k_rope), dim=-1) if k_rope.shape[-1] else c_kv + quantize_nvfp4_rows_to_cache( + latent, out_loc, rows, self.latent_scale(layer_id), + self.latent_block_scale(layer_id), + ) + return if self.kv_quant == "fp8": from freetoken.kernel.triton.kv_quant import quantize_rows_to_cache @@ -130,6 +154,7 @@ def rebuild(self, num_pages: int) -> None: callers re-derive views per forward, same contract as MHAKVCache.rebuild).""" self._kv_buffer = None self._scale_buffer = None + self._block_scale_buffer = None if self._device.type == "cuda": torch.cuda.synchronize(self._device) torch.cuda.empty_cache() @@ -156,6 +181,8 @@ def unit_bytes(self) -> tuple[int, int]: kv = int(buf.numel() * buf.element_size()) // tokens if self._scale_buffer is not None: kv += int(self._scale_buffer.numel() * self._scale_buffer.element_size()) // tokens + if self._block_scale_buffer is not None: + kv += self._block_scale_buffer.numel() // tokens return kv, 0 # -- pool properties ---------------------------------------------------------- @@ -265,6 +292,21 @@ def _index_rows(self, num_pages: int) -> int: # 1/ratio shadow of every token slot + one scratch row per request slot. return num_pages * self._page_size // self._index_ratio + self._num_req_slots + @classmethod + def kv_cost(cls, config) -> tuple[int, int, int, int]: + per_page, fixed, page_size, reserve = super().kv_cost(config) + for spec in config.model_config.kv_cache_group_specs(): + if spec.mla and spec.index_ratio > 1: + row_bytes = spec.index_head_dim * spec.num_index_layers * 2 + fixed += (config.max_running_req + 1) * row_bytes * (2 * spec.index_ratio + 1) + return per_page, fixed, page_size, reserve + + def unit_bytes(self) -> tuple[int, int]: + # Scratch rows and the two tail rings are fixed costs, not token capacity. + kv, swa = MLAKVCache.unit_bytes(self) + index_bytes = self._num_index_layers * self._index_head_dim * 2 // self._index_ratio + return kv + index_bytes, swa + @property def cmp_scratch_base(self) -> int: """First scratch row (== shadow row count); request ``table_idx`` offsets it.""" diff --git a/tests/attention/test_dsa_kpool.py b/tests/attention/test_dsa_kpool.py index 963df83aa..bed1bbdae 100644 --- a/tests/attention/test_dsa_kpool.py +++ b/tests/attention/test_dsa_kpool.py @@ -49,8 +49,8 @@ def _args(num_layers=1): ) -@pytest.fixture() -def harness(monkeypatch): +@pytest.fixture(params=["none", "nvfp4"]) +def harness(monkeypatch, request): from freetoken.attention.dsa_indexer_kpool import Glm5NextDSABackend from freetoken.kvcache.dsa_pool import KpoolDSAKVCache @@ -58,7 +58,7 @@ def harness(monkeypatch): latent_dim=LATENT, num_layers=1, num_pages=8, page_size=64, dtype=torch.bfloat16, device=torch.device(DEV), index_head_dim=DI, num_index_layers=1, - index_ratio=KPOOL, num_req_slots=4, + index_ratio=KPOOL, num_req_slots=4, kv_quant=request.param, ) page_table = torch.full((4, 512), -1, dtype=torch.int32, device=DEV) page_table[0, :512] = torch.arange(512, dtype=torch.int32, device=DEV) @@ -76,14 +76,15 @@ def harness(monkeypatch): return backend, pool, ape -def test_fp8_latent_cache_keeps_kpool_index_tiers_bf16(): +@pytest.mark.parametrize("kv_quant", ["fp8", "nvfp4"]) +def test_quantized_latent_cache_keeps_kpool_index_tiers_bf16(kv_quant): from freetoken.kvcache.dsa_pool import KpoolDSAKVCache pool = KpoolDSAKVCache( latent_dim=LATENT, num_layers=1, num_pages=8, page_size=64, dtype=torch.bfloat16, device=torch.device(DEV), index_head_dim=DI, num_index_layers=1, - index_ratio=KPOOL, num_req_slots=4, kv_quant="fp8", + index_ratio=KPOOL, num_req_slots=4, kv_quant=kv_quant, ) assert pool.latent_rows(0).dtype is torch.uint8 assert pool.latent_scale(0).dtype is torch.float32 @@ -129,10 +130,15 @@ def _rand_seq(total, seed=1): def _run(backend, batch, d, sl, ape): + d["kv_quant"] = backend.kvcache.kv_quant + backend.prepare_metadata(batch) + return _forward(backend, batch, d, sl, ape) + + +def _forward(backend, batch, d, sl, ape): from freetoken.attention.dsa import DSAIndexerInputs t = batch.positions.shape[0] - backend.prepare_metadata(batch) return backend.mla_forward( d["q_nope"][sl], d["q_nope"].new_empty(t, H, 0), d["c_kv"][sl], d["c_kv"].new_empty(t, 0), @@ -165,6 +171,14 @@ def _ref_scores(d, ape, q_idx_t, w_t, n_pools): def _ref_attend(d, q_t, positions): """Full softmax MLA over latent rows at ``positions`` for one query [H, LATENT].""" lat = d["c_kv"][positions].float() # [n, LATENT] + if d.get("kv_quant") == "nvfp4": + from tests.kernels.test_kv_nvfp4 import _reference + + packed, block, row = _reference(lat) + codes = torch.stack((packed & 15, packed >> 4), -1).flatten(-2).long() + grid = lat.new_tensor([0, .5, 1, 1.5, 2, 3, 4, 6, + 0, -.5, -1, -1.5, -2, -3, -4, -6]) + lat = grid[codes] * block.view(torch.float8_e4m3fn).float().repeat_interleave(16, -1) * row[:, None] logits = q_t.float() @ lat.T * SM_SCALE # [H, n] p = torch.softmax(logits, dim=-1) return (p @ lat).to(torch.bfloat16) @@ -383,3 +397,52 @@ def test_padding_and_empty_batch_leave_shadow_rows_clean(): ratio=KPOOL, ) assert torch.equal(slab, before) + + +def test_decode_graph_replay_tracks_slots_and_lengths(harness): + backend, pool, ape = harness + total, extra = 60, 6 + d = _rand_seq(total + extra, seed=81) + _run(backend, _prefill_batch(0, total), d, slice(0, total), ape) + static = {k: v[total:total + 1].clone() for k, v in d.items() if isinstance(v, torch.Tensor)} + batch = _decode_batch(total) + batch.size = batch.padded_size = 1 + backend.init_capture_graph(512, [1]) + backend.prepare_for_capture(batch) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + _forward(backend, batch, static, slice(None), ape) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out = _forward(backend, batch, static, slice(None), ape) + base = 0 + for pos in range(total, total + extra): + if pos == total + 2: + from freetoken.attention.dsa import get_global_ctx + + base = 256 + pool.latent_rows(0)[base:base + 128].copy_(pool.latent_rows(0)[:128]) + if pool.kv_quant == "nvfp4": + pool.latent_scale(0)[base:base + 128].copy_(pool.latent_scale(0)[:128]) + pool.latent_block_scale(0)[base:base + 128].copy_(pool.latent_block_scale(0)[:128]) + pool.index_k_cache(0)[64:96].copy_(pool.index_k_cache(0)[:32]) + pool.tail_k(0)[1].copy_(pool.tail_k(0)[0]) + pool.tail_gate(0)[1].copy_(pool.tail_gate(0)[0]) + get_global_ctx().page_table[1, :128] = torch.arange(base, base + 128, device=DEV) + batch.active_table_idx.fill_(1) + batch.padded_reqs[0].table_idx = 1 + for key, tensor in static.items(): + tensor.copy_(d[key][pos:pos + 1]) + batch.positions.fill_(pos) + batch.out_loc.fill_(base + pos) + batch.padded_reqs[0].device_len = pos + 1 + backend.prepare_metadata(batch) + backend.prepare_for_replay(batch) + graph.replay() + selected = _ref_selected_positions(d, ape, d["qi"][pos], d["wi"][pos], pos) + ref = _ref_attend(d, d["q_nope"][pos], selected) + torch.testing.assert_close(out[0], ref, atol=3e-2, rtol=1e-2) + backend.reset_capture() diff --git a/tests/engine/test_kv_quant_config.py b/tests/engine/test_kv_quant_config.py index 34a980aec..ba4196be0 100644 --- a/tests/engine/test_kv_quant_config.py +++ b/tests/engine/test_kv_quant_config.py @@ -49,6 +49,7 @@ def _model_config(kind): ), "mla": (_spec("full", AttnType.MLA, mla=True),), "dsa": (_spec("full", AttnType.DSA, mla=True, index_head_dim=128),), + "kpool": (_spec("full", AttnType.DSA, mla=True, index_head_dim=128, index_ratio=4),), "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),), @@ -57,7 +58,7 @@ def _model_config(kind): mc.has_swa_attention = True if kind == "dsv4": mc.dsv4_args = SimpleNamespace(window_size=128) - if kind == "qsa": + if kind in ("qsa", "kpool"): mc.has_linear_attention = True mc.kv_cache_group_specs = lambda: specs return mc @@ -118,7 +119,7 @@ def test_only_the_backends_that_read_scales_declare_fp8_support(): for name in SUPPORTED_ATTENTION_BACKENDS.supported_names() if attention_backend_info(name).supports_nvfp4_kv } - assert nvfp4 == {"triton", "qsa_sparse"} + assert nvfp4 == {"triton", "qsa_sparse", "dsa"} def test_auto_avoids_the_fast_backends_for_fp8(monkeypatch): @@ -167,7 +168,7 @@ def test_nvfp4_auto_selects_triton(monkeypatch): assert config.attention_backend == "triton" -@pytest.mark.parametrize("kind", ["mla", "dsa", "dsv4", "bsa"]) +@pytest.mark.parametrize("kind", ["dsv4", "bsa"]) def test_nvfp4_rejects_unsupported_pools_before_allocation(monkeypatch, kind): from freetoken.engine.engine import _adjust_config from freetoken.kvcache import create_kvcache_pool @@ -224,25 +225,17 @@ def test_nvfp4_accepts_qsa(monkeypatch): assert config.attention_backend == "qsa_sparse" -@pytest.mark.parametrize("kind", ["mla", "dsa"]) +@pytest.mark.parametrize("kind", ["mla", "dsa", "kpool"]) @pytest.mark.parametrize("backend", ["auto", "dsa"]) -def test_mla_and_dsa_select_the_scale_reading_backend(monkeypatch, kind, backend): +@pytest.mark.parametrize("kv_quant", ["fp8", "nvfp4"]) +def test_mla_and_dsa_select_the_scale_reading_backend(monkeypatch, kind, backend, kv_quant): from freetoken.engine.engine import _adjust_config _patch_fast_machine(monkeypatch) - config = _config(kind, attention_backend=backend, kv_quant="fp8") + config = _config(kind, attention_backend=backend, kv_quant=kv_quant) _adjust_config(config) assert config.attention_backend == "dsa" - - -@pytest.mark.parametrize("kind", ["mla", "dsa"]) -def test_explicit_dsa_backend_does_not_enable_nvfp4(monkeypatch, kind): - from freetoken.engine.engine import _adjust_config - - _patch_fast_machine(monkeypatch) - config = _config(kind, attention_backend="dsa", kv_quant="nvfp4") - with pytest.raises(ValueError, match="kv-cache-dtype nvfp4"): - _adjust_config(config) + assert config.page_size == (64 if kind == "kpool" else 1) @pytest.mark.parametrize("kind", ["dsv4", "bsa"]) diff --git a/tests/kernels/test_kv_nvfp4.py b/tests/kernels/test_kv_nvfp4.py index a91be7dab..662b44bf2 100644 --- a/tests/kernels/test_kv_nvfp4.py +++ b/tests/kernels/test_kv_nvfp4.py @@ -47,6 +47,94 @@ def _decode(pool, which, layer=1): return grid[code] * block.repeat_interleave(16, -1) * row[..., None] +def _decode_latent(pool, layer=0): + codes = pool.latent_rows(layer) + code = torch.stack((codes & 15, codes >> 4), -1).flatten(-2).long() + grid = torch.tensor([0, .5, 1, 1.5, 2, 3, 4, 6, + 0, -.5, -1, -1.5, -2, -3, -4, -6], device=codes.device) + block = pool.latent_block_scale(layer).view(torch.float8_e4m3fn).float() + return grid[code] * block.repeat_interleave(16, -1) * pool.latent_scale(layer)[:, None] + + +@pytest.mark.parametrize("rope", [0, 64]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_latent_scatter_matches_reference_and_preserves_prefix(rope, dtype): + from freetoken.kvcache.dsa_pool import MLAKVCache + + torch.manual_seed(73) + dim = 512 + rope + pool = MLAKVCache(dim, 4, 2, 64, dtype, torch.device("cuda"), + layer_ids=(1, 3), kv_quant="nvfp4") + x = torch.randn(7, dim + 32, device="cuda", dtype=dtype)[:, :dim] + x[0].zero_() + x[1, :16] *= 100 + x[2] *= 1e-5 + loc = torch.tensor([64, 3, 127, 14, 6, 90, 31], device="cuda") + pool.store_kv(x[:4, :512], x[:4, 512:], loc[:4], 3) + before = [v.clone() for v in (pool.latent_rows(3), pool.latent_scale(3), + pool.latent_block_scale(3))] + pool.store_kv(x[4:, :512], x[4:, 512:], loc[4:], 3) + pool.store_kv(x[:0, :512], x[:0, 512:], loc[:0], 3) + for got, saved in zip((pool.latent_rows(3), pool.latent_scale(3), + pool.latent_block_scale(3)), before): + torch.testing.assert_close(got[loc[:4]], saved[loc[:4]], rtol=0, atol=0) + packed, block, row = _reference(x) + torch.testing.assert_close(pool.latent_rows(3)[loc], packed) + torch.testing.assert_close(pool.latent_block_scale(3)[loc], block) + torch.testing.assert_close(pool.latent_scale(3)[loc], row) + assert torch.count_nonzero(pool.latent_rows(1)) == 0 + assert torch.isfinite(_decode_latent(pool, 3)).all() + assert pool.k_cache(3).data_ptr() == pool.v_cache(3).data_ptr() + # Reused physical slots replace codes and both scales together. + x[4:].mul_(0.125) + pool.store_kv(x[4:, :512], x[4:, 512:], loc[4:], 3) + packed, block, row = _reference(x[4:]) + torch.testing.assert_close(pool.latent_rows(3)[loc[4:]], packed) + torch.testing.assert_close(pool.latent_block_scale(3)[loc[4:]], block) + torch.testing.assert_close(pool.latent_scale(3)[loc[4:]], row) + + +@pytest.mark.parametrize("rope", [0, 64]) +@pytest.mark.parametrize("splits", [0, 4]) +@pytest.mark.parametrize("queries,broadcast", [(1, False), (5, False), (5, True)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +def test_sparse_mla_nvfp4_matches_restored_reference(rope, splits, queries, broadcast, dtype): + from freetoken.kernel.triton.glm_dsa_sparse import glm_dsa_sparse_attn + from freetoken.kvcache.dsa_pool import MLAKVCache + + torch.manual_seed(74) + dim, n = 512 + rope, 67 + pool = MLAKVCache(dim, 1, 2, 64, torch.bfloat16, torch.device("cuda"), kv_quant="nvfp4") + x = torch.randn(n, dim, device="cuda", dtype=torch.bfloat16) + x *= torch.linspace(.2, 2, n, device="cuda")[:, None] + if rope: + x[:, 512:] *= 3 + loc = torch.randperm(128, device="cuda")[:n] + pool.store_kv(x[:, :512], x[:, 512:], loc, 0) + q = torch.randn(2, queries, 19, dim, device="cuda", dtype=dtype) + sel = loc.repeat(2, 1 if broadcast else queries, 1).to(torch.int32) + sel[1] = sel[1].flip(-1) + sel[..., 5] = -1 + cnt = torch.full((2, queries), n, device="cuda", dtype=torch.int32) + cnt[0, 0] = 0 + cnt[1, 0] = 13 + out = glm_dsa_sparse_attn( + q, pool.latent_rows(0), sel, .04, counts=cnt, d_v=512, + pool_scale=pool.latent_scale(0), pool_block_scale=pool.latent_block_scale(0), + kv_quant="nvfp4", force_splits=splits, + ) + decoded = _decode_latent(pool) + ref = torch.zeros_like(out, dtype=torch.float32) + for b in range(2): + for m in range(queries): + rows = sel[b, 0 if broadcast else m, :int(cnt[b, m])] + rows = rows[rows >= 0].long() + if rows.numel(): + kv = decoded[rows] + ref[b, m] = (q[b, m].float() @ kv.T * .04).softmax(-1) @ kv[:, :512] + torch.testing.assert_close(out.float(), ref, atol=2e-2, rtol=1e-2) + + @pytest.mark.parametrize("dim", [16, 48, 64, 128, 256, 512]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) def test_scatter_matches_independent_reference(dim, dtype): diff --git a/tests/kvcache/test_dsa_pool.py b/tests/kvcache/test_dsa_pool.py index 28b130497..8bcdb8d5e 100644 --- a/tests/kvcache/test_dsa_pool.py +++ b/tests/kvcache/test_dsa_pool.py @@ -174,3 +174,57 @@ def test_rebuild_shrink_and_engine_wiring(): pool.rebuild_from_config(config=None, num_pages=63) assert pool.latent_rows(0).shape[0] == 64 # 63 + 1 dummy page assert pool.index_k_cache(0).shape[0] == 64 + + +@pytest.mark.parametrize("kv_quant", ["none", "fp8", "nvfp4"]) +@pytest.mark.parametrize("ratio", [1, 4]) +def test_latent_budget_matches_allocations_and_rebuild(kv_quant, ratio): + from types import SimpleNamespace + from freetoken.attention import AttnType + from freetoken.kvcache.dsa_pool import DSAKVCache, KpoolDSAKVCache + from freetoken.models.config import KVCacheGroupSpec + + cls = KpoolDSAKVCache if ratio > 1 else DSAKVCache + spec = KVCacheGroupSpec( + name="full", layer_ids=(3, 7), num_kv_heads=1, head_dim=512, sliding_window=None, + mla=True, index_head_dim=128, num_index_layers=2, index_ratio=ratio, + attn_type=AttnType.DSA, + ) + cfg = SimpleNamespace( + kv_quant=kv_quant, dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1), + max_running_req=3, page_size=64, + model_config=SimpleNamespace(kv_cache_group_specs=lambda: (spec,)), + ) + extra = dict(index_ratio=ratio, num_req_slots=4) if ratio > 1 else {} + pool = cls(512, 8, 2, 64, torch.bfloat16, torch.device("cuda"), + index_head_dim=128, num_index_layers=2, layer_ids=(3, 7), + kv_quant=kv_quant, **extra) + page_bytes, fixed, page_size, _ = cls.kv_cost(cfg) + for pages in (2, 5, 1): + pool.rebuild_from_config(cfg, pages) + allocated_pages = pages + 1 + buffers = (pool._kv_buffer, pool._scale_buffer, pool._block_scale_buffer, + pool._index_k_buffer) + if ratio > 1: + buffers += (pool._tail_k, pool._tail_gate) + assert pool.cmp_scratch_base == allocated_pages * 64 // ratio + assert pool.tail_k(0).dtype == torch.bfloat16 + actual = sum(b.numel() * b.element_size() for b in buffers if b is not None) + assert actual == allocated_pages * page_bytes + fixed + assert pool.unit_bytes() == (page_bytes // page_size, 0) + assert pool.latent_rows(7).shape == (allocated_pages * 64, 256 if kv_quant == "nvfp4" else 512) + loc = torch.tensor([allocated_pages * 64 - 1], device="cuda") + x = torch.ones(1, 512, device="cuda", dtype=torch.bfloat16) + pool.store_kv(x, x[:, :0], loc, 7) + if kv_quant == "nvfp4": + from tests.kernels.test_kv_nvfp4 import _decode_latent + + torch.testing.assert_close(_decode_latent(pool, 7)[loc], x.float()) + assert pool.k_cache(7).data_ptr() == pool.v_cache(7).data_ptr() + + +def test_nvfp4_latent_rejects_partial_blocks(): + from freetoken.kvcache.dsa_pool import MLAKVCache + + with pytest.raises(ValueError, match="divisible by 16"): + MLAKVCache(72, 1, 1, 1, torch.bfloat16, torch.device("cpu"), kv_quant="nvfp4") diff --git a/tests/models/test_glm5_next_config.py b/tests/models/test_glm5_next_config.py index 8f0bc518e..3a41eed99 100644 --- a/tests/models/test_glm5_next_config.py +++ b/tests/models/test_glm5_next_config.py @@ -23,6 +23,28 @@ _KDA_IDS = tuple(i for i in range(_NUM_LAYERS) if i not in _DSA_IDS) +def test_nvfp4_pool_factory_preserves_glm5_hybrid_geometry(): + import torch + from freetoken.kvcache import create_kvcache_pool + from freetoken.kvcache.dsa_pool import KpoolDSAKVCache + + cfg = parse_config(_hf_config()) + pool = create_kvcache_pool( + cfg, num_pages=2, page_size=64, dtype=torch.bfloat16, + device=torch.device("cpu"), num_req_slots=3, kv_quant="nvfp4", + ) + assert isinstance(pool, KpoolDSAKVCache) + spec, = cfg.kv_cache_group_specs() + assert pool.num_layers == len(_DSA_IDS) + for layer in _DSA_IDS: + assert pool.latent_rows(layer).shape == (128, spec.head_dim // 2) + assert pool.latent_block_scale(layer).shape == (128, spec.head_dim // 16) + with pytest.raises(KeyError): + pool.latent_rows(_KDA_IDS[0]) + assert pool.index_k_cache(0).dtype == torch.bfloat16 + assert pool.tail_gate(0).dtype == torch.bfloat16 + + def _layer_types() -> list[str]: return [ "deepseek_sparse_attention" if i in _DSA_IDS else "linear_attention" diff --git a/tests/models/test_glm_dsa.py b/tests/models/test_glm_dsa.py index 0dceabadd..f09a73822 100644 --- a/tests/models/test_glm_dsa.py +++ b/tests/models/test_glm_dsa.py @@ -225,7 +225,7 @@ def test_splitk_matches_single_program(): assert (o_single.float() - o_split.float()).abs().max().item() < 1e-2 -def _make_backend(dsa: bool, latent=80, dv=64, idx_dim=32, idx_heads=16, topk=64, pages=400): +def _make_backend(dsa: bool, latent=80, dv=64, idx_dim=32, idx_heads=16, topk=64, pages=400, kv_quant="none"): """Minimal ctx + pool + DSAAttnBackend (no engine).""" from types import SimpleNamespace @@ -239,9 +239,9 @@ def _make_backend(dsa: bool, latent=80, dv=64, idx_dim=32, idx_heads=16, topk=64 ctx.page_table = torch.zeros(4, pages, dtype=torch.int32, device="cuda") if dsa: ctx.kv_cache = DSAKVCache(latent, 2, pages, 1, torch.bfloat16, torch.device("cuda"), - index_head_dim=idx_dim, num_index_layers=1) + index_head_dim=idx_dim, num_index_layers=1, kv_quant=kv_quant) else: - ctx.kv_cache = MLAKVCache(latent, 2, pages, 1, torch.bfloat16, torch.device("cuda")) + ctx.kv_cache = MLAKVCache(latent, 2, pages, 1, torch.bfloat16, torch.device("cuda"), kv_quant=kv_quant) set_global_ctx(ctx) args = SimpleNamespace( kv_lora_rank=dv, qk_rope_head_dim=latent - dv, qk_head_dim=latent, @@ -258,7 +258,8 @@ def _ref_attend(q_cat, pool_rows, live_rows, scale, dv): return s.softmax(-1) @ k[:, :dv] -def test_backend_ragged_prefill_identity_and_selection(): +@pytest.mark.parametrize("kv_quant", ["none", "nvfp4"]) +def test_backend_ragged_prefill_identity_and_selection(kv_quant): """Two-request ragged prefill through the BACKEND (page-table slicing, counts = positions + 1, per-request segmentation, leader/follower reuse): request A stays under index_topk (selection == identity == dense), request B @@ -267,7 +268,7 @@ def test_backend_ragged_prefill_identity_and_selection(): torch.manual_seed(5) dv, dr, h, idx_h, idx_d, topk = 64, 16, 8, 16, 32, 64 - backend, ctx = _make_backend(dsa=True, topk=topk) + backend, ctx = _make_backend(dsa=True, topk=topk, kv_quant=kv_quant) pool = ctx.kv_cache scale = backend.sm_scale @@ -311,6 +312,10 @@ def test_backend_ragged_prefill_identity_and_selection(): # request A (kv <= topk): selection covers all live -> equals dense reference q_cat = torch.cat([q_nope, q_pe], -1) slab = pool.latent_rows(0) + if kv_quant == "nvfp4": + from tests.kernels.test_kv_nvfp4 import _decode_latent + + slab = _decode_latent(pool) for j in range(8): # A's queries, positions 32..39 live = ctx.page_table[0, : 33 + j] ref = _ref_attend(q_cat[j], slab, live, scale, dv) @@ -335,15 +340,17 @@ def test_backend_ragged_prefill_identity_and_selection(): assert (o0.float() - o1.float()).abs().max().item() < 3e-2 # identity wiring (dense ablation): same batch through an MLAKVCache backend - backend_d, ctx_d = _make_backend(dsa=False) + backend_d, ctx_d = _make_backend(dsa=False, kv_quant=kv_quant) ctx_d.page_table.copy_(ctx.page_table) - for lid in (0, 1): - ctx_d.kv_cache._kv_buffer.copy_(pool._kv_buffer) + ctx_d.kv_cache._kv_buffer.copy_(pool._kv_buffer) + if kv_quant == "nvfp4": + ctx_d.kv_cache._scale_buffer.copy_(pool._scale_buffer) + ctx_d.kv_cache._block_scale_buffer.copy_(pool._block_scale_buffer) batch_d = SimpleNamespace(reqs=reqs, positions=positions, out_loc=out_loc, active_table_idx=None, attn_metadata=None) backend_d.prepare_metadata(batch_d) od = backend_d.mla_forward(q_nope, q_pe, c_kv, k_rope, 0, batch_d, indexer_inputs=None) - slab_d = ctx_d.kv_cache.latent_rows(0) + slab_d = _decode_latent(ctx_d.kv_cache) if kv_quant == "nvfp4" else ctx_d.kv_cache.latent_rows(0) for j in range(8): live = ctx_d.page_table[0, : 33 + j] ref = _ref_attend(q_cat[j], slab_d, live, scale, dv)