feat(kvcache): --kv-placement host keeps the QSA K/V in pinned host memory

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 |

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.

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,11 +35,11 @@
 
 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
-from freetoken.utils import init_logger
+from freetoken.utils import div_ceil, init_logger
 
 from .base import AttentionSpec, BaseAttnBackend, BaseAttnMetadata
 
@@ -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 = [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]
+        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",
+]

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)
 

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 only
+    # 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

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"<i{host.element_size()}",
+            "data": (device_ptr(host), False),
+            "version": 3,
+        }
+
+
+def mapped_cuda_view(host: torch.Tensor, device: torch.device) -> 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.
 

python/freetoken/kernel/triton/qsa/__init__.py

@@ -1,13 +1,14 @@
 """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
 from .topk import qsa_block_topk, qsa_block_topk_scratch_width
 
 __all__ = [
     "expand_qsa_block_indices",
+    "fold_attention_pass",
     "qsa_block_topk",
     "qsa_block_topk_scratch_width",
     "qsa_compress_groups",