From d07d0a761fe2060faf1a10d8457d7e3e9bdf8de5 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 29 Sep 2026 04:21:26 +0300 Subject: [PATCH 1/4] feat(kvcache): --kv-placement host keeps the QSA K/V in pinned host memory --- python/freetoken/attention/qsa_sparse.py | 123 +++++++++++++++-- python/freetoken/engine/config.py | 3 + python/freetoken/engine/engine.py | 6 +- python/freetoken/kernel/pinned.py | 25 ++++ .../freetoken/kernel/triton/qsa/__init__.py | 3 +- python/freetoken/kernel/triton/qsa/attend.py | 127 +++++++++++++++++- python/freetoken/kvcache/__init__.py | 24 ++++ python/freetoken/kvcache/base.py | 38 ++++-- python/freetoken/kvcache/mha_pool.py | 44 ++++-- python/freetoken/kvcache/qsa_pool.py | 89 +++++++++++- python/freetoken/server/args.py | 9 ++ tests/engine/test_cache_budget.py | 3 + tests/kvcache/test_pool_sizing_surface.py | 3 +- tests/kvcache/test_qsa_pool.py | 98 +++++++++++++- tests/models/qwen4_exp/common.py | 3 + tests/models/qwen4_exp/test_qsa_backend.py | 107 +++++++++++++++ tests/models/qwen4_exp/test_qsa_kernels.py | 107 +++++++++++++++ 17 files changed, 753 insertions(+), 59 deletions(-) diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py index 8ce63c4242..865f02bb59 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -20,7 +20,9 @@ clamped to ``kvlen // index_ratio`` -- slab rows are never cleared, so stale rows must stay unreachable), take the top ``index_budget // index_ratio`` blocks, expand them to token indices plus the causal tail of the open group; -6. attend to exactly those tokens. +6. attend to exactly those tokens -- in place, or, when the pool keeps its K/V in host memory + and this is a prefill, from device windows the batch's pages are staged into + (``plan_staging_windows``). Addressing: the engine pins ``page_size == 64`` (this backend's ``page_sizes``), so a group of ``index_ratio`` tokens never straddles a page and ``block_table[req, p] = page_table[req, p * @@ -33,7 +35,7 @@ import os from dataclasses import dataclass -from typing import TYPE_CHECKING, Callable, List +from typing import TYPE_CHECKING, Callable, List, Sequence import torch from freetoken.core import Batch, get_global_ctx @@ -69,6 +71,45 @@ def _resolve_block_topk() -> Callable | None: return qsa_block_topk +def staging_moves_less(rows: int, select_width: int, cached_tokens: int) -> bool: + """Whether staging a prefill's cached pages moves fewer bytes over PCIe than reading them in + place: staging moves every cached token once, reading in place moves each row's selection.""" + return rows * select_width > cached_tokens + + +@dataclass(frozen=True) +class StagingWindow: + """One device window of a host-placed prefill: the physical pages gathered into it, in + window order, and the block table that reads them there (-1 for every page outside it).""" + + source_pages: torch.Tensor # [n] int64 + block_table: torch.Tensor # [bs, pages] int32 + + +def plan_staging_windows( + block_table: torch.Tensor, page_counts: Sequence[int], capacity: int +) -> list[StagingWindow]: + """Split the batch's cached pages, request by request in logical order, into windows of at + most ``capacity`` pages. One window attends in one pass; more fold by log-sum-exp.""" + device = block_table.device + counts = torch.tensor(page_counts, dtype=torch.int64) + total = int(counts.sum()) + requests = torch.repeat_interleave(torch.arange(len(page_counts)), counts) + first = torch.repeat_interleave(torch.cumsum(counts, 0) - counts, counts) + logical = torch.arange(total) - first + requests, logical = requests.to(device), logical.to(device) + physical = block_table[requests, logical].to(torch.int64) + windows = [] + for start in range(0, total, capacity): + span = slice(start, min(start + capacity, total)) + table = torch.full_like(block_table, -1) + table[requests[span], logical[span]] = torch.arange( + span.stop - span.start, dtype=block_table.dtype, device=device + ) + windows.append(StagingWindow(physical[span], table)) + return windows + + @dataclass class QSASparseMetadata(BaseAttnMetadata): # fmt: off @@ -92,6 +133,8 @@ class QSASparseMetadata(BaseAttnMetadata): rope_rows: torch.Tensor | None = None # [T] int32 arange q_rope_cache: torch.Tensor | None = None # [T, rotary_dim] float32 k_rope_cache: torch.Tensor | None = None # [T, rotary_dim] float32 + # host-placed prefill only, built once per forward: the device windows its pages stage into + staging: list[StagingWindow] | None = None # fmt: on def get_last_indices(self, bs: int) -> torch.Tensor: @@ -282,8 +325,6 @@ def qsa_forward( layer_id: int, batch: Batch, ) -> torch.Tensor: - from freetoken.kernel.triton.qsa import qsa_sparse_paged_attention - md = batch.attn_metadata assert isinstance(md, QSASparseMetadata) slot = self._idx_slot[layer_id] @@ -295,18 +336,66 @@ def qsa_forward( # capture batch runs its warmup and its capture through ONE metadata object, and a # cached plan would bake the warmup's addresses into the graph. self._plan_index_writes(md, batch) + md.staging = self._plan_staging(md, rows=q.shape[0]) self._update_index_cache(index, md, slot) indices = self._select(index, md, slot) - return qsa_sparse_paged_attention( - q, - self.kvcache.k_cache(layer_id), - self.kvcache.v_cache(layer_id), - indices, - md.block_table, - md.token_to_req, - torch.empty_like(q), - ) + return self._attend(q, layer_id, indices, md) + + def _plan_staging(self, md: QSASparseMetadata, rows: int) -> list[StagingWindow] | None: + """The windows a host-placed prefill stages its pages into, or None to read them in place.""" + if self.kvcache.placement == "device" or md.is_decode: + return None + page_counts = [-(-length // self.page_size) for length in md.kv_len_cpu.tolist()] + if not staging_moves_less(rows, self.select_width, sum(page_counts) * self.page_size): + return None + capacity = self.kvcache.staging_window()[0].shape[0] + return plan_staging_windows(md.block_table, page_counts, capacity) + + def _attend( + self, q: torch.Tensor, layer_id: int, indices: torch.Tensor, md: QSASparseMetadata + ) -> torch.Tensor: + from freetoken.kernel.triton.qsa import fold_attention_pass, qsa_sparse_paged_attention + + k_cache, v_cache = self.kvcache.k_cache(layer_id), self.kvcache.v_cache(layer_id) + out = torch.empty_like(q) + if md.staging is None: + return qsa_sparse_paged_attention( + q, k_cache, v_cache, indices, md.block_table, md.token_to_req, out + ) + window_k, window_v = self.kvcache.staging_window() + if len(md.staging) == 1: + (window,) = md.staging + self._stage(window, k_cache, v_cache, window_k, window_v) + return qsa_sparse_paged_attention( + q, window_k, window_v, indices, window.block_table, md.token_to_req, out + ) + total, total_lse, part_lse = self.kvcache.fold_workspace(q.shape[0]) + total.zero_() + total_lse.fill_(-float("inf")) + last = len(md.staging) - 1 + for pass_index, window in enumerate(md.staging): + self._stage(window, k_cache, v_cache, window_k, window_v) + qsa_sparse_paged_attention( + q, window_k, window_v, indices, window.block_table, md.token_to_req, out, + lse=part_lse, + ) + fold_attention_pass( + total, total_lse, out, part_lse, result=out if pass_index == last else None + ) + return out + + @staticmethod + def _stage( + window: StagingWindow, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + window_k: torch.Tensor, + window_v: torch.Tensor, + ) -> None: + pages = window.source_pages.numel() + torch.index_select(k_cache, 0, window.source_pages, out=window_k[:pages]) + torch.index_select(v_cache, 0, window.source_pages, out=window_v[:pages]) def _plan_index_writes(self, md: QSASparseMetadata, batch: Batch) -> None: """Per-token slab row and ring row for this forward; the other QSA layers reuse it @@ -539,4 +628,10 @@ def reset_capture(self) -> None: self._graph = {} -__all__ = ["QSASparseAttnBackend", "QSASparseMetadata"] +__all__ = [ + "QSASparseAttnBackend", + "QSASparseMetadata", + "StagingWindow", + "plan_staging_windows", + "staging_moves_less", +] diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index b6fa8b8949..6c77d05d53 100644 --- a/python/freetoken/engine/config.py +++ b/python/freetoken/engine/config.py @@ -7,6 +7,7 @@ import torch from freetoken.distributed import DistributedInfo +from freetoken.kvcache.base import KVPlacement from freetoken.layers.quantization import set_quant_config from freetoken.mm.config import ENCODER_SECTIONS, MultimodalConfig from freetoken.models.register import EncoderSpec, ModelSpec, _load_attr, checkpoint_quant_config, get_model_spec @@ -94,6 +95,8 @@ class EngineConfig: # KV capacity in tokens; resolved into num_page_override by _adjust_config once page_size # is final. Mutually exclusive with num_page_override. num_token_override: int | None = None + # --kv-placement: "host" keeps the paged K/V in pinned host memory the GPU reads over PCIe + kv_placement: KVPlacement = "device" # Runtime knobs of the multimodal path; the architecture side (vision_config, mrope) lives in ModelConfig. mm: MultimodalConfig = field(default_factory=MultimodalConfig) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 3988caf602..3ecaeeccc9 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -29,7 +29,7 @@ from .config import EngineConfig from .graph import GraphRunner, get_free_memory from .sample import BatchSamplingArgs, Sampler -from freetoken.kvcache import create_kv_pool, resolve_pool_class +from freetoken.kvcache import check_kv_placement, create_kv_pool, resolve_pool_class from freetoken.kvcache.base import CacheRebuildRejected from freetoken.kvcache.cache_status import _supports_swa_ratio from freetoken.kvcache.linear_state_pool import ( @@ -1535,6 +1535,10 @@ def override(attr: str, value: Any): # this is dangerous, use with caution is_moe = getattr(model_config, "is_moe", False) expert_quant = getattr(model_config, "expert_quant", "none") + # Before any weight is read: a placement the pool cannot serve would otherwise fail + # minutes later, after the expert banks are loaded. + check_kv_placement(model_config, config.kv_placement) + if not is_moe: # A dense model has no routed experts: the MoE knobs are inert, and the offload family # is worse than inert -- engine init would build an expert cache for a model that has diff --git a/python/freetoken/kernel/pinned.py b/python/freetoken/kernel/pinned.py index a21e78ac4c..a804051693 100644 --- a/python/freetoken/kernel/pinned.py +++ b/python/freetoken/kernel/pinned.py @@ -56,6 +56,31 @@ def _host_ptr_identity() -> bool: return bool(_load_pinned_extension().host_ptr_identity()) +class _MappedHostArray: + """``__cuda_array_interface__`` over a pinned+mapped host tensor, so torch aliases it as a + CUDA tensor without a copy; the alias holds this object, and this object holds the host.""" + + def __init__(self, host: torch.Tensor) -> None: + self.host = host + self.__cuda_array_interface__ = { + "shape": tuple(host.shape), + "strides": None, + # bfloat16 has no typestr: export the same-width integer and view the dtype back + "typestr": f" torch.Tensor: + """A CUDA tensor over pinned+mapped host memory: kernels read and write it in place over + PCIe, and it keeps ``host`` alive for as long as it or any view of it lives.""" + with torch.cuda.device(device): + # Inside the device context: torch reports no pointer as pinned before its CUDA context exists. + assert host.is_pinned() and host.is_contiguous(), "only contiguous pinned+mapped host memory aliases" + return torch.as_tensor(_MappedHostArray(host), device=device).view(host.dtype) + + def device_ptr(t: torch.Tensor) -> int: """Base address of ``t`` as the GPU must dereference it. diff --git a/python/freetoken/kernel/triton/qsa/__init__.py b/python/freetoken/kernel/triton/qsa/__init__.py index 91753b7922..e7f78e2b79 100644 --- a/python/freetoken/kernel/triton/qsa/__init__.py +++ b/python/freetoken/kernel/triton/qsa/__init__.py @@ -1,6 +1,6 @@ """Triton kernels for Qwen3.8-Flash-Next QSA sparse attention.""" -from .attend import qsa_sparse_paged_attention +from .attend import fold_attention_pass, qsa_sparse_paged_attention from .compress import qsa_compress_groups, qsa_index_norm_rope, qsa_store_rows from .expand import expand_qsa_block_indices from .score import qsa_mqa_paged @@ -8,6 +8,7 @@ __all__ = [ "expand_qsa_block_indices", + "fold_attention_pass", "qsa_block_topk", "qsa_block_topk_scratch_width", "qsa_compress_groups", diff --git a/python/freetoken/kernel/triton/qsa/attend.py b/python/freetoken/kernel/triton/qsa/attend.py index 541e27c680..5a60fdaacd 100644 --- a/python/freetoken/kernel/triton/qsa/attend.py +++ b/python/freetoken/kernel/triton/qsa/attend.py @@ -21,6 +21,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + lse_ptr, stride_q_row, stride_q_head, stride_k_block, @@ -46,6 +47,7 @@ def _qsa_sparse_paged_gqa_splitk_kernel( NUM_TILES: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, + WRITE_LSE: tl.constexpr, ) -> None: # row * stride can overflow int32 for large row counts. row = tl.program_id(0).to(tl.int64) @@ -143,6 +145,11 @@ def _qsa_sparse_paged_gqa_splitk_kernel( 0.0, ) output_mask = head_offsets[:, None] < GROUP_SIZE + partial_lse = tl.where( + has_values, + max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)), + -float("inf"), + ) if NUM_SPLITS == 1: tl.store( output_ptr @@ -152,12 +159,13 @@ def _qsa_sparse_paged_gqa_splitk_kernel( normalized_output, mask=output_mask, ) + if WRITE_LSE: + tl.store( + lse_ptr + row * NUM_QUERY_HEADS + first_head + head_offsets, + partial_lse, + mask=head_offsets < GROUP_SIZE, + ) else: - partial_lse = tl.where( - has_values, - max_value + tl.math.log2(tl.maximum(normalizer, 1.0e-20)), - -float("inf"), - ) tl.store( partial_output_ptr + ( @@ -185,6 +193,7 @@ def _qsa_merge_splitk_kernel( partial_output_ptr, partial_lse_ptr, output_ptr, + lse_ptr, stride_output_row, stride_output_head, num_rows, @@ -192,6 +201,7 @@ def _qsa_merge_splitk_kernel( NUM_QUERY_HEADS: tl.constexpr, NUM_SPLITS: tl.constexpr, BLOCK_SPLITS: tl.constexpr, + WRITE_LSE: tl.constexpr, ) -> None: row = tl.program_id(0).to(tl.int64) head = tl.program_id(1) @@ -222,6 +232,11 @@ def _qsa_merge_splitk_kernel( output_ptr + row * stride_output_row + head * stride_output_head + dim_offsets, merged, ) + if WRITE_LSE: + tl.store( + lse_ptr + row * NUM_QUERY_HEADS + head, + tl.where(has_values, lse_max + tl.math.log2(tl.maximum(denominator, 1.0e-20)), -float("inf")), + ) def qsa_sparse_paged_attention( @@ -232,8 +247,10 @@ def qsa_sparse_paged_attention( block_table: torch.Tensor, token_to_req: torch.Tensor, out: torch.Tensor | None = None, + lse: torch.Tensor | None = None, ) -> torch.Tensor: - """Run sparse GQA directly over paged BF16 K/V caches.""" + """Run sparse GQA directly over paged BF16 K/V caches; ``lse``, when given, receives each + head's log2-sum-exp so passes over disjoint pages fold (``fold_attention_pass``).""" 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") @@ -256,6 +273,8 @@ def qsa_sparse_paged_attention( if out is None: out = torch.empty_like(q) assert out.shape == q.shape and out.dtype == q.dtype and out.stride(2) == 1 + if lse is not None: + assert lse.shape == q.shape[:2] and lse.dtype == torch.float32 and lse.is_contiguous() if not q.shape[0]: return out @@ -309,6 +328,7 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + out if lse is None else lse, q.stride(0), q.stride(1), k_cache.stride(0), @@ -334,6 +354,7 @@ def qsa_sparse_paged_attention( NUM_TILES=num_tiles, BLOCK_M=block_m, BLOCK_N=block_n, + WRITE_LSE=lse is not None, num_warps=partial_warps, num_stages=2, ) @@ -344,6 +365,7 @@ def qsa_sparse_paged_attention( partial_output, partial_lse, out, + out if lse is None else lse, out.stride(0), out.stride(1), q.shape[0], @@ -351,10 +373,101 @@ def qsa_sparse_paged_attention( NUM_QUERY_HEADS=q.shape[1], NUM_SPLITS=num_splits, BLOCK_SPLITS=triton.next_power_of_2(num_splits), + WRITE_LSE=lse is not None, num_warps=2, num_stages=1, ) return out -__all__ = ["qsa_sparse_paged_attention"] +@triton.jit +def _fold_attention_pass_kernel( + total_ptr, + total_lse_ptr, + part_ptr, + part_lse_ptr, + result_ptr, + stride_total_row, + stride_total_head, + stride_part_row, + stride_part_head, + stride_result_row, + stride_result_head, + stride_lse_row, + HEAD_DIM: tl.constexpr, + BLOCK_D: tl.constexpr, + FINAL: tl.constexpr, +) -> None: + row = tl.program_id(0) + head = tl.program_id(1) + lse_at = row * stride_lse_row + head + total_lse = tl.load(total_lse_ptr + lse_at) + part_lse = tl.load(part_lse_ptr + lse_at) + peak = tl.maximum(total_lse, part_lse) + empty = peak == -float("inf") + shift = tl.where(empty, 0.0, peak) + keep = tl.exp2(total_lse - shift) + take = tl.exp2(part_lse - shift) + norm = keep + take + safe_norm = tl.where(empty, 1.0, norm) + keep = tl.where(empty, 0.0, keep / safe_norm) + take = tl.where(empty, 0.0, take / safe_norm) + + dims = tl.arange(0, BLOCK_D) + inside = dims < HEAD_DIM + total_at = total_ptr + row * stride_total_row + head * stride_total_head + dims + total = tl.load(total_at, mask=inside, other=0.0) + part = tl.load( + part_ptr + row * stride_part_row + head * stride_part_head + dims, mask=inside, other=0.0 + ).to(tl.float32) + folded = total * keep + part * take + if FINAL: + result_at = result_ptr + row * stride_result_row + head * stride_result_head + dims + tl.store(result_at, folded.to(result_ptr.dtype.element_ty), mask=inside) + else: + tl.store(total_at, folded, mask=inside) + tl.store(total_lse_ptr + lse_at, shift + tl.log2(norm)) + + +def fold_attention_pass( + total: torch.Tensor, + total_lse: torch.Tensor, + part: torch.Tensor, + part_lse: torch.Tensor, + result: torch.Tensor | None = None, +) -> None: + """Fold one pass into ``total`` (float32) by log2-sum-exp share; given ``result`` (which may + be ``part``), the last fold writes the sum there and leaves ``total`` as it was.""" + rows, heads, head_dim = part.shape + target = total if result is None else result + if total.shape != part.shape or target.shape != part.shape or total.dtype != torch.float32: + raise ValueError("fold_attention_pass takes a float32 total shaped like the pass") + if total_lse.shape != (rows, heads) or part_lse.shape != (rows, heads): + raise ValueError("fold_attention_pass takes one log2-sum-exp per row and head") + if total_lse.stride() != part_lse.stride() or total_lse.stride(1) != 1: + raise ValueError("fold_attention_pass takes row-major log2-sum-exp rows of one stride") + if total.stride(2) != 1 or part.stride(2) != 1 or target.stride(2) != 1: + raise ValueError("fold_attention_pass takes head-contiguous outputs") + if not rows or not heads: + return + _fold_attention_pass_kernel[(rows, heads)]( + total, + total_lse, + part, + part_lse, + target, + total.stride(0), + total.stride(1), + part.stride(0), + part.stride(1), + target.stride(0), + target.stride(1), + total_lse.stride(0), + HEAD_DIM=head_dim, + BLOCK_D=triton.next_power_of_2(head_dim), + FINAL=result is not None, + num_warps=1, + ) + + +__all__ = ["fold_attention_pass", "qsa_sparse_paged_attention"] diff --git a/python/freetoken/kvcache/__init__.py b/python/freetoken/kvcache/__init__.py index e0cf606db1..723c3bf4a8 100644 --- a/python/freetoken/kvcache/__init__.py +++ b/python/freetoken/kvcache/__init__.py @@ -12,6 +12,7 @@ BaseCacheHandle, BaseKVCachePool, BasePrefixCache, + KVPlacement, MatchResult, SizeInfo, ) @@ -80,6 +81,20 @@ def resolve_pool_class(model_config: ModelConfig) -> type[BaseKVCachePool]: return MHAKVCache +def check_kv_placement(model_config: ModelConfig, placement: KVPlacement) -> None: + """Refuse a K/V placement the model's pool family cannot serve: only a sparse attention + reads few enough cached tokens per query for host memory to feed it.""" + if placement == "device": + return + from .qsa_pool import QSAKVCache + + if resolve_pool_class(model_config) is not QSAKVCache: + raise ValueError( + f"--kv-placement {placement} is served by the QSA sparse pool only; this model " + "attends densely and would read its whole cache over PCIe for every query" + ) + + 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 @@ -136,6 +151,8 @@ 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 + placement=config.kv_placement, + prefill_rows=config.max_forward_len, ) @@ -147,7 +164,10 @@ def create_kvcache_pool( device: torch.device, num_swa_tokens: int | None = None, num_req_slots: int | None = None, + placement: KVPlacement = "device", + prefill_rows: int = 0, ) -> BaseKVCachePool: + check_kv_placement(model_config, placement) if model_config.has_swa_attention: from .hybrid_swa_pool import HybridSWAKVCache @@ -224,6 +244,9 @@ def create_kvcache_pool( num_req_slots=num_req_slots, layer_ids=spec.layer_ids, mrope=model_config.model_is_mrope, + placement=placement, + num_qo_heads=model_config.num_qo_heads, + prefill_rows=prefill_rows, ) if len(kv_specs) == 1 and kv_specs[0].mla: @@ -305,6 +328,7 @@ def create_prefix_cache( __all__ = [ + "check_kv_placement", "create_kv_pool", "create_kvcache_pool", "create_prefix_cache", diff --git a/python/freetoken/kvcache/base.py b/python/freetoken/kvcache/base.py index 569da1cb6c..66ff756f3e 100644 --- a/python/freetoken/kvcache/base.py +++ b/python/freetoken/kvcache/base.py @@ -2,13 +2,16 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import ClassVar, NamedTuple +from typing import ClassVar, Literal, NamedTuple, TypeAlias import torch from freetoken.utils import div_even, init_logger, mem_GB logger = init_logger(__name__) +# Where a pool keeps its paged K/V: device memory, or pinned host memory the GPU reads in place. +KVPlacement: TypeAlias = Literal["device", "host"] + class CacheRebuildRejected(Exception): """A runtime cache rebuild was rejected BEFORE any destructive free (e.g. the @@ -16,23 +19,32 @@ class CacheRebuildRejected(Exception): this is recoverable, unlike a failure after the free.""" -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.) - - ``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.""" - per_token = ( +def spec_kv_slab_bytes_per_token(spec, config) -> int: + """One paged-KV group's K/V slab bytes per token: (1|2 slabs) x head_dim x local kv heads + x dtype x layers.""" + return ( (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 * spec.num_layers ) - return per_token + spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio + + +def spec_index_bytes_per_token(spec) -> int: + """The bf16 index-key slab's bytes per token when the spec carries indexer dims. + (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.""" + return spec.index_head_dim * spec.num_index_layers * 2 // spec.index_ratio + + +def spec_kv_bytes_per_token(spec, config) -> int: + """One paged-KV group's bytes per token: its K/V slab plus its index-key slab. Pure + per-spec arithmetic -- pool families compose it over THEIR OWN groups; no family + branching here.""" + return spec_kv_slab_bytes_per_token(spec, config) + spec_index_bytes_per_token(spec) @dataclass(frozen=True) @@ -91,7 +103,7 @@ def solve_num_pages(cls, config, available_memory: int) -> int: real_kv_size = num_pages * cache_per_page + fixed_cache_size logger.info( f"Allocating {num_pages * config.page_size} tokens for KV cache, " - f"K + V = {mem_GB(real_kv_size)}" + f"{mem_GB(real_kv_size)} of it on the device" ) return num_pages diff --git a/python/freetoken/kvcache/mha_pool.py b/python/freetoken/kvcache/mha_pool.py index 8ed280b966..ca0613416d 100644 --- a/python/freetoken/kvcache/mha_pool.py +++ b/python/freetoken/kvcache/mha_pool.py @@ -6,7 +6,7 @@ from freetoken.distributed import get_tp_info from freetoken.utils import div_even -from .base import BaseKVCachePool +from .base import BaseKVCachePool, KVPlacement class MHAKVCache(BaseKVCachePool): @@ -20,6 +20,9 @@ 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. + + ``placement="host"`` keeps the K/V slabs in pinned+mapped host memory behind a CUDA alias: + every kernel reads and writes them in place over PCIe, and they cost no device memory. """ def __init__( @@ -32,6 +35,7 @@ def __init__( dtype: torch.dtype, device: torch.device, layer_ids: Sequence[int] | None = None, + placement: KVPlacement = "device", ) -> None: tp_info = get_tp_info() local_kv_heads = div_even(num_kv_heads, tp_info.size, allow_replicate=True) @@ -47,16 +51,26 @@ 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._placement: KVPlacement = placement + self._device = device + self._kv_buffer = self._allocate_kv( + (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), dtype ) 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 _allocate_kv(self, shape: tuple[int, ...], dtype: torch.dtype) -> torch.Tensor: + match self._placement: + case "device": + return torch.empty(shape, device=self._device, dtype=dtype) + case "host": + from freetoken.kernel.pinned import alloc_pinned_tensor, mapped_cuda_view + + return mapped_cuda_view(alloc_pinned_tensor(*shape, dtype=dtype), self._device) + case _: + raise ValueError(f"unknown KV placement {self._placement!r}") + def rebuild(self, num_pages: int) -> None: """Reallocate the KV buffer for ``num_pages`` pages IN PLACE. @@ -66,17 +80,14 @@ def rebuild(self, num_pages: int) -> None: """ _, 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 - if device.type == "cuda": - torch.cuda.synchronize(device) + if self._device.type == "cuda": + torch.cuda.synchronize(self._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._kv_buffer = self._allocate_kv( + (2, num_storage_layers, num_pages, page_size, local_kv_heads, head_dim), dtype ) self._k_buffer = self._kv_buffer[0] self._v_buffer = self._kv_buffer[1] @@ -99,10 +110,17 @@ def rebuild_from_config( self.rebuild(num_pages + 1) # +1 for the dummy page (matches create_kvcache_pool) def unit_bytes(self) -> tuple[int, int]: + """Device bytes per token; host-placed K/V costs none.""" + if self._placement == "host": + return 0, 0 buf = self._kv_buffer tokens = int(buf.shape[2]) * int(buf.shape[3]) return int(buf.numel() * buf.element_size()) // tokens, 0 + @property + def placement(self) -> KVPlacement: + return self._placement + def _dense(self, layer_id: int) -> int: if self._layer_map is None: return layer_id diff --git a/python/freetoken/kvcache/qsa_pool.py b/python/freetoken/kvcache/qsa_pool.py index 7de70977d5..9498888058 100644 --- a/python/freetoken/kvcache/qsa_pool.py +++ b/python/freetoken/kvcache/qsa_pool.py @@ -18,6 +18,12 @@ The slab is amortized into the per-token KV price (``unit_bytes``); the ring and scratch are fixed and priced through ``kv_cost``'s ``fixed_cache_size``. + +Placed on the host (``--kv-placement host``), the K/V slabs live in pinned host memory and only +the index tiers stay on the device, so the context is bounded by host RAM. Decode reads its +selection in place over PCIe. A prefill cannot: every query row reads its own selection, so a +chunk would move each cached token over the bus hundreds of times. It stages the batch's pages +into a device window once per layer instead (``staging_window``), a fixed device cost. """ from __future__ import annotations @@ -26,11 +32,18 @@ from typing import Sequence import torch +from freetoken.utils import init_logger, mem_GB +from .base import KVPlacement from .mha_pool import MHAKVCache +logger = init_logger(__name__) + # The index tiers are always 2-byte (compute dtype); spec_kv_bytes_per_token budgets the same. _INDEX_DTYPE_BYTES = 2 +# Tokens of one layer's K/V a host-placed pool stages on the device for a prefill. Wider costs +# VRAM the expert cache could hold; a context past it is attended in several passes. +STAGING_TOKENS = 1 << 17 # t/h/w int32 rope position kept per KV slot on mrope models _ROPE_POS_BYTES = 3 * 4 @@ -66,7 +79,16 @@ def __init__( ring_capacity: int | None = None, layer_ids: Sequence[int] | None = None, mrope: bool = False, + placement: KVPlacement = "device", + num_qo_heads: int = 0, + prefill_rows: int = 0, ) -> None: + """``num_qo_heads`` and ``prefill_rows`` size a host-placed pool's fold accumulator: the + query heads of one attention layer, and the most rows one forward attends.""" + if placement == "host" and (num_qo_heads < 1 or prefill_rows < 1): + raise ValueError( + "a host-placed QSA pool needs num_qo_heads and prefill_rows for its fold accumulator" + ) 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( @@ -102,9 +124,38 @@ def __init__( dtype=dtype, device=device, layer_ids=layer_ids, + placement=placement, ) self._zero_kv_slabs() self._alloc_index_tiers(num_pages) + self._staging: torch.Tensor | None = None + self._fold: tuple[torch.Tensor, torch.Tensor] | None = None + if placement == "host": + self._staging = self._alloc_staging_window() + self._fold = self._alloc_fold_workspace(prefill_rows, num_qo_heads, head_dim) + logger.info( + f"K/V in pinned host memory: {mem_GB(self.host_bytes)}; the device keeps the " + f"index tiers, a {STAGING_TOKENS}-token staging window and a {prefill_rows}-row " + "fold accumulator" + ) + + def _alloc_staging_window(self) -> torch.Tensor: + _, _, _, page_size, kv_heads, head_dim = self._kv_buffer.shape + return torch.zeros( + (2, STAGING_TOKENS // page_size, page_size, kv_heads, head_dim), + dtype=self._kv_buffer.dtype, + device=self._device, + ) + + def _alloc_fold_workspace( + self, rows: int, heads: int, head_dim: int + ) -> tuple[torch.Tensor, torch.Tensor]: + """float32 output and log2-sum-exp rows a prefill spanning several windows folds its + passes into; held for the pool's life so the deepest prefill allocates nothing more.""" + return ( + torch.zeros((rows, heads, head_dim), dtype=torch.float32, device=self._device), + torch.zeros((2, rows, heads), dtype=torch.float32, device=self._device), + ) def _zero_kv_slabs(self) -> None: # Defense-in-depth: the attend kernels pos-mask every K/V load (the real fix for @@ -160,7 +211,9 @@ def rebuild(self, num_pages: int) -> None: @classmethod def kv_cost(cls, config) -> tuple[int, int, int, int]: - from .base import spec_kv_bytes_per_token + """Device bytes only: host-placed K/V costs none per token, and its staging window and + fold accumulator are fixed.""" + from .base import spec_index_bytes_per_token, spec_kv_slab_bytes_per_token from freetoken.attention import AttnType num_req_slots = config.max_running_req + 1 @@ -169,7 +222,14 @@ def kv_cost(cls, config) -> tuple[int, int, int, int]: for spec in config.model_config.kv_cache_group_specs(): if spec.is_swa: continue - per_token += spec_kv_bytes_per_token(spec, config) + slab = spec_kv_slab_bytes_per_token(spec, config) + per_token += spec_index_bytes_per_token(spec) + if config.kv_placement == "host": + fixed += slab * STAGING_TOKENS // spec.num_layers + # float32: the output rows [rows, heads, head_dim] and two lse rows [rows, heads] + fixed += config.max_forward_len * config.model_config.num_qo_heads * (spec.head_dim + 2) * 4 + else: + per_token += slab if spec.attn_type is AttnType.QSA: # One index-key row = all index layers at one position. row = spec.index_head_dim * spec.num_index_layers * _INDEX_DTYPE_BYTES @@ -179,8 +239,8 @@ def kv_cost(cls, config) -> tuple[int, int, int, int]: return per_token * config.page_size, fixed, config.page_size, 0 def unit_bytes(self) -> tuple[int, int]: - # Only the shadow slab scales with pages, and only its non-scratch rows; the ring and - # the scratch rows are the fixed term kv_cost reports separately. + # Only the shadow slab scales with pages, and only its non-scratch rows; the ring, the + # scratch rows and a host pool's staging window are the fixed term kv_cost reports. kv, swa = super().unit_bytes() tokens = int(self._kv_buffer.shape[2]) * int(self._kv_buffer.shape[3]) slab = ( @@ -191,6 +251,27 @@ def unit_bytes(self) -> tuple[int, int]: ) return kv + slab // tokens + (_ROPE_POS_BYTES if self._mrope else 0), swa + def staging_window(self) -> tuple[torch.Tensor, torch.Tensor]: + """Device K and V pages ``[pages, page_size, kv_heads, head_dim]`` a host-placed pool's + prefill gathers one layer's pages into; shared by every layer, one at a time.""" + assert self._staging is not None, "only a host-placed pool stages its K/V" + return self._staging[0], self._staging[1] + + def fold_workspace(self, rows: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """float32 ``total [rows, heads, head_dim]``, ``total_lse`` and ``part_lse [rows, heads]`` + a host-placed prefill folds its window passes into; the caller initializes them.""" + assert self._fold is not None, "only a host-placed pool folds window passes" + total, lse = self._fold + assert rows <= total.shape[0], f"{rows} rows exceed the {total.shape[0]} prefill rows" + return total[:rows], lse[0, :rows], lse[1, :rows] + + @property + def host_bytes(self) -> int: + """Pinned host bytes the K/V slabs hold; zero on the device.""" + if self.placement == "device": + return 0 + return self._kv_buffer.numel() * self._kv_buffer.element_size() + def cmp_k_cache(self, slot: int) -> torch.Tensor: """Compressed index keys of one sparse layer: ``[rows, index_head_dim]``.""" return self._cmp_k_buffer[slot] diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 7c0ed9b047..e62d1a13d8 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -432,6 +432,15 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--kv-placement", + choices=["device", "host"], + default=ServerArgs.kv_placement, + help="Where the paged K/V lives. 'host' keeps it in pinned host memory the GPU reads " + "over PCIe, so the context is bounded by host RAM instead of VRAM; the attention " + "backend must read sparsely (qsa_sparse).", + ) + parser.add_argument( "--page-size", type=int, diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a77409f013..f6463bc897 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -147,6 +147,7 @@ class Cfg: moe_cpu_layers = None num_page_override = None num_token_override = None + kv_placement = "device" @property def model_config(self): @@ -220,6 +221,7 @@ class Cfg: attention_backend = "fi" num_page_override = None num_token_override = 5000 + kv_placement = "device" @property def model_config(self): @@ -459,6 +461,7 @@ class Cfg: attention_backend = "triton" num_page_override = None num_token_override = None + kv_placement = "device" max_seq_len_override = None @property diff --git a/tests/kvcache/test_pool_sizing_surface.py b/tests/kvcache/test_pool_sizing_surface.py index 33f27983de..49571847f5 100644 --- a/tests/kvcache/test_pool_sizing_surface.py +++ b/tests/kvcache/test_pool_sizing_surface.py @@ -215,6 +215,7 @@ def test_create_kv_pool_builds_the_right_family(): config = SimpleNamespace( model_config=mc, page_size=1, cache_type="radix", max_running_req=2, swa_full_tokens_ratio=1.0, swa_num_pages_override=None, max_seq_len=64, + max_forward_len=64, kv_placement="device", ) pool = create_kv_pool(config, num_pages=8, device=torch.device("cpu"), dtype=torch.bfloat16) assert isinstance(pool, MHAKVCache) @@ -232,7 +233,7 @@ def test_create_kv_pool_builds_the_right_family(): config2 = SimpleNamespace( model_config=mc2, page_size=1, cache_type="swa_radix", max_running_req=2, swa_full_tokens_ratio=0.5, swa_num_pages_override=None, max_seq_len=64, - max_extend_tokens=64, + max_extend_tokens=64, max_forward_len=64, kv_placement="device", ) pool2 = create_kv_pool(config2, num_pages=8, device=torch.device("cpu"), dtype=torch.bfloat16) assert isinstance(pool2, HybridSWAKVCache) diff --git a/tests/kvcache/test_qsa_pool.py b/tests/kvcache/test_qsa_pool.py index dea722f631..69109d7d0b 100644 --- a/tests/kvcache/test_qsa_pool.py +++ b/tests/kvcache/test_qsa_pool.py @@ -15,7 +15,11 @@ import torch from freetoken.attention import AttnType -from freetoken.kvcache.base import spec_kv_bytes_per_token +from freetoken.kvcache.base import ( + spec_index_bytes_per_token, + spec_kv_bytes_per_token, + spec_kv_slab_bytes_per_token, +) from freetoken.kvcache.qsa_pool import QSAKVCache from freetoken.models.config import KVCacheGroupSpec @@ -37,7 +41,11 @@ def _tp(monkeypatch): ) -def _pool(num_pages=4, page_size=64, index_ratio=4, num_req_slots=4, ring_capacity=None): +QO_HEADS = 4 +PREFILL_ROWS = 16 + + +def _pool(num_pages=4, page_size=64, index_ratio=4, num_req_slots=4, ring_capacity=None, placement="device"): return QSAKVCache( num_kv_heads=2, num_layers=8, @@ -52,6 +60,9 @@ def _pool(num_pages=4, page_size=64, index_ratio=4, num_req_slots=4, ring_capaci num_req_slots=num_req_slots, ring_capacity=ring_capacity, layer_ids=(1, 3, 5, 7), + placement=placement, + num_qo_heads=QO_HEADS, + prefill_rows=PREFILL_ROWS, ) @@ -70,8 +81,11 @@ def _spec(*, index_ratio=4, attn_type=AttnType.QSA, num_kv_heads=2, head_dim=64, ) -def _config(spec, *, page_size=64, max_running_req=3): - mc = SimpleNamespace(num_layers=8, has_swa_attention=False, has_linear_attention=True, model_is_mrope=False) +def _config(spec, *, page_size=64, max_running_req=3, kv_placement="device"): + mc = SimpleNamespace( + num_layers=8, has_swa_attention=False, has_linear_attention=True, model_is_mrope=False, + num_qo_heads=QO_HEADS, + ) mc.kv_cache_group_specs = lambda: (spec,) return SimpleNamespace( model_config=mc, @@ -79,6 +93,8 @@ def _config(spec, *, page_size=64, max_running_req=3): dtype=torch.bfloat16, tp_info=SimpleNamespace(size=1), max_running_req=max_running_req, + max_forward_len=PREFILL_ROWS, + kv_placement=kv_placement, ) @@ -203,7 +219,7 @@ def test_resolve_pool_class_and_factory(): mc = SimpleNamespace( model_is_mrope=False, num_layers=8, has_swa_attention=False, has_linear_attention=True, - num_kv_heads=2, head_dim=64, dsv4_args=None, + num_kv_heads=2, head_dim=64, dsv4_args=None, num_qo_heads=QO_HEADS, ) mc.kv_cache_group_specs = lambda: (spec,) assert resolve_pool_class(mc) is QSAKVCache @@ -217,3 +233,75 @@ def test_resolve_pool_class_and_factory(): with pytest.raises(ValueError, match="num_req_slots"): create_kvcache_pool(mc, num_pages=4, page_size=64, dtype=torch.bfloat16, device=DEV) + + +# ------------------------------------------------------------------------- host placement + + +@pytest.fixture +def host_memory_on_the_cpu(monkeypatch): + """Host placement without a GPU: pinned memory is plain memory and its alias is itself.""" + import freetoken.kernel.pinned as pinned + import freetoken.kvcache.qsa_pool as qsa_pool + + monkeypatch.setattr(pinned, "alloc_pinned_tensor", lambda *shape, dtype: torch.empty(shape, dtype=dtype)) + monkeypatch.setattr(pinned, "mapped_cuda_view", lambda host, device: host) + monkeypatch.setattr(qsa_pool, "STAGING_TOKENS", 128) + + +def test_host_placement_keeps_the_index_tiers_and_one_window_on_the_device(host_memory_on_the_cpu): + pool = _pool(num_pages=4, placement="host") + assert pool.placement == "host" + # K/V: 2 slabs x 4 layers x 4 pages x 64 tokens x 2 heads x 64 dims x 2 bytes, all on the host + assert pool.host_bytes == 2 * 4 * 4 * 64 * 2 * 64 * 2 + assert _pool(num_pages=4).host_bytes == 0 + kv_bytes, swa_bytes = pool.unit_bytes() + assert (kv_bytes, swa_bytes) == (spec_index_bytes_per_token(_spec()), 0) + window_k, window_v = pool.staging_window() + assert window_k.shape == window_v.shape == (128 // 64, 64, 2, 64) + assert pool.cmp_k_cache(0).shape == (4 * 64 // 4 + 4, 32) + total, total_lse, part_lse = pool.fold_workspace(5) + assert total.shape == (5, QO_HEADS, 64) and total.dtype == torch.float32 + assert total_lse.shape == part_lse.shape == (5, QO_HEADS) and total_lse.dtype == torch.float32 + with pytest.raises(AssertionError, match="prefill rows"): + pool.fold_workspace(PREFILL_ROWS + 1) + + pool.rebuild(8) + assert pool.k_cache(1).shape == (8, 64, 2, 64) + assert pool.host_bytes == 2 * 4 * 8 * 64 * 2 * 64 * 2 + assert pool.staging_window()[0].shape == (2, 64, 2, 64) + assert pool.fold_workspace(PREFILL_ROWS)[0].shape == (PREFILL_ROWS, QO_HEADS, 64) + with pytest.raises(AssertionError, match="host-placed"): + _pool().staging_window() + with pytest.raises(AssertionError, match="host-placed"): + _pool().fold_workspace(1) + + +def test_host_placement_prices_its_device_workspace_instead_of_every_token(host_memory_on_the_cpu): + spec = _spec() + config = _config(spec) + device_page, device_fixed, _, _ = QSAKVCache.kv_cost(config) + host_page, host_fixed, _, _ = QSAKVCache.kv_cost(_config(spec, kv_placement="host")) + slab = spec_kv_slab_bytes_per_token(spec, config) + assert device_page - host_page == slab * 64 + assert host_page == spec_index_bytes_per_token(spec) * 64 + # one staging window of one layer's K/V, and a float32 fold accumulator with its two lse rows + window = slab // spec.num_layers * 128 + fold = PREFILL_ROWS * QO_HEADS * (64 + 2) * 4 + assert host_fixed - device_fixed == window + fold + pool = _pool(placement="host") + workspace = (*pool.staging_window(), *pool.fold_workspace(PREFILL_ROWS)) + assert host_fixed - device_fixed == sum(t.numel() * t.element_size() for t in workspace) + + +def test_only_the_sparse_pool_takes_host_placement(): + from freetoken.kvcache import check_kv_placement + + sparse = SimpleNamespace(model_is_mrope=False, dsv4_args=None) + sparse.kv_cache_group_specs = lambda: (_spec(),) + dense = SimpleNamespace(model_is_mrope=False, dsv4_args=None) + dense.kv_cache_group_specs = lambda: (_spec(attn_type=AttnType.FULL, index_ratio=1),) + check_kv_placement(sparse, "host") + check_kv_placement(dense, "device") + with pytest.raises(ValueError, match="QSA sparse pool only"): + check_kv_placement(dense, "host") diff --git a/tests/models/qwen4_exp/common.py b/tests/models/qwen4_exp/common.py index 0a52fb0e01..0480f8c9fb 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, + placement: str = "device", ) -> None: from freetoken.attention.qsa_sparse import QSASparseAttnBackend from freetoken.kvcache import create_kvcache_pool @@ -170,6 +171,8 @@ def __init__( dtype=dtype, device=self.device, num_req_slots=self.num_req_slots, + placement=placement, + prefill_rows=8192, ) 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 236610791d..e630e87832 100644 --- a/tests/models/qwen4_exp/test_qsa_backend.py +++ b/tests/models/qwen4_exp/test_qsa_backend.py @@ -249,3 +249,110 @@ def test_two_qsa_layers_keep_separate_slab_slots(monkeypatch): slab = fixture.pool.cmp_k_cache assert not torch.equal(slab(0), slab(1)) + + +# ------------------------------------------------------------------ K/V in host memory + + +def test_staging_windows_follow_each_request_in_logical_order(): + from freetoken.attention.qsa_sparse import plan_staging_windows + + # two requests: 3 cached pages at physical 7, 2, 9 and 2 cached pages at 4, 5 + block_table = torch.tensor([[7, 2, 9, 0], [4, 5, 0, 0]], dtype=torch.int32) + whole = plan_staging_windows(block_table, [3, 2], capacity=8) + assert len(whole) == 1 + assert whole[0].source_pages.tolist() == [7, 2, 9, 4, 5] + assert whole[0].block_table.tolist() == [[0, 1, 2, -1], [3, 4, -1, -1]] + + split = plan_staging_windows(block_table, [3, 2], capacity=2) + assert [w.source_pages.tolist() for w in split] == [[7, 2], [9, 4], [5]] + assert [w.block_table.tolist() for w in split] == [ + [[0, 1, -1, -1], [-1, -1, -1, -1]], + [[-1, -1, 0, -1], [1, -1, -1, -1]], + [[-1, -1, -1, -1], [-1, 0, -1, -1]], + ] + # every cached page is staged exactly once across the windows + staged = sorted(p for w in split for p in w.source_pages.tolist()) + assert staged == sorted([7, 2, 9, 4, 5]) + + +def test_a_prefill_stages_only_when_its_rows_select_more_than_is_cached(): + from freetoken.attention.qsa_sparse import staging_moves_less + + select_width = 2051 + assert staging_moves_less(rows=8192, select_width=select_width, cached_tokens=1 << 20) + assert not staging_moves_less(rows=20, select_width=select_width, cached_tokens=1 << 16) + assert not staging_moves_less(rows=32, select_width=64, cached_tokens=32 * 64) + assert staging_moves_less(rows=33, select_width=64, cached_tokens=32 * 64) + + +@requires_cuda +def test_a_prefill_folded_over_windows_allocates_nothing_one_window_does_not(monkeypatch): + """The startup fit measures one staging window; a prefill folding several must not reach + past that peak, or the first long prompt takes the memory the fit kept free.""" + import freetoken.kvcache.qsa_pool as qsa_pool + from freetoken.attention.qsa_sparse import QSASparseAttnBackend + + original = QSASparseAttnBackend._attend + peaks: dict[int, int] = {} + + def measured(self, q, layer_id, indices, md): + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + out = original(self, q, layer_id, indices, md) + torch.cuda.synchronize() + peaks[len(md.staging)] = torch.cuda.max_memory_allocated() - before + return out + + monkeypatch.setattr(QSASparseAttnBackend, "_attend", measured) + config = parsed_config() + length = 5000 + for window_tokens in (1 << 17, 128): + monkeypatch.setattr(qsa_pool, "STAGING_TOKENS", window_tokens) + fixture = Fixture(config, num_pages=128, placement="host") + attn = fixture.layer(QSA_LAYER) + x = _inputs(fixture, [length])[0] + attn.forward(x, fixture.batch([fixture.req(0, 0, length)], "prefill")) + del fixture, attn, x + assert sorted(peaks) == [1, 40] + output_bytes = length * config.num_qo_heads * config.head_dim * 2 + assert peaks[1] >= output_bytes, "the measured step must include the attend's own output" + assert peaks[40] <= peaks[1] + (1 << 20), ( + f"folding 40 windows peaked at {peaks[40] >> 20} MiB, one window at {peaks[1] >> 20} MiB" + ) + + +@requires_cuda +@pytest.mark.parametrize("window_tokens", [1 << 17, 128], ids=["one-window", "folded-windows"]) +def test_host_placement_matches_device_placement(monkeypatch, window_tokens: int): + """K/V in host memory: a prefill staged into device windows -- folded by log-sum-exp when the + context spans several -- and a decode read in place both equal the device-placed pool.""" + import freetoken.kvcache.qsa_pool as qsa_pool + + monkeypatch.setattr(qsa_pool, "STAGING_TOKENS", window_tokens) + config = parsed_config() + length, steps = 5000, 3 + outputs = {} + for placement in ("device", "host"): + fixture = Fixture(config, num_pages=128, placement=placement) + attn = fixture.layer(QSA_LAYER) + x = _inputs(fixture, [length], extra=steps)[0] + req = fixture.req(0, 0, length) + prefill = fixture.batch([req], "prefill") + got = [attn.forward(x[:length], prefill)] + windows = prefill.attn_metadata.staging + if placement == "host": + pages, window_pages = -(-length // 64), window_tokens // 64 + assert windows is not None and len(windows) == -(-pages // window_pages) + else: + assert windows is None + for step in range(steps): + fixture.step(req) + got.append(attn.forward(x[length + step : length + step + 1], fixture.batch([req], "decode"))) + outputs[placement] = got + for step, (device_out, host_out) in enumerate(zip(outputs["device"], outputs["host"])): + if window_tokens >= length or step > 0: + assert torch.equal(host_out, device_out), f"output {step} differs" + else: + torch.testing.assert_close(host_out.float(), device_out.float(), rtol=1e-2, atol=1e-2) diff --git a/tests/models/qwen4_exp/test_qsa_kernels.py b/tests/models/qwen4_exp/test_qsa_kernels.py index 6715519f05..5d86f2c645 100644 --- a/tests/models/qwen4_exp/test_qsa_kernels.py +++ b/tests/models/qwen4_exp/test_qsa_kernels.py @@ -501,3 +501,110 @@ def test_capture_graph_provisions_the_block_topk_scratch(): assert backend._scratch("topk_scratch", 2, width, dtype=torch.int32).data_ptr() == ( static.data_ptr() ) + + +# -------------------------------------------------------------------------------------- +# Passes over disjoint pages of one selection (a host-placed prefill's staging windows) +# -------------------------------------------------------------------------------------- + + +def _attention_pass(scores: torch.Tensor, values: torch.Tensor): + """Softmax attention over log2-domain scores [rows, heads, n] and values [rows, n, dim]: + the normalized output and each head's log2-sum-exp, as the attend kernel reports them.""" + lse = torch.logsumexp(scores * math.log(2.0), dim=-1) / math.log(2.0) + weights = torch.exp2(scores - lse.unsqueeze(-1)).nan_to_num(0.0) + return torch.einsum("rhn,rnd->rhd", weights, values), lse + + +@requires_cuda +def test_folding_passes_equals_one_pass_over_their_union(): + from freetoken.kernel.triton.qsa import fold_attention_pass + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(0) + rows, heads, tokens, dim = 5, 3, 40, 256 + scores = torch.randn(rows, heads, tokens, device=device, generator=generator) * 4 + values = torch.randn(rows, tokens, dim, device=device, generator=generator) + # row 0 attends nothing in any pass; row 1 only in the last one + scores[0] = -float("inf") + scores[1, :, :31] = -float("inf") + whole, whole_lse = _attention_pass(scores, values) + + total = torch.zeros(rows, heads, dim, device=device) + total_lse = torch.full((rows, heads), -float("inf"), device=device) + for part in (slice(0, 7), slice(7, 31), slice(31, 40)): + out, lse = _attention_pass(scores[..., part], values[:, part]) + fold_attention_pass(total, total_lse, out.bfloat16(), lse) + + assert not total.isnan().any() + torch.testing.assert_close(total, whole, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(total_lse, whole_lse, rtol=1e-6, atol=1e-6) + assert torch.equal(total[0], torch.zeros(heads, dim, device=device)) + assert (total_lse[0] == -float("inf")).all() + + +@requires_cuda +def test_the_last_pass_folds_into_its_own_output(): + """The final fold writes the sum in the output's dtype over the pass it reads, and leaves the + accumulator alone: the prefill needs no float32 result and no copy of one.""" + from freetoken.kernel.triton.qsa import fold_attention_pass + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(1) + rows, heads, tokens, dim = 7, 4, 64, 256 + scores = torch.randn(rows, heads, tokens, device=device, generator=generator) * 4 + values = torch.randn(rows, tokens, dim, device=device, generator=generator) + scores[2] = -float("inf") + whole, _ = _attention_pass(scores, values) + + total = torch.zeros(rows, heads, dim, device=device) + total_lse = torch.full((rows, heads), -float("inf"), device=device) + first, first_lse = _attention_pass(scores[..., :20], values[:, :20]) + fold_attention_pass(total, total_lse, first.bfloat16(), first_lse) + kept, kept_lse = total.clone(), total_lse.clone() + last, last_lse = _attention_pass(scores[..., 20:], values[:, 20:]) + out = last.bfloat16() + fold_attention_pass(total, total_lse, out, last_lse, result=out) + + torch.testing.assert_close(out.float(), whole, rtol=2e-2, atol=2e-2) + assert torch.equal(out[2], torch.zeros(heads, dim, dtype=out.dtype, device=device)) + assert torch.equal(total, kept) and torch.equal(total_lse, kept_lse) + + +@requires_cuda +@pytest.mark.parametrize("rows", [3, 300], ids=["split-merge", "one-split"]) +def test_attend_reports_the_log2_sum_exp_its_passes_fold_by(rows: int): + """The kernel's lse output makes two passes over disjoint pages fold into the pass over all + of them, on both its split-and-merge path (few rows) and its unsplit path (many).""" + from freetoken.kernel.triton.qsa import fold_attention_pass, qsa_sparse_paged_attention + + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(5) + pages, kv_heads, q_heads, dim, width = 8, 2, 4, 256, 300 + k_cache = torch.randn(pages, PAGE_SIZE, kv_heads, dim, device=device, generator=generator).bfloat16() + v_cache = torch.randn(pages, PAGE_SIZE, kv_heads, dim, device=device, generator=generator).bfloat16() + q = torch.randn(rows, q_heads, dim, device=device, generator=generator).bfloat16() + indices = torch.randint(0, pages * PAGE_SIZE, (rows, width), device=device, generator=generator).to(torch.int32) + block_table = torch.arange(pages, dtype=torch.int32, device=device).unsqueeze(0) + token_to_req = torch.zeros(rows, dtype=torch.int32, device=device) + + whole = qsa_sparse_paged_attention(q, k_cache, v_cache, indices, block_table, token_to_req) + whole_lse = torch.empty(rows, q_heads, dtype=torch.float32, device=device) + again = qsa_sparse_paged_attention(q, k_cache, v_cache, indices, block_table, token_to_req, lse=whole_lse) + assert torch.equal(again, whole) + + keys = k_cache.view(-1, kv_heads, dim)[indices.long()].float() # [rows, width, kv, dim] + group = q_heads // kv_heads + scores = torch.einsum("rhd,rnhd->rhn", q.float(), keys.repeat_interleave(group, dim=2)) + reference_lse = torch.logsumexp(scores * dim**-0.5, dim=-1) / math.log(2.0) + torch.testing.assert_close(whole_lse, reference_lse, rtol=1e-3, atol=1e-2) + + total = torch.zeros(q.shape, dtype=torch.float32, device=device) + total_lse = torch.full((rows, q_heads), -float("inf"), device=device) + part_lse = torch.empty_like(total_lse) + for keep in (torch.arange(pages) < 3, torch.arange(pages) >= 3): + table = torch.where(keep.to(device), block_table, -1).to(torch.int32) + part = qsa_sparse_paged_attention(q, k_cache, v_cache, indices, table, token_to_req, lse=part_lse) + fold_attention_pass(total, total_lse, part, part_lse) + torch.testing.assert_close(total, whole.float(), rtol=2e-2, atol=2e-2) + torch.testing.assert_close(total_lse, whole_lse, rtol=1e-4, atol=1e-3) From 5eb704032ec2deb79bc2093c438d4754986e22a5 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Thu, 1 Oct 2026 10:33:08 +0300 Subject: [PATCH 2/4] refactor(qsa): count a request's cached pages with div_ceil --- python/freetoken/attention/qsa_sparse.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/freetoken/attention/qsa_sparse.py b/python/freetoken/attention/qsa_sparse.py index 865f02bb59..88c9aeadfe 100644 --- a/python/freetoken/attention/qsa_sparse.py +++ b/python/freetoken/attention/qsa_sparse.py @@ -39,7 +39,7 @@ import torch from freetoken.core import Batch, get_global_ctx -from freetoken.utils import init_logger +from freetoken.utils import div_ceil, init_logger from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata @@ -346,7 +346,7 @@ def _plan_staging(self, md: QSASparseMetadata, rows: int) -> list[StagingWindow] """The windows a host-placed prefill stages its pages into, or None to read them in place.""" if self.kvcache.placement == "device" or md.is_decode: return None - page_counts = [-(-length // self.page_size) for length in md.kv_len_cpu.tolist()] + page_counts = [div_ceil(length, self.page_size) for length in md.kv_len_cpu.tolist()] if not staging_moves_less(rows, self.select_width, sum(page_counts) * self.page_size): return None capacity = self.kvcache.staging_window()[0].shape[0] From 2812257f1c2424288342205333a19f892379de08 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 13 Oct 2026 12:51:18 +0300 Subject: [PATCH 3/4] docs(kvcache): describe --kv-placement in the cli reference and the Flash-Next notes --- docs/cli.md | 1 + docs/models.md | 5 ++++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/cli.md b/docs/cli.md index 2166e9e86e..f4f856b8dc 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -71,6 +71,7 @@ ft serve --model ... --gpu GPU-9e8d7c6b # the same card by UUID (a unique prefi |---|---|---| | `--memory-ratio` | 0.9 | Fraction of free VRAM the engine may use (weights + MoE cache + KV) | | `--num-pages` / `--num-tokens` | auto | KV capacity override in pages / tokens (mutually exclusive; auto sizes from VRAM left after weights and MoE cache) | +| `--kv-placement` | device | `host` keeps the paged K/V in pinned host memory the GPU reads over PCIe, so the context is bounded by host RAM instead of VRAM; served by the QSA sparse pool only (Qwen3.8-Flash-Next) | | `--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` | | `--attention-backend`, `--attn` | auto | `trtllm`/`fi`/`fa`/`triton`/`dsv4_sparse`/`dsa`; `prefill,decode` pair allowed; auto picks per model + GPU | diff --git a/docs/models.md b/docs/models.md index cdadc2dcb1..d09da25159 100644 --- a/docs/models.md +++ b/docs/models.md @@ -119,4 +119,7 @@ unset on PCIe GPUs: wider grids add no bandwidth there. template kwargs with thinking on, `"chat_template_kwargs": {"enable_thinking": true, "reasoning_effort": 37}`, and the checkpoint's encoder validates it. Runtime window-cache controls use 128-token pages. -- Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM. +- Qwen3.8-Flash-Next keeps a 47.7 GiB PLE n-gram table pinned in host RAM. With `--kv-placement host` + its paged K/V lives there too, so the context is bounded by host RAM rather than VRAM: decode reads + the tokens its indexer selects over PCIe, and a prefill stages each layer's cached pages into a + fixed device window. From f5fa784b8cff778eeee43dca90d4eada0ec9b5c3 Mon Sep 17 00:00:00 2001 From: Yevhenii Date: Tue, 13 Oct 2026 12:52:10 +0300 Subject: [PATCH 4/4] docs(kvcache): complete a host pool's fixed device costs and the placement check's reason --- python/freetoken/engine/engine.py | 4 ++-- python/freetoken/kvcache/qsa_pool.py | 12 +++++++----- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 3ecaeeccc9..1e4e2ac878 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -1535,8 +1535,8 @@ def override(attr: str, value: Any): # this is dangerous, use with caution is_moe = getattr(model_config, "is_moe", False) expert_quant = getattr(model_config, "expert_quant", "none") - # Before any weight is read: a placement the pool cannot serve would otherwise fail - # minutes later, after the expert banks are loaded. + # Before any weight is read: a placement the pool cannot serve would otherwise fail only + # after the expert banks are loaded. check_kv_placement(model_config, config.kv_placement) if not is_moe: diff --git a/python/freetoken/kvcache/qsa_pool.py b/python/freetoken/kvcache/qsa_pool.py index 9498888058..db857c175a 100644 --- a/python/freetoken/kvcache/qsa_pool.py +++ b/python/freetoken/kvcache/qsa_pool.py @@ -19,11 +19,13 @@ The slab is amortized into the per-token KV price (``unit_bytes``); the ring and scratch are fixed and priced through ``kv_cost``'s ``fixed_cache_size``. -Placed on the host (``--kv-placement host``), the K/V slabs live in pinned host memory and only -the index tiers stay on the device, so the context is bounded by host RAM. Decode reads its +Placed on the host (``--kv-placement host``), the K/V slabs live in pinned host memory and the +index tiers stay on the device, so the context is bounded by host RAM. Decode reads its selection in place over PCIe. A prefill cannot: every query row reads its own selection, so a chunk would move each cached token over the bus hundreds of times. It stages the batch's pages -into a device window once per layer instead (``staging_window``), a fixed device cost. +into a device window once per layer instead (``staging_window``) and folds the passes of a +context wider than one window in a float32 accumulator (``fold_workspace``): both are fixed +device costs. """ from __future__ import annotations @@ -239,8 +241,8 @@ def kv_cost(cls, config) -> tuple[int, int, int, int]: return per_token * config.page_size, fixed, config.page_size, 0 def unit_bytes(self) -> tuple[int, int]: - # Only the shadow slab scales with pages, and only its non-scratch rows; the ring, the - # scratch rows and a host pool's staging window are the fixed term kv_cost reports. + # Only the shadow slab scales with pages, and only its non-scratch rows; the ring, the scratch + # rows and a host pool's staging window and fold accumulator are the fixed term kv_cost reports. kv, swa = super().unit_bytes() tokens = int(self._kv_buffer.shape[2]) * int(self._kv_buffer.shape[3]) slab = (