fix(dsv4): honor prefill limits and derive SWA capacity

docs/cli.md

@@ -53,10 +53,19 @@ parsers all resolve automatically from the checkpoint and the GPU.
 |---|---|---|
 | `--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) |
+| `--swa-num-tokens` | derived/auto | Expert absolute window/SWA override; DSV4 normally derives the minimum from `--max-prefill-length`; an override must align to the SWA page size and cannot be smaller than the derived requirement |
 | `--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 |
 
+For DSV4, `--max-prefill-length` is the normal single sizing knob. After resolving the model's
+window page, concurrency, and radix-retained working set, startup derives the minimum SWA pool that
+can hold roughly twice that chunk before reclamation. `--swa-num-tokens` is the expert override
+(and startup counterpart of `ft ctl cache --swa`); it must be an exact multiple of the DSV4 window
+size and cannot be smaller than the derived requirement. Status and rebuild results always report
+requested, pool-cap, and effective prefill tokens, whether SWA was derived or explicit, and which
+limit is binding.
+
 ### MoE offload
 
 See [models.md](models.md#moe-backends) for what each backend does.
@@ -153,4 +162,3 @@ expert format + GPU name, so a profile from different hardware is ignored
 rather than misapplied. Selection flags: `--dtype`, `--model`, `--formats`,
 `--isa`; decision rule: `--threshold` (default 2.0 — recommend hybrid when CPU
 bandwidth > 2× PCIe).
-

python/freetoken/cache_report.py

@@ -188,6 +188,16 @@ def format_cache_status(doc: dict, *, prefix: str = "cache: ") -> str:
     budget = _int(geometry, "cache_budget_bytes")
 
     header = f"{prefix}state={(doc or {}).get('state', 'serving')}"
+    requested_prefill = _int(geometry, "requested_prefill_tokens")
+    pool_prefill = _int(geometry, "pool_prefill_cap_tokens")
+    effective_prefill = _int(geometry, "effective_prefill_tokens")
+    if requested_prefill or effective_prefill:
+        header += (
+            f", prefill={effective_prefill} tok"
+            f" (requested {requested_prefill}, pool cap {pool_prefill or 'none'}, "
+            f"source {geometry.get('swa_capacity_source', 'none')}, "
+            f"reason {geometry.get('prefill_limiting_reason', 'none')})"
+        )
     if known:
         header += f", {format_bytes(sum(known))} allocated"
         if budget > 0:

python/freetoken/control_cli.py

@@ -29,6 +29,18 @@ def __init__(self, message: str, *, exit_code: int = 1) -> None:
         self.exit_code = exit_code
 
 
+def _prefill_summary(doc: dict[str, Any]) -> str:
+    requested = doc.get("requested_prefill_tokens", "unknown")
+    pool_cap = doc.get("pool_prefill_cap_tokens", "unknown")
+    effective = doc.get("effective_prefill_tokens", "unknown")
+    source = doc.get("swa_capacity_source", "unknown")
+    reason = doc.get("prefill_limiting_reason", "unknown")
+    return (
+        f"prefill requested={requested} pool_cap={pool_cap} effective={effective} "
+        f"source={source} reason={reason}"
+    )
+
+
 def parse_count(token: str) -> int:
     match = re.fullmatch(r"\s*(\d+(?:\.\d+)?)\s*([kKmM]?)\s*", token)
     if not match:
@@ -59,7 +71,11 @@ def _decode_error_body(raw: bytes) -> str:
         for key in ("error", "detail", "message"):
             value = doc.get(key)
             if value:
-                return str(value)
+                suffix = (
+                    f" ({_prefill_summary(doc)})"
+                    if "effective_prefill_tokens" in doc else ""
+                )
+                return f"{value}{suffix}"
         status = doc.get("status")
         if status:
             return str(status)
@@ -252,6 +268,7 @@ def _format_rebuild(doc: dict[str, Any]) -> str:
         f"kv={doc.get('num_pages', 'unknown')}",
         f"mamba={doc.get('mamba_slots', 'unknown')}",
         f"swa={doc.get('num_swa_pages', 'unknown')}",
+        _prefill_summary(doc),
     ]
     if doc.get("error"):
         parts.append(f"error={doc['error']}")

python/freetoken/engine/config.py

@@ -65,13 +65,20 @@ class EngineConfig:
     # Window/full ratio for the SWA radix cache (`--cache-type radix` on SWA models) and the DSV4
     # window tier: the DEFAULT window-pool size = max(working-set floor, ratio x full-pool tokens).
     # < 1.0 trades retained window-prefix capacity for memory savings; must be in (0, 1]. It is the
-    # DSV4 window/full ratio directly. Used only when swa_num_pages_override is None (a runtime
-    # rebuild can pin an absolute window instead).
+    # DSV4 fallback is superseded by its max-prefill-derived absolute window; generic radix-SWA
+    # still uses the ratio until an explicit startup/live override pins an absolute window.
     swa_full_tokens_ratio: float = 0.2
     # Absolute window-pool size in the pool's own pages (usable, dummy excluded); None -> use the
     # ratio default above. A runtime cache rebuild sets this (num_swa_pages) to pin the window
     # regardless of the full anchor; the ratio is the startup default and the fallback.
     swa_num_pages_override: int | None = None
+    # User-facing absolute window-pool capacity in tokens. Config resolution converts this to
+    # swa_num_pages_override after the model-specific SWA page unit is known. Kept separately
+    # until then so DSV4's 128-token window pages are validated rather than silently rounded.
+    swa_num_token_override: int | None = None
+    # Current sizing source vocabulary: derived | explicit | none. DSV4 derives from max prefill;
+    # generic radix-SWA derives from its ratio; expert startup/live overrides are explicit.
+    swa_capacity_source: str = "none"
     distributed_timeout: float = 60.0
     use_dummy_weight: bool = False
     use_pynccl: bool = True

python/freetoken/engine/engine.py

@@ -823,6 +823,7 @@ def rebuild_runtime_cache(
         # FrozenInstanceError, which here aborts the rebuild after the CUDA graphs are gone (→ 503).
         if num_swa_pages is not None:
             object.__setattr__(config, "swa_num_pages_override", num_swa_pages)
+            object.__setattr__(config, "swa_capacity_source", "explicit")
         if moe_cache_size is not None:
             assert self.moe_offload_cache is not None, "no MoE offload cache to resize"
             self.moe_offload_cache.rebuild(moe_cache_size)
@@ -983,8 +984,7 @@ def _resolve_cache_type(has_linear_attention: bool, requested: str) -> str:
 def _adjust_dsv4_config(config: EngineConfig, override) -> None:
     """DSV4 engine-config reconciliation at config-resolution time (before the pool exists).
     Syncs the resolved runtime config into the opaque ``dsv4_args`` payload, sets
-    page_size to the window page P, forces single-chunk prefill, and clamps cuda_graph_bs/max_bs to
-    the DSV4 decode batch size.
+    page_size to the window page P, and clamps cuda_graph_bs/max_bs to the DSV4 decode batch size.
     """
     model_config = config.model_config
     model_config.dsv4_args.max_seq_len = config.max_seq_len
@@ -1001,11 +1001,8 @@ def _adjust_dsv4_config(config: EngineConfig, override) -> None:
     if getattr(config, "cache_type", "radix") != "naive":
         override("cache_type", "swa_radix")
     # 'radix' (SWARadixCache on the full-loc currency, carry-aware re-prefill) is the default and is
-    # honored, as is an explicit 'naive'. Don't let max_extend_tokens force a second chunk within
-    # one prompt (the pool's prefill_chunk_budget still chunks prompts larger than the window
-    # pool); prefill batches ragged (bs>=1), each segment resuming from its own cached_len.
-    if getattr(config, "max_extend_tokens", 0) < config.max_seq_len:
-        override("max_extend_tokens", config.max_seq_len)
+    # honored, as is an explicit 'naive'. max_extend_tokens is a user-visible scheduler bound and
+    # remains authoritative; the pool's prefill_chunk_budget may narrow it further.
 
     # DSV4 decode batches at most max_running_req rows; its full-loc snapshot is sized to that,
     # so a graph bs above it would exceed the backend's captured snapshot rows. Clamp any
@@ -1366,6 +1363,67 @@ def override(attr: str, value: Any):  # this is dangerous, use with caution
             )
         override("num_page_override", config.num_token_override // config.page_size)
 
+    # Resolve window capacity only after model-specific page-size/cache reconciliation. DSV4's
+    # ordinary one-knob path derives the minimum window that can honor max_extend_tokens; an
+    # explicit token/page capacity is an expert override and must be at least that large.
+    swa_tokens = getattr(config, "swa_num_token_override", None)
+    swa_pages = getattr(config, "swa_num_pages_override", None)
+    explicit_swa = swa_tokens is not None or swa_pages is not None
+    if explicit_swa:
+        supports_absolute_swa = is_dsv4 or (
+            has_swa_attention and getattr(config, "cache_type", "radix") == "swa_radix"
+        )
+        if not supports_absolute_swa:
+            raise ValueError(
+                "absolute SWA capacity requires DSV4 or a sliding-window model with "
+                "--cache-type radix"
+            )
+        if swa_tokens is not None and swa_pages is not None:
+            raise ValueError(
+                "swa_num_token_override and swa_num_pages_override are mutually exclusive"
+            )
+        if swa_pages is not None and swa_pages <= 0:
+            raise ValueError(f"swa_num_pages_override must be positive, got {swa_pages}")
+        if swa_tokens is not None:
+            if swa_tokens <= 0:
+                raise ValueError(f"swa_num_token_override must be positive, got {swa_tokens}")
+            swa_page_size = model_config.dsv4_args.window_size if is_dsv4 else 1
+            if swa_tokens % swa_page_size != 0:
+                lower = swa_tokens // swa_page_size * swa_page_size
+                upper = (swa_tokens // swa_page_size + 1) * swa_page_size
+                raise ValueError(
+                    f"--swa-num-tokens {swa_tokens} is not a multiple of the resolved SWA "
+                    f"page size {swa_page_size}; nearest valid values: {lower} or {upper}"
+                )
+            swa_pages = swa_tokens // swa_page_size
+            override("swa_num_pages_override", swa_pages)
+        override("swa_capacity_source", "explicit")
+
+    if is_dsv4:
+        from freetoken.kvcache.dsv4_cost_model import (
+            dsv4_required_swa_pages,
+        )
+
+        P = model_config.dsv4_args.window_size
+        radix = config.cache_type != "naive"
+        requested_prefill = int(getattr(config, "max_extend_tokens", 8192))
+        required_pages = dsv4_required_swa_pages(
+            requested_prefill, config.max_running_req, radix, P
+        )
+        if swa_pages is None:
+            swa_pages = required_pages
+            override("swa_num_pages_override", swa_pages)
+            override("swa_capacity_source", "derived")
+        else:
+            if int(swa_pages) < required_pages:
+                raise ValueError(
+                    f"--swa-num-tokens {int(swa_pages) * P} is below the minimum "
+                    f"{required_pages * P} SWA tokens required by --max-prefill-length "
+                    f"{requested_prefill} with max_running_req={config.max_running_req}"
+                )
+    elif has_swa_attention and getattr(config, "cache_type", None) == "swa_radix" and not explicit_swa:
+        override("swa_capacity_source", "derived")
+
     # The rope cos/sin table is baked to rotary_config.max_position, and neither rope kernel
     # bounds-checks the position it gathers with -- a longer ceiling reads past the table.
     # DSV4 is exempt: it sizes its own table from the resolved max_seq_len (_adjust_dsv4_config).

python/freetoken/kvcache/cache_status.py

@@ -145,7 +145,7 @@ def _swa() -> int:
     return floors
 
 
-def compute_cache_pools(engine: "Engine") -> Dict[str, int]:
+def compute_cache_pools(engine: "Engine") -> Dict[str, Any]:
     """The ACTUAL pool sizes allocated at load, in API units. The frontend otherwise only
     learns them from per-generation UserReply snapshots — i.e. not until the first chat —
     so this seeds /v1/cache/status geometry with truth from the moment the server is ready.
@@ -154,12 +154,32 @@ def compute_cache_pools(engine: "Engine") -> Dict[str, int]:
     pools = {
         "num_pages": 0, "page_size": 0, "moe_cache_size": 0, "num_mamba_slots": 0,
         "swa_page_size": 0, "num_swa_pages": 0,
+        "requested_prefill_tokens": 0, "pool_prefill_cap_tokens": 0,
+        "effective_prefill_tokens": 0, "swa_capacity_source": "none",
+        "prefill_limiting_reason": "none",
     }
     try:
         config = engine.config
         pools["num_pages"] = int(engine.num_pages or 0)
         if config is not None:
             pools["page_size"] = int(config.page_size or 0)
+            requested = int(getattr(config, "max_extend_tokens", 0) or 0)
+            pool_cap = int(
+                getattr(getattr(engine, "kv_cache", None), "prefill_chunk_budget", 0) or 0
+            )
+            mc = config.model_config
+            pools.update(prefill_geometry(
+                requested,
+                pool_cap,
+                getattr(config, "swa_capacity_source", "none"),
+                has_window_pool=bool(
+                    getattr(mc, "dsv4_args", None) is not None
+                    or (
+                        getattr(mc, "has_swa_attention", False)
+                        and getattr(config, "cache_type", None) == "swa_radix"
+                    )
+                ),
+            ))
             # Window pool: its own page unit (swa_page_size -- DSV4 windows are P-token pages,
             # radix-SWA is token-granular page_size 1) and the concrete current size in that unit
             # (num_swa_pages, usable count). Same source as the scheduler's _current_cache_geometry.
@@ -184,6 +204,34 @@ def compute_cache_pools(engine: "Engine") -> Dict[str, int]:
     return pools
 
 
+def prefill_geometry(
+    requested_tokens: int,
+    pool_cap_tokens: int,
+    swa_capacity_source: str,
+    *,
+    has_window_pool: bool,
+) -> Dict[str, Any]:
+    """Canonical requested/pool/effective prefill readout for startup and live rebuilds."""
+    requested = max(0, int(requested_tokens or 0))
+    pool_cap = max(0, int(pool_cap_tokens or 0))
+    effective = min(requested, pool_cap) if pool_cap else requested
+    if requested <= 0:
+        reason = "none"
+    elif pool_cap and pool_cap < requested:
+        reason = "swa_pool"
+    elif pool_cap and pool_cap == requested:
+        reason = "requested_and_swa_pool"
+    else:
+        reason = "requested_limit"
+    return {
+        "requested_prefill_tokens": requested,
+        "pool_prefill_cap_tokens": pool_cap,
+        "effective_prefill_tokens": effective,
+        "swa_capacity_source": str(swa_capacity_source) if has_window_pool else "none",
+        "prefill_limiting_reason": reason,
+    }
+
+
 def compute_cache_status_meta(engine: "Engine") -> Dict[str, Any]:
     """Full readiness ("meta", …) payload for the desktop cache panel: the per-unit VRAM costs
     (compute_cache_unit_bytes) plus the post-weights pre-pool free-VRAM baseline (the sliders'
@@ -193,11 +241,21 @@ def compute_cache_status_meta(engine: "Engine") -> Dict[str, Any]:
     meta: Dict[str, Any] = dict(compute_cache_unit_bytes(engine))
     meta["free_vram_bytes"] = _pool_budget_free_vram_bytes(engine)
     meta["floors"] = compute_cache_floors(engine)
-    meta["pools"] = compute_cache_pools(engine)
-    # Current window/full reuse ratio (the tunable knob), for DSV4 and radix-SWA; 0.0 otherwise.
+    pools = meta["pools"] = compute_cache_pools(engine)
+    # Effective allocated window/full ratio, for DSV4 and radix-SWA; 0.0 otherwise. An absolute
+    # derived/explicit window supersedes the configured fallback ratio, so report live geometry.
     cfg = engine.config
     has_swa_ratio = cfg is not None and _supports_swa_ratio(cfg)
-    meta["swa_full_tokens_ratio"] = float(cfg.swa_full_tokens_ratio) if has_swa_ratio else 0.0
+    ratio = 0.0
+    if has_swa_ratio:
+        try:
+            is_dsv4 = getattr(cfg.model_config, "dsv4_args", None) is not None
+            full = int(pools["num_pages"]) if is_dsv4 else int(pools["num_pages"]) * int(pools["page_size"])
+            window = int(pools["num_swa_pages"])
+            ratio = min(1.0, window / full) if full > 0 else float(cfg.swa_full_tokens_ratio)
+        except Exception:  # noqa: BLE001 -- best-effort readiness metadata
+            ratio = float(cfg.swa_full_tokens_ratio)
+    meta["swa_full_tokens_ratio"] = ratio
     # Exact total cache VRAM budget (all pools) the engine honors: memory_ratio of the
     # post-weights baseline minus weights — the same figure the rebuild fit-check uses, with
     # fixed=0 so it's the whole-cache ceiling (KV + MoE + Mamba + SWA), not the KV+MoE remainder.