diff --git a/docs/cli.md b/docs/cli.md index cf4b27a2..26c99047 100644 --- a/docs/cli.md +++ b/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). - diff --git a/python/freetoken/cache_report.py b/python/freetoken/cache_report.py index cffe76b8..701607d9 100644 --- a/python/freetoken/cache_report.py +++ b/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: diff --git a/python/freetoken/control_cli.py b/python/freetoken/control_cli.py index 50a32aae..6ac0d305 100644 --- a/python/freetoken/control_cli.py +++ b/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']}") diff --git a/python/freetoken/engine/config.py b/python/freetoken/engine/config.py index 543012f3..012b623b 100644 --- a/python/freetoken/engine/config.py +++ b/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 diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index ff3c985f..5dafaf3b 100644 --- a/python/freetoken/engine/engine.py +++ b/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). diff --git a/python/freetoken/kvcache/cache_status.py b/python/freetoken/kvcache/cache_status.py index 10169d65..7e1a3aba 100644 --- a/python/freetoken/kvcache/cache_status.py +++ b/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. diff --git a/python/freetoken/kvcache/dsv4_cost_model.py b/python/freetoken/kvcache/dsv4_cost_model.py index 2c084aff..218869f7 100644 --- a/python/freetoken/kvcache/dsv4_cost_model.py +++ b/python/freetoken/kvcache/dsv4_cost_model.py @@ -38,6 +38,26 @@ def dsv4_reserved_window_pages(max_running_req: int, radix: bool) -> int: return 2 * (max_running_req + 1) + (3 * max_running_req if radix else 0) + 1 +def dsv4_prefill_chunk_budget( + num_swa_pages: int, max_running_req: int, radix: bool, P: int = 128 +) -> int: + """Effective DSV4 prefill cap from a USABLE window-page count. + + A prefill forward peaks at about twice the incoming chunk before between-chunk reclamation; + the concurrent decode/radix working set is unavailable to that peak. + """ + reserved = dsv4_reserved_window_pages(max_running_req, radix) + return max(P, (int(num_swa_pages) - reserved) // 2 * P) + + +def dsv4_required_swa_pages( + prefill_tokens: int, max_running_req: int, radix: bool, P: int = 128 +) -> int: + """Minimum USABLE SWA pages needed to honor a requested prefill chunk ceiling.""" + chunk_pages = -(-int(prefill_tokens) // P) + return dsv4_reserved_window_pages(max_running_req, radix) + 2 * chunk_pages + + def ring_size_for_ratio(ratio: int) -> int: """Compress-state ring slots per window page (non-speculative).""" if ratio == 4: @@ -235,20 +255,31 @@ def dsv4_solve_num_pages( floor_win_pages: int, P: int = 128, n_scratch: int = 1, + num_swa_pages: int | None = None, ) -> DSV4PoolSizes: """Largest budget-respecting pool: max ``num_pages`` with exact - ``dsv4_pool_bytes(sizes(num_pages, win=max(floor, ceil(r*num)))) <= available_bytes``. + ``dsv4_pool_bytes(sizes(num_pages, win)) <= available_bytes``. - The window floor is honored in PAGES (never by inflating ``swa_ratio``) and the total is - byte-checked: at small budgets the window pins at ``floor_win_pages`` and the full/cmp/idx - anchor SHRINKS to fit, instead of every tier inflating past the budget. Raises ``ValueError`` - when even the minimal pool does not fit. + With no absolute target, ``win=max(floor, ceil(ratio*num_pages))``. ``num_swa_pages`` is an + absolute USABLE-page pin, so its physical window is ``target + 1`` for the dummy page and no + longer scales with the full-history anchor. The total is byte-checked in either case. Raises + ``ValueError`` when even the minimal geometry does not fit. """ + pinned_win_pages = ( + max(floor_win_pages, int(num_swa_pages) + 1) + if num_swa_pages is not None else None + ) + def _sizes(num: int) -> DSV4PoolSizes: - win = max(floor_win_pages, (round(swa_ratio * num * P) + P - 1) // P) + win = ( + pinned_win_pages + if pinned_win_pages is not None + else max(floor_win_pages, (round(swa_ratio * num * P) + P - 1) // P) + ) return dsv4_pool_sizes(num, args, swa_ratio, P=P, n_win_pages=win) - lo = max(floor_win_pages, 2) # full history must at least cover the window working set + # Full history must cover both the working-set floor and an absolute window pin. + lo = max(floor_win_pages, pinned_win_pages or 0, 2) if dsv4_pool_bytes(_sizes(lo), args, n_scratch) > available_bytes: raise ValueError( f"DSV4 KV budget {available_bytes} bytes cannot fit the minimal pool " @@ -270,13 +301,32 @@ def _sizes(num: int) -> DSV4PoolSizes: _AUTO_KV_SLACK_BYTES = 2 << 30 # absorbs plan-vs-measured drift (observed ~265MiB) and leaves a usable pool -def dsv4_auto_cost_model(args, swa_ratio, floor_win_pages, P=128, n_scratch=1): +def dsv4_auto_cost_model( + args, swa_ratio, floor_win_pages, P=128, n_scratch=1, num_swa_pages: int | None = None +): """Affine (cache_per_page, fixed_cache_size, min_reserve_tokens) for the MoE-first auto planner: exact marginal per-page cost across all tiers (+ the full_to_window mapping), a fixed intercept anchored at the minimal viable pool, and a reserve floor covering the window working set plus a slack absorbing plan-vs-measured drift. Conservative at the shipped swa_ratio; an extreme swa_ratio can dip <1% under exact (harmless -- num_pages is re-solved byte-exactly from - measured memory).""" + measured memory). + + An absolute ``num_swa_pages`` target is a fixed window tier, so only the full-history tiers are + marginal. That geometry is exactly affine in usable full pages; price its fixed window and + sentinel/scratch rows in the intercept instead of applying the ratio model. + """ + if num_swa_pages is not None: + win_pages = max(floor_win_pages, int(num_swa_pages) + 1) + per_page = dsv4_cache_per_page(args, 0.0, P) + P * _INT64_BYTES + n0 = max(win_pages, 2) # physical full pages, including the dummy + base = dsv4_pool_bytes( + dsv4_pool_sizes(n0, args, swa_ratio, P=P, n_win_pages=win_pages), args, n_scratch + ) + slack_pages = -(-_AUTO_KV_SLACK_BYTES // per_page) + usable0 = n0 - 1 + min_reserve_tokens = (usable0 + slack_pages) * P + return per_page, max(0, base - usable0 * per_page), min_reserve_tokens + per_page = dsv4_cache_per_page(args, swa_ratio, P) + P * _INT64_BYTES n0 = max(floor_win_pages, 2) win0 = max(floor_win_pages, (round(swa_ratio * n0 * P) + P - 1) // P) @@ -295,6 +345,8 @@ def dsv4_auto_cost_model(args, swa_ratio, floor_win_pages, P=128, n_scratch=1): "dsv4_kv_unit_bytes", "dsv4_pool_bytes", "dsv4_pool_sizes", + "dsv4_prefill_chunk_budget", + "dsv4_required_swa_pages", "dsv4_reserved_window_pages", "dsv4_solve_num_pages", "dsv4_window_unit_bytes", @@ -316,16 +368,17 @@ def _dsv4_swa_ratio(config) -> float: def _dsv4_window_floor_pages(config, P: int) -> int: """Minimum window pages = the live sliding working set the window pool must always hold: - one prefill chunk's reach (capped at 8 pages == 1024 tok; chunked prefill bounds the rest) - + each running request's 128-tail (<=2 pages mid-boundary, concurrency-scaled) + one - radix-locked live tail page per concurrent cached prompt + the reserved dummy. This is the - only HARD floor on DSV4 KV sizing -- the full-history (cmp/idx) capacity above it is purely - memory-derived, and a request longer than it is gracefully gated by ``available_size``.""" + twice one bounded prefill reach (both sides coexist before reclamation), the concurrent + decode/radix reserve, and one physical dummy page. This is the only HARD floor on DSV4 KV + sizing -- full-history capacity above it is memory-derived, and larger prefills are chunked.""" prefill_reach_pages = (config.max_seq_len + P - 1) // P # DSV4 config resolution rewrites cache_type to "swa_radix" (engine._adjust_dsv4_config), # so the radix term must key on "not naive" -- `== "radix"` was always False here. radix = config.cache_type != "naive" - return min(prefill_reach_pages, 8) + dsv4_reserved_window_pages(config.max_running_req, radix) + usable = dsv4_required_swa_pages( + min(prefill_reach_pages, 8) * P, config.max_running_req, radix, P + ) + return usable + 1 # physical dummy page def _dsv4_pool_sizes(config, num_pages: int, num_swa_pages: int | None = None): @@ -352,8 +405,18 @@ def _dsv4_pool_sizes(config, num_pages: int, num_swa_pages: int | None = None): else config.swa_num_pages_override ) if target is not None: - # Absolute window: `target` usable pages + 1 dummy, floored and capped at the full anchor. - win = min(num_pages, max(floor_pages, int(target) + 1)) + # Absolute window: `target` usable pages + 1 dummy. Unlike ratio sizing, an absolute pin + # must never be silently truncated to the full-history anchor; every startup/rebuild path + # reaches this shared geometry boundary before allocation. + target = int(target) + full_usable_pages = num_pages - 1 + if target > full_usable_pages: + raise ValueError( + f"full-history capacity {full_usable_pages} pages " + f"({full_usable_pages * P} tokens) is smaller than absolute DSV4 SWA " + f"capacity {target} pages ({target * P} tokens)" + ) + win = max(floor_pages, target + 1) else: # Default: ratio x full, rounded UP to whole window pages (floor applied ONCE, in pages). win = max(floor_pages, (round(swa_ratio * num_pages * P) + P - 1) // P) diff --git a/python/freetoken/kvcache/dsv4_paged_pool.py b/python/freetoken/kvcache/dsv4_paged_pool.py index 3ce23eea..83824825 100644 --- a/python/freetoken/kvcache/dsv4_paged_pool.py +++ b/python/freetoken/kvcache/dsv4_paged_pool.py @@ -322,7 +322,9 @@ def kv_cost(cls, config) -> tuple[int, int, int, int]: P = dsv4_args.window_size floor = _dsv4_window_floor_pages(config, P) per_page, fixed, min_reserve_tokens = dsv4_auto_cost_model( - dsv4_args, _dsv4_swa_ratio(config), floor, P=P, n_scratch=config.max_running_req + 1 + dsv4_args, _dsv4_swa_ratio(config), floor, P=P, + n_scratch=config.max_running_req + 1, + num_swa_pages=config.swa_num_pages_override, ) return per_page, fixed, config.page_size, min_reserve_tokens @@ -346,6 +348,7 @@ def solve_num_pages(cls, config, available_memory: int) -> int: available_memory, dsv4_args, _dsv4_swa_ratio(config), floor_win_pages=_dsv4_window_floor_pages(config, P), P=P, n_scratch=config.max_running_req + 1, + num_swa_pages=config.swa_num_pages_override, ) # The solver fits PHYSICAL pages to memory; one is the dummy page, so the # usable (advertised) count is one less. @@ -403,9 +406,12 @@ def validate_rebuild( # Size the pool a KV/window rebuild would build: the target anchor (or current) with # the target window (or current), computed BEFORE the config is mutated. target_pages = num_pages if num_pages is not None else current_num_pages - kv_sizes = _dsv4_pool_sizes( - config, target_pages + 1, num_swa_pages=num_swa_pages - ) # +1 for dummy page + try: + kv_sizes = _dsv4_pool_sizes( + config, target_pages + 1, num_swa_pages=num_swa_pages + ) # +1 for dummy page + except ValueError as exc: + raise CacheRebuildRejected(str(exc)) from exc else: # MoE-only rebuild keeps the CURRENT pool: budget-check against its live sizes # (reflects DSV4_FORCE_SMALL_POOL and the physical dummy page). @@ -463,7 +469,7 @@ def _init_paged_state(self, max_running_req: int, radix: bool) -> None: LAST window page are the reserved dummy region: page_table's dummy row points at ``full_token - P`` (== the generic ``fill_(num_tokens)`` convention with num_tokens = the allocatable token count), permanently bound so graph-padded rows scatter to a real slot.""" - from .dsv4_cost_model import dsv4_reserved_window_pages + from .dsv4_cost_model import dsv4_prefill_chunk_budget P = self.P self._paged_params = (int(max_running_req), bool(radix)) self.full_to_window.fill_(-1) @@ -473,8 +479,9 @@ def _init_paged_state(self, max_running_req: int, radix: bool) -> None: # only between chunks; peak ~2x the chunk), so reserve the concurrent working set and # halve the rest -- the same formula the bespoke manager used. n_win_pages = (self.sizes.n_win_slots // P) - 1 - reserved = dsv4_reserved_window_pages(max_running_req, radix) - self._chunk_budget = max(P, (n_win_pages - reserved) // 2 * P) + self._chunk_budget = dsv4_prefill_chunk_budget( + n_win_pages, max_running_req, radix, P + ) @property def swa_num_tokens(self) -> int: diff --git a/python/freetoken/message/frontend.py b/python/freetoken/message/frontend.py index 24725567..18de66bf 100644 --- a/python/freetoken/message/frontend.py +++ b/python/freetoken/message/frontend.py @@ -67,4 +67,9 @@ class CacheRebuildReply(BaseFrontendMsg): num_pages: int = 0 mamba_slots: int = 0 num_swa_pages: int = 0 + requested_prefill_tokens: int = 0 + pool_prefill_cap_tokens: int = 0 + effective_prefill_tokens: int = 0 + swa_capacity_source: str = "none" + prefill_limiting_reason: str = "none" error: str | None = None diff --git a/python/freetoken/message/tokenizer.py b/python/freetoken/message/tokenizer.py index 33b75c78..2afae1e4 100644 --- a/python/freetoken/message/tokenizer.py +++ b/python/freetoken/message/tokenizer.py @@ -99,6 +99,11 @@ class CacheRebuildResultMsg(BaseTokenizerMsg): num_pages: int = 0 mamba_slots: int = 0 num_swa_pages: int = 0 + requested_prefill_tokens: int = 0 + pool_prefill_cap_tokens: int = 0 + effective_prefill_tokens: int = 0 + swa_capacity_source: str = "none" + prefill_limiting_reason: str = "none" error: str | None = None diff --git a/python/freetoken/scheduler/cache.py b/python/freetoken/scheduler/cache.py index 9be235b2..f443dfae 100644 --- a/python/freetoken/scheduler/cache.py +++ b/python/freetoken/scheduler/cache.py @@ -50,11 +50,6 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ # lifecycle (alloc_swa / out-of-window free / free-on-finish). is_swa gates only the extra # SWARadixCache reuse machinery (tree match/insert/evict_swa/swa_uuid lock). self.swa_paged = swa_pool is not None and getattr(swa_pool, "swa_paged", False) - # Owned-pool capability pickup: a plugged-in swa pool may cap the prefill chunk (DSV4: - # ~half the window working set). Instance attrs shadow the class defaults; absent - # attributes leave the defaults untouched (Gemma4). - if swa_pool is not None: - self.prefill_chunk_budget = getattr(swa_pool, "prefill_chunk_budget", None) self.prefix_cache = self._make_prefix_cache(device, page_size, type) self.device = device self.num_pages = num_pages @@ -64,7 +59,17 @@ def __init__(self, num_pages: int, page_size: int, page_table: torch.Tensor, typ # ----- capability hooks (defaults; plugged-in pools may narrow them) ----- supports_runtime_rebuild = True - prefill_chunk_budget = None # generic shared page pool: no per-model prefill chunk cap + @property + def prefill_chunk_budget(self) -> int | None: + """Current pool-provided prefill cap, if any. + + Pools rebuild in place, so this must remain a live delegation rather than a snapshot taken + during CacheManager construction. Generic shared pools expose no cap and return None. + """ + return ( + getattr(self.swa_pool, "prefill_chunk_budget", None) + if self.swa_pool is not None else None + ) def page_usage(self) -> tuple[int, int]: """(used_pages, total_pages): allocated, non-evictable pages over the pool total diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 35541161..cdd1a53f 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -129,6 +129,28 @@ def __init__(self, config: SchedulerConfig): self.prefill_budget = ( min(config.max_extend_tokens, _chunk_cap) if _chunk_cap else config.max_extend_tokens ) + from freetoken.kvcache.cache_status import prefill_geometry + + prefill = prefill_geometry( + config.max_extend_tokens, + _chunk_cap or 0, + getattr(config, "swa_capacity_source", "none"), + has_window_pool=bool( + getattr(config.model_config, "dsv4_args", None) is not None + or ( + getattr(config.model_config, "has_swa_attention", False) + and getattr(config, "cache_type", None) == "swa_radix" + ) + ), + ) + logger.info_rank0( + "Prefill chunk budget: " + f"requested={prefill['requested_prefill_tokens']}, " + f"pool_cap={prefill['pool_prefill_cap_tokens']}, " + f"effective={prefill['effective_prefill_tokens']}, " + f"source={prefill['swa_capacity_source']}, " + f"reason={prefill['prefill_limiting_reason']}" + ) self.config = config self.status_reporter = SchedulerStatusReporter( log=logger.info_rank0, @@ -619,6 +641,11 @@ def _reply_rebuild(self, request_id: str, status: str, error: str | None = None) num_pages=geo["num_pages"], mamba_slots=geo["num_mamba_slots"] or 0, num_swa_pages=geo["num_swa_pages"] or 0, + requested_prefill_tokens=geo["requested_prefill_tokens"], + pool_prefill_cap_tokens=geo["pool_prefill_cap_tokens"], + effective_prefill_tokens=geo["effective_prefill_tokens"], + swa_capacity_source=geo["swa_capacity_source"], + prefill_limiting_reason=geo["prefill_limiting_reason"], error=error, ) ] @@ -641,6 +668,7 @@ def _execute_pending_rebuild(self) -> None: # the prefix cache that a successful resize of just the requested pool preserves. snapshot = self._current_cache_geometry() prior = {k: snapshot[k] for k, v in requested.items() if v is not None} + prior_swa_source = snapshot["swa_capacity_source"] # Cleared here, set by engine.rebuild_runtime_cache at its point of no return — lets the # except below tell a pre-teardown failure (engine untouched) from a mid-teardown one. self.engine.rebuild_teardown_started = False @@ -677,6 +705,8 @@ def _execute_pending_rebuild(self) -> None: logger.error(f"cache rebuild failed: {e!r} — rolling back to the previous geometry") try: self.rebuild_cache(**prior) + if "num_swa_pages" in prior: + object.__setattr__(self.config, "swa_capacity_source", prior_swa_source) except Exception as e2: # noqa: BLE001 — rollback failed too; genuinely unrecoverable logger.error(f"cache rebuild rollback failed: {e2!r} — server latched failed") self._reply_rebuild( @@ -714,11 +744,20 @@ def _current_cache_geometry(self) -> dict: getattr(config, "cache_type", None) == "swa_radix" ): # usable window tokens = pool tokens minus the slot-0 sentinel num_swa_pages = max(0, int(getattr(eng.kv_cache, "swa_num_tokens", 0) or 0) - 1) + from freetoken.kvcache.cache_status import prefill_geometry + + prefill = prefill_geometry( + config.max_extend_tokens, + int(getattr(eng.kv_cache, "prefill_chunk_budget", 0) or 0), + getattr(config, "swa_capacity_source", "none"), + has_window_pool=num_swa_pages is not None, + ) return dict( num_pages=eng.num_pages, moe_cache_size=eng.moe_offload_cache.cache_size if eng.moe_offload_cache is not None else None, num_mamba_slots=(eng.linear_state_pool.num_slots - 1) if eng.linear_state_pool is not None else None, num_swa_pages=num_swa_pages, + **prefill, ) def _log_cache_geometry(self, event: str) -> None: @@ -741,6 +780,13 @@ def _log_cache_geometry(self, event: str) -> None: f"swa {pools['num_swa_pages']} pages" f" ({swa_tokens} tokens, {_gib(swa_tokens * unit['swa_bytes_per_token'])})" ) + parts.append( + f"prefill {pools['effective_prefill_tokens']} tokens" + f" (requested {pools['requested_prefill_tokens']}, " + f"pool cap {pools['pool_prefill_cap_tokens'] or 'none'}, " + f"source {pools['swa_capacity_source']}, " + f"reason {pools['prefill_limiting_reason']})" + ) if pools["num_mamba_slots"]: parts.append( f"mamba {pools['num_mamba_slots']} slots" diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index 80dbf6ab..46e3c6eb 100644 --- a/python/freetoken/server/api_server.py +++ b/python/freetoken/server/api_server.py @@ -275,6 +275,11 @@ def _resolve_rebuild(self, msg: CacheRebuildReply) -> None: "num_pages": msg.num_pages, "mamba_slots": msg.mamba_slots, "num_swa_pages": msg.num_swa_pages, + "requested_prefill_tokens": msg.requested_prefill_tokens, + "pool_prefill_cap_tokens": msg.pool_prefill_cap_tokens, + "effective_prefill_tokens": msg.effective_prefill_tokens, + "swa_capacity_source": msg.swa_capacity_source, + "prefill_limiting_reason": msg.prefill_limiting_reason, "error": msg.error, } fut = self.rebuild_futures.pop(msg.request_id, None) @@ -495,6 +500,40 @@ class CacheRebuildRequest(BaseModel): timeout: float = 300.0 +def _rebuild_prefill_fields(state: Any) -> Dict[str, Any]: + """Best-effort current five-field prefill snapshot for every rebuild exit path.""" + keys = ( + "requested_prefill_tokens", "pool_prefill_cap_tokens", "effective_prefill_tokens", + "swa_capacity_source", "prefill_limiting_reason", + ) + try: + geo = cache_geometry(state) + return {key: geo[key] for key in keys} + except Exception: # loading/dummy test states may not have stats or complete config + from freetoken.kvcache.cache_status import prefill_geometry + + config = getattr(state, "config", None) + pools = getattr(state, "cache_pools", None) or {} + last = getattr(state, "last_rebuild", None) or {} + requested = int( + last.get("requested_prefill_tokens") + or pools.get("requested_prefill_tokens", 0) + or getattr(config, "max_extend_tokens", 0) or 0 + ) + pool_cap = int( + last.get("pool_prefill_cap_tokens") + or pools.get("pool_prefill_cap_tokens", 0) or 0 + ) + source = str( + last.get("swa_capacity_source") + or pools.get("swa_capacity_source") + or getattr(config, "swa_capacity_source", "none") or "none" + ) + return prefill_geometry( + requested, pool_cap, source, has_window_pool=source != "none" + ) + + async def dispatch_rebuild( state: FrontendManager, *, @@ -531,7 +570,10 @@ async def dispatch_rebuild( # maintenance forever with no reply ever arriving to clear it) and surface the error. state.rebuild_futures.pop(request_id, None) state.maintenance_state = "serving" - return {"status": "failed", "error": f"failed to dispatch rebuild: {e!r}"} + return { + "status": "failed", "error": f"failed to dispatch rebuild: {e!r}", + **_rebuild_prefill_fields(state), + } try: return await asyncio.wait_for(fut, timeout=timeout) except asyncio.TimeoutError: @@ -541,7 +583,10 @@ async def dispatch_rebuild( # blocked; the eventual CacheRebuildReply flips it to serving/failed via # _resolve_rebuild. Drop the now-cancelled future so it does not linger. state.rebuild_futures.pop(request_id, None) - return {"status": "timeout", "request_id": request_id} + return { + "status": "timeout", "request_id": request_id, + **_rebuild_prefill_fields(state), + } def _resolve_num_swa_pages(state: FrontendManager, req: CacheRebuildRequest) -> int | None: @@ -571,32 +616,51 @@ async def cache_rebuild(req: CacheRebuildRequest): state = get_global_state() if state.maintenance_state == "loading": return JSONResponse( - {"status": "loading", "error": "model is still loading; cannot rebuild cache yet"}, + { + "status": "loading", "error": "model is still loading; cannot rebuild cache yet", + **_rebuild_prefill_fields(state), + }, status_code=503, ) if state.maintenance_state == "failed": return JSONResponse( - {"status": "failed", "error": "server latched in maintenance; restart required"}, + { + "status": "failed", "error": "server latched in maintenance; restart required", + **_rebuild_prefill_fields(state), + }, status_code=503, ) if state.maintenance_state == "rebuilding": return JSONResponse( - {"status": "busy", "error": "a cache rebuild is already in progress"}, + { + "status": "busy", "error": "a cache rebuild is already in progress", + **_rebuild_prefill_fields(state), + }, status_code=409, ) if state.maintenance_state == "stopping": return JSONResponse( - {"status": "busy", "error": "engine stop is in progress"}, + { + "status": "busy", "error": "engine stop is in progress", + **_rebuild_prefill_fields(state), + }, status_code=409, ) if req.num_swa_pages is not None and req.swa_full_tokens_ratio is not None: return JSONResponse( - {"status": "failed", "error": "pass num_swa_pages OR swa_full_tokens_ratio, not both"}, + { + "status": "failed", + "error": "pass num_swa_pages OR swa_full_tokens_ratio, not both", + **_rebuild_prefill_fields(state), + }, status_code=422, ) if req.swa_full_tokens_ratio is not None and not 0.0 < req.swa_full_tokens_ratio <= 1.0: return JSONResponse( - {"status": "failed", "error": "swa_full_tokens_ratio must be in (0, 1]"}, + { + "status": "failed", "error": "swa_full_tokens_ratio must be in (0, 1]", + **_rebuild_prefill_fields(state), + }, status_code=422, ) result = await dispatch_rebuild( @@ -786,6 +850,27 @@ def cache_geometry(state: Any) -> dict: # Concrete current window size in that unit (usable pages): the last rebuild's value if it # pinned/derived one, else the load-time pool size. 0 for models without a window pool. "num_swa_pages": int(last.get("num_swa_pages") or pools.get("num_swa_pages", 0) or 0), + "requested_prefill_tokens": int( + pools.get("requested_prefill_tokens", 0) + or getattr(config, "max_extend_tokens", 0) or 0 + ), + "pool_prefill_cap_tokens": int( + last.get("pool_prefill_cap_tokens") + or pools.get("pool_prefill_cap_tokens", 0) or 0 + ), + "effective_prefill_tokens": int( + last.get("effective_prefill_tokens") + or pools.get("effective_prefill_tokens", 0) + or getattr(config, "max_extend_tokens", 0) or 0 + ), + "swa_capacity_source": str( + last.get("swa_capacity_source") + or pools.get("swa_capacity_source", "none") or "none" + ), + "prefill_limiting_reason": str( + last.get("prefill_limiting_reason") + or pools.get("prefill_limiting_reason", "none") or "none" + ), # Engine's exact total cache VRAM budget (all pools), from the ("meta", …) ack. 0 when # unknown (pre-budget engine) — the desktop then reverse-derives it from the limits. "cache_budget_bytes": int(getattr(state, "cache_budget_bytes", 0) or 0), diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index b2857b75..d25643bb 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -338,6 +338,19 @@ def _infer_reasoning_parser(model_path: str) -> str | None: ), ) + parser.add_argument( + "--swa-num-tokens", + type=_positive_int, + dest="swa_num_token_override", + default=ServerArgs.swa_num_token_override, + help=( + "Expert override for absolute window (SWA) capacity in tokens. DSV4 normally " + "derives the minimum from --max-prefill-length; an override must be large enough " + "to honor that request and align to the resolved SWA page size (DSV4: its window " + "size; radix-SWA: 1)." + ), + ) + parser.add_argument( "--page-size", type=int, diff --git a/python/freetoken/tokenizer/server.py b/python/freetoken/tokenizer/server.py index 530e862d..5f7c3097 100644 --- a/python/freetoken/tokenizer/server.py +++ b/python/freetoken/tokenizer/server.py @@ -191,6 +191,11 @@ def tokenize_worker( num_pages=m.num_pages, mamba_slots=m.mamba_slots, num_swa_pages=m.num_swa_pages, + requested_prefill_tokens=m.requested_prefill_tokens, + pool_prefill_cap_tokens=m.pool_prefill_cap_tokens, + effective_prefill_tokens=m.effective_prefill_tokens, + swa_capacity_source=m.swa_capacity_source, + prefill_limiting_reason=m.prefill_limiting_reason, error=m.error, ) ) diff --git a/tests/engine/test_attention_backend_matrix.py b/tests/engine/test_attention_backend_matrix.py index f27640b9..e88464ea 100644 --- a/tests/engine/test_attention_backend_matrix.py +++ b/tests/engine/test_attention_backend_matrix.py @@ -152,6 +152,19 @@ def test_auto_dsv4_sets_window_page_size(monkeypatch): assert config.page_size == 128 +def test_swa_capacity_source_vocabulary_for_generic_and_dense_models(monkeypatch): + from freetoken.engine.engine import _adjust_config + + _patch_env(monkeypatch) + swa = _config("swa", attention_backend="triton") + _adjust_config(swa) + assert swa.swa_capacity_source == "derived" + + dense = _config("full", attention_backend="triton") + _adjust_config(dense) + assert dense.swa_capacity_source == "none" + + @pytest.mark.parametrize( "kind, backend", [ diff --git a/tests/engine/test_cache_budget.py b/tests/engine/test_cache_budget.py index a164f0b4..5a070a67 100644 --- a/tests/engine/test_cache_budget.py +++ b/tests/engine/test_cache_budget.py @@ -166,6 +166,99 @@ def test_adjust_config_allows_auto_for_dsv4(): assert cfg.moe_backend == "offload" assert cfg.moe_cache_auto is True # resolved later at engine init, not here assert cfg.page_size == 128 # DSV4's KV page is the P-token window page + assert cfg.swa_capacity_source == "derived" + + +def test_adjust_config_preserves_dsv4_prefill_limit(): + # DSV4's pool may narrow the scheduler bound, but config resolution must not replace an + # explicit chunk limit with the full context length. + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg(max_seq_len=524288, max_extend_tokens=8192) + _adjust_config(cfg) + assert cfg.max_extend_tokens == 8192 + + +def test_adjust_config_resolves_dsv4_swa_tokens_to_window_pages(): + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg(swa_num_token_override=19200) + _adjust_config(cfg) + assert cfg.swa_num_pages_override == 150 + # Keep the user-facing value intact as the record of what was requested. + assert cfg.swa_num_token_override == 19200 + assert cfg.swa_capacity_source == "explicit" + + +def test_adjust_config_derives_swa_from_prefill_as_the_normal_dsv4_path(): + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg(max_running_req=4, max_extend_tokens=24576) + _adjust_config(cfg) + assert getattr(cfg, "swa_num_token_override", None) is None + assert cfg.swa_num_pages_override == 407 + assert cfg.swa_capacity_source == "derived" + + +def test_adjust_config_rejects_swa_too_small_for_requested_prefill(): + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg( + max_running_req=4, max_extend_tokens=8192, swa_num_token_override=19200 + ) + with pytest.raises( + ValueError, + match=r"19200 is below the minimum 19328 SWA tokens required", + ): + _adjust_config(cfg) + + +def test_adjust_config_accepts_exact_swa_requirement_for_prefill(): + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg( + max_running_req=4, max_extend_tokens=8192, swa_num_token_override=19328 + ) + _adjust_config(cfg) + assert cfg.swa_num_pages_override == 151 + + +def test_adjust_config_rejects_tiny_override_before_working_set_floor_inflation(): + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg( + max_seq_len=128, + max_running_req=4, + max_extend_tokens=128, + swa_num_token_override=128, + ) + with pytest.raises(ValueError, match=r"128 is below the minimum 3200 SWA tokens required"): + _adjust_config(cfg) + + +def test_adjust_config_rejects_misaligned_dsv4_swa_tokens(): + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg(swa_num_token_override=19199) + with pytest.raises(ValueError, match="resolved SWA page size 128"): + _adjust_config(cfg) + + +def test_adjust_config_rejects_nonpositive_programmatic_swa_tokens(): + from freetoken.engine.engine import _adjust_config + + cfg = _dsv4_adjust_cfg(swa_num_token_override=0) + with pytest.raises(ValueError, match="must be positive"): + _adjust_config(cfg) + + +def test_adjust_config_rejects_absolute_swa_capacity_without_window_pool(): + from freetoken.engine.engine import _adjust_config + + cfg = _generic_rotary_cfg(max_position=1024, override=1024) + object.__setattr__(cfg, "swa_num_token_override", 19200) + with pytest.raises(ValueError, match="requires DSV4 or a sliding-window model"): + _adjust_config(cfg) def test_adjust_config_resolves_num_tokens_for_dsv4(): diff --git a/tests/kvcache/test_cache_unit_bytes.py b/tests/kvcache/test_cache_unit_bytes.py index c0eec34e..fc84e4c6 100644 --- a/tests/kvcache/test_cache_unit_bytes.py +++ b/tests/kvcache/test_cache_unit_bytes.py @@ -255,8 +255,10 @@ def test_compute_cache_pools_reads_load_time_allocations(): num_pages=90112, config=SimpleNamespace( page_size=1, + max_extend_tokens=8192, model_config=SimpleNamespace(dsv4_args=None, has_swa_attention=False), ), + kv_cache=SimpleNamespace(prefill_chunk_budget=8064), moe_offload_cache=SimpleNamespace(cache_size=526, bank_caches={}), # num_slots includes the reserved padding sink; API units are usable slots. linear_state_pool=SimpleNamespace(num_slots=65), @@ -268,6 +270,11 @@ def test_compute_cache_pools_reads_load_time_allocations(): "num_mamba_slots": 64, "swa_page_size": 0, "num_swa_pages": 0, + "requested_prefill_tokens": 8192, + "pool_prefill_cap_tokens": 8064, + "effective_prefill_tokens": 8064, + "swa_capacity_source": "none", + "prefill_limiting_reason": "swa_pool", } @@ -282,9 +289,57 @@ def test_compute_cache_pools_zero_for_missing_pools(): "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", } +def test_cache_status_reports_effective_prefill_clamp(): + from freetoken.cache_report import format_cache_status + + rendered = format_cache_status({ + "state": "serving", + "geometry": { + "requested_prefill_tokens": 8192, + "pool_prefill_cap_tokens": 8064, + "effective_prefill_tokens": 8064, + "swa_capacity_source": "explicit", + "prefill_limiting_reason": "swa_pool", + }, + }) + assert ( + "prefill=8064 tok (requested 8192, pool cap 8064, " + "source explicit, reason swa_pool)" + ) in rendered + + +def test_cache_status_always_reports_equal_and_generic_prefill_geometry(): + from freetoken.cache_report import format_cache_status + + equal = format_cache_status({"geometry": { + "requested_prefill_tokens": 8192, + "pool_prefill_cap_tokens": 8192, + "effective_prefill_tokens": 8192, + "swa_capacity_source": "derived", + "prefill_limiting_reason": "requested_and_swa_pool", + }}) + assert ( + "requested 8192, pool cap 8192, source derived, reason requested_and_swa_pool" + ) in equal + + generic = format_cache_status({"geometry": { + "requested_prefill_tokens": 8192, + "pool_prefill_cap_tokens": 0, + "effective_prefill_tokens": 8192, + "swa_capacity_source": "none", + "prefill_limiting_reason": "requested_limit", + }}) + assert "requested 8192, pool cap none, source none, reason requested_limit" in generic + + def test_status_meta_includes_pools(): eng = _mha_engine() eng.config = _config() @@ -292,3 +347,30 @@ def test_status_meta_includes_pools(): meta = compute_cache_status_meta(eng) assert meta["pools"]["num_pages"] == 100 assert meta["pools"]["page_size"] == 16 + + +def test_status_meta_reports_effective_derived_dsv4_ratio(monkeypatch): + import freetoken.kvcache.cache_status as cache_status + + monkeypatch.setattr(cache_status, "compute_cache_unit_bytes", lambda engine: {}) + monkeypatch.setattr(cache_status, "compute_cache_floors", lambda engine: {}) + config = SimpleNamespace( + page_size=128, max_extend_tokens=24576, max_running_req=4, + cache_type="swa_radix", swa_capacity_source="derived", + swa_full_tokens_ratio=0.2, memory_ratio=0.9, + model_config=SimpleNamespace( + dsv4_args=SimpleNamespace(window_size=128), has_swa_attention=False, + ), + ) + engine = SimpleNamespace( + config=config, num_pages=4096, + kv_cache=SimpleNamespace( + sizes=SimpleNamespace(n_win_pages=408), prefill_chunk_budget=24576, + ), + moe_offload_cache=None, linear_state_pool=None, + _post_weights_free=0, _baseline_free=0, _weights_bytes=0, + ) + + meta = compute_cache_status_meta(engine) + assert meta["pools"]["num_swa_pages"] == 407 + assert meta["swa_full_tokens_ratio"] == pytest.approx(407 / 4096) diff --git a/tests/kvcache/test_pool_sizing_surface.py b/tests/kvcache/test_pool_sizing_surface.py index d50a5be4..5764a2c5 100644 --- a/tests/kvcache/test_pool_sizing_surface.py +++ b/tests/kvcache/test_pool_sizing_surface.py @@ -112,6 +112,107 @@ def test_dsv4_kv_cost_and_floor_parity(): DSV4PagedKVCache.solve_num_pages(_dsv4_config(num_page_override=1), 0) +def test_dsv4_prefill_budget_and_required_window_are_inverse_at_page_boundaries(): + from freetoken.kvcache.dsv4_cost_model import ( + dsv4_prefill_chunk_budget, + dsv4_required_swa_pages, + ) + + # max_running_req=4 + radix reserves 23 usable pages. A 64-page (8192-token) + # chunk therefore needs 23 + 2*64 = 151 usable SWA pages. + assert dsv4_prefill_chunk_budget(150, 4, True) == 8064 + assert dsv4_required_swa_pages(8192, 4, True) == 151 + assert dsv4_prefill_chunk_budget(151, 4, True) == 8192 + assert dsv4_required_swa_pages(24576, 4, True) == 407 + assert dsv4_prefill_chunk_budget(407, 4, True) == 24576 + + +def test_dsv4_physical_floor_can_hold_its_minimum_prefill_reach(): + from freetoken.kvcache.dsv4_cost_model import ( + _dsv4_window_floor_pages, + dsv4_prefill_chunk_budget, + ) + + config = _dsv4_config() + object.__setattr__(config, "max_seq_len", 128) + object.__setattr__(config, "max_running_req", 4) + floor_physical = _dsv4_window_floor_pages(config, config.page_size) + assert floor_physical == 26 # 25 usable (23 reserve + twice one 128-token chunk) + dummy + assert dsv4_prefill_chunk_budget(floor_physical - 1, 4, True) == 128 + + +def test_dsv4_pinned_window_cost_matches_allocated_geometry(): + from freetoken.kvcache.dsv4_cost_model import _dsv4_pool_sizes, dsv4_pool_bytes + from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + + config = _dsv4_config() + object.__setattr__(config, "swa_full_tokens_ratio", 0.9) # must lose to the absolute pin + object.__setattr__(config, "swa_num_pages_override", 150) + per_page, fixed, page_size, _reserve = DSV4PagedKVCache.kv_cost(config) + + for usable_pages in (150, 512, 4096): + sizes = _dsv4_pool_sizes(config, usable_pages + 1) + exact = dsv4_pool_bytes( + sizes, config.model_config.dsv4_args, config.max_running_req + 1 + ) + assert sizes.n_win_pages == 151 # 150 usable + the dummy; ratio is not consulted + assert usable_pages * per_page + fixed == exact + assert page_size == config.page_size + + +def test_dsv4_pinned_window_solver_honors_exact_geometry(): + from freetoken.kvcache.dsv4_cost_model import _dsv4_pool_sizes, dsv4_pool_bytes + from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + + config = _dsv4_config() + object.__setattr__(config, "swa_full_tokens_ratio", 0.01) + object.__setattr__(config, "swa_num_pages_override", 150) + target_pages = 512 + target = _dsv4_pool_sizes(config, target_pages + 1) + budget = dsv4_pool_bytes( + target, config.model_config.dsv4_args, config.max_running_req + 1 + ) + + assert DSV4PagedKVCache.solve_num_pages(config, budget) == target_pages + + +def test_dsv4_explicit_full_capacity_must_cover_absolute_window(): + from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + + config = _dsv4_config(num_page_override=150) + object.__setattr__(config, "swa_num_pages_override", 150) + assert DSV4PagedKVCache.solve_num_pages(config, available_memory=0) == 150 + + object.__setattr__(config, "num_page_override", 149) + with pytest.raises(ValueError, match="is smaller than absolute DSV4 SWA capacity"): + DSV4PagedKVCache.solve_num_pages(config, available_memory=0) + + +def test_dsv4_live_rebuild_rejects_full_capacity_below_absolute_window(): + from freetoken.kvcache.base import CacheRebuildRejected + from freetoken.kvcache.dsv4_paged_pool import DSV4PagedKVCache + + config = _dsv4_config() + object.__setattr__(config, "memory_ratio", 1.0) + object.__setattr__(config, "swa_num_pages_override", 150) + pool = object.__new__(DSV4PagedKVCache) + common = dict( + target_moe=0, per_expert_bytes=0, baseline_free=10**15, + weights_bytes=0, current_num_pages=512, + ) + + # Equality is valid for a combined rebuild: both counts are usable pages. + pool.validate_rebuild(config, num_pages=150, num_swa_pages=150, **common) + + # A combined full+window resize must not silently truncate the requested window. + with pytest.raises(CacheRebuildRejected, match="is smaller than absolute DSV4 SWA capacity"): + pool.validate_rebuild(config, num_pages=149, num_swa_pages=150, **common) + + # A KV-only shrink must honor the existing absolute pin from startup or an earlier rebuild. + with pytest.raises(CacheRebuildRejected, match="is smaller than absolute DSV4 SWA capacity"): + pool.validate_rebuild(config, num_pages=149, num_swa_pages=None, **common) + + def test_startup_kv_budget_composition(): # The engine-side budget must stay ratio*old - (old - new); a sign/order slip here # mis-sizes every model's startup KV pool with the whole CPU suite green. diff --git a/tests/scheduler/test_cache_rebuild.py b/tests/scheduler/test_cache_rebuild.py index 0c058070..b1dc5cd2 100644 --- a/tests/scheduler/test_cache_rebuild.py +++ b/tests/scheduler/test_cache_rebuild.py @@ -145,26 +145,82 @@ def test_rebuild_cache_refreshes_prefill_budget(monkeypatch): monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None) + from freetoken.scheduler.cache import CacheManager + + pool = SimpleNamespace(prefill_chunk_budget=5000, swa_paged=False) + page_table = _page_table(4, 64) + cache_manager = CacheManager( + num_pages=8, page_size=2, page_table=page_table, type="radix", swa_pool=pool + ) + sched = Scheduler.__new__(Scheduler) sched.prefill_manager = SimpleNamespace(runnable=False) sched.decode_manager = SimpleNamespace(runnable=False) sched.device = torch.device("cpu") sched.config = SimpleNamespace(tp_info=SimpleNamespace(size=1), max_extend_tokens=100_000) + def rebuild_runtime_cache(**kw): + # DSV4 pools rebuild in place; their capability changes without replacing the object held + # by CacheManager. + pool.prefill_chunk_budget = 1000 + sched.engine = SimpleNamespace( - rebuild_runtime_cache=lambda **kw: None, num_pages=32, page_table=None + rebuild_runtime_cache=rebuild_runtime_cache, num_pages=8, page_table=page_table ) - # engine.page_table unchanged across the (stubbed) rebuild -> no token_pool re-point. - sched.table_manager = SimpleNamespace(page_table=None) - # DSV4-like manager: prefill_chunk_budget tracks the (about-to-shrink) window pool; no shared - # page table, so rebuild_cache's prefix-cache rebuild branch is skipped. - cache_manager = SimpleNamespace( - prefill_chunk_budget=5000, rebuild=lambda *a: None, check_integrity=lambda: None) + # A window-only rebuild leaves the generic page table and token pool in place. + sched.table_manager = SimpleNamespace(page_table=page_table) sched.cache_manager = cache_manager - sched.table_manager.rebuild = lambda pt: None - sched.table_manager.token_pool = None sched.prefill_budget = min(sched.config.max_extend_tokens, cache_manager.prefill_chunk_budget) assert sched.prefill_budget == 5000 - cache_manager.prefill_chunk_budget = 1000 # the (stubbed) engine rebuild shrank the pool - Scheduler.rebuild_cache(sched, num_pages=16) + Scheduler.rebuild_cache(sched, num_swa_pages=10) + assert cache_manager.prefill_chunk_budget == 1000 assert sched.prefill_budget == 1000 # tracks the shrunk cap, not the stale 5000 + + +def test_destructive_swa_rollback_restores_prior_capacity_source(): + from types import SimpleNamespace + + from freetoken.scheduler.scheduler import Scheduler + + sched = Scheduler.__new__(Scheduler) + sched._pending_rebuild = SimpleNamespace( + request_id="r1", moe_cache_size=None, num_pages=None, + num_mamba_slots=None, num_swa_pages=10, + ) + sched.config = SimpleNamespace( + tp_info=SimpleNamespace(size=1), swa_capacity_source="derived" + ) + sched.engine = SimpleNamespace(rebuild_teardown_started=False) + snapshot = { + "moe_cache_size": 4602, "num_pages": 4096, "num_mamba_slots": None, + "num_swa_pages": 407, "requested_prefill_tokens": 24576, + "pool_prefill_cap_tokens": 24576, "effective_prefill_tokens": 24576, + "swa_capacity_source": "derived", + "prefill_limiting_reason": "requested_and_swa_pool", + } + sched._current_cache_geometry = lambda: dict(snapshot) + calls = [] + + def rebuild_cache(**targets): + calls.append(targets) + object.__setattr__(sched.config, "swa_capacity_source", "explicit") + if len(calls) == 1: + sched.engine.rebuild_teardown_started = True + raise RuntimeError("injected post-teardown failure") + + replies = [] + sched.rebuild_cache = rebuild_cache + sched._log_cache_geometry = lambda event: None + sched._reply_rebuild = lambda request_id, status, error=None: replies.append( + (request_id, status, error) + ) + + Scheduler._execute_pending_rebuild(sched) + + assert calls == [ + {"moe_cache_size": None, "num_pages": None, "num_mamba_slots": None, + "num_swa_pages": 10}, + {"num_swa_pages": 407}, + ] + assert sched.config.swa_capacity_source == "derived" + assert replies[0][1] == "rejected" diff --git a/tests/server/test_cache_args.py b/tests/server/test_cache_args.py new file mode 100644 index 00000000..ca8cc693 --- /dev/null +++ b/tests/server/test_cache_args.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from freetoken.server.args import parse_args + + +def test_parse_absolute_swa_token_capacity(): + args, run_shell = parse_args([ + "--model", "/tmp", + "--dtype", "bfloat16", + "--tool-call-parser", "llama3", + "--reasoning-parser", "off", + "--swa-num-tokens", "19200", + ]) + + assert run_shell is False + assert args.swa_num_token_override == 19200 + # Conversion waits for model config resolution, where the SWA page unit is known. + assert args.swa_num_pages_override is None + + +def test_rebuild_and_rejection_output_include_complete_prefill_geometry(): + from freetoken.control_cli import _decode_error_body, _format_rebuild + + doc = { + "status": "rejected", + "error": "requested cache does not fit", + "requested_prefill_tokens": 24576, + "pool_prefill_cap_tokens": 19200, + "effective_prefill_tokens": 19200, + "swa_capacity_source": "explicit", + "prefill_limiting_reason": "swa_pool", + } + expected = ( + "prefill requested=24576 pool_cap=19200 effective=19200 " + "source=explicit reason=swa_pool" + ) + assert expected in _format_rebuild(doc) + assert expected in _decode_error_body(__import__("json").dumps(doc).encode()) diff --git a/tests/server/test_rebuild_maintenance.py b/tests/server/test_rebuild_maintenance.py index a1d10a7e..b87ad457 100644 --- a/tests/server/test_rebuild_maintenance.py +++ b/tests/server/test_rebuild_maintenance.py @@ -59,6 +59,11 @@ def _reply(request_id, status, **over): num_pages=0, mamba_slots=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", error=None, ) base.update(over) @@ -80,6 +85,10 @@ async def _run(): result = asyncio.run(_run()) assert result["status"] == "failed" assert "zmq push failed" in result["error"] + assert set(( + "requested_prefill_tokens", "pool_prefill_cap_tokens", "effective_prefill_tokens", + "swa_capacity_source", "prefill_limiting_reason", + )).issubset(result) assert state.maintenance_state == "serving" assert state.rebuild_futures == {} # no dangling future leaked @@ -98,6 +107,7 @@ async def _run(): result = asyncio.run(_run()) assert result["status"] == "timeout" + assert result["swa_capacity_source"] == "none" assert state.maintenance_state == "rebuilding" # still gated on purpose assert state.rebuild_futures == {} # cancelled future dropped, not leaked @@ -254,6 +264,7 @@ def test_cache_rebuild_guarded_during_loading(): client = TestClient(api.app) r = client.post("/v1/cache/rebuild", json={}) assert r.status_code == 503 + assert r.json()["prefill_limiting_reason"] == "none" assert "loading" in r.json().get("error", "").lower() finally: api._GLOBAL_STATE = prev