From 4d6d63fd372dc03d8f1e00d0d4ed4ddf5da643ad Mon Sep 17 00:00:00 2001 From: Aswinkumar Date: Sun, 30 Aug 2026 15:22:29 -0700 Subject: [PATCH 1/4] Add KV ladder support in argument parser Added arguments for enabling KV ladder and setting ladder step size. --- python/freetoken/server/args.py | 41 +++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index 6696f65dd..01dc9f217 100644 --- a/python/freetoken/server/args.py +++ b/python/freetoken/server/args.py @@ -32,6 +32,11 @@ class ServerArgs(SchedulerConfig): # Default max output (decode) tokens for a request that omits one. None falls back to the # adapter's built-in default (32k). max_output_tokens: int | None = None + # Grow KV at request boundaries, funding each growth by shrinking the GPU MoE expert + # cache. Currently restricted to single-request serving. + enable_kv_ladder: bool = False + # KV tokens per ladder rung. The automatic startup reserve is twice this value. + ladder_step_size: int = 32_768 # Report the prefix-cache hit in each response's usage block (OpenAI # prompt_tokens_details.cached_tokens, Anthropic cache_read_input_tokens, Responses # input_tokens_details.cached_tokens). Mirrors sglang's --enable-cache-report. @@ -266,6 +271,27 @@ def _infer_reasoning_parser(model_path: str) -> str | None: help="Default max output tokens for requests that omit one (default 32k).", ) + parser.add_argument( + "--enable-kv-ladder", + action="store_true", + default=ServerArgs.enable_kv_ladder, + help=( + "Start KV at two ladder steps and grow it by one step whenever " + "prompt + max output reaches the current bound, shrinking the MoE cache to fit. " + "Requires --moe-cache-auto, --max-running-requests 1, and TP=1." + ), + ) + + parser.add_argument( + "--ladder-step-size", + type=_positive_int, + default=ServerArgs.ladder_step_size, + help=( + "KV tokens per automatic ladder step (default 32768); the ladder starts with " + "twice this capacity. Only used with --enable-kv-ladder." + ), + ) + parser.add_argument( "--memory-ratio", type=float, @@ -681,6 +707,21 @@ def _infer_reasoning_parser(model_path: str) -> str | None: if is_offload_moe_backend(kwargs["moe_backend"]) and _no_cache_flag: kwargs["moe_cache_auto"] = True + if kwargs["enable_kv_ladder"]: + if kwargs["max_running_req"] != 1: + parser.error("--enable-kv-ladder requires --max-running-requests 1") + if kwargs["tensor_parallel_size"] != 1: + parser.error("--enable-kv-ladder currently requires --tensor-parallel-size 1") + if kwargs["moe_backend"] in ("cpu", "fused"): + parser.error("--enable-kv-ladder requires the offload or hybrid MoE backend") + if not kwargs["moe_cache_auto"]: + parser.error("--enable-kv-ladder requires --moe-cache-auto") + # --moe-cache-auto consumes this floor during startup, giving the ladder its first + # 2x-step rung without a wasteful post-startup graph recapture. + kwargs["kv_reserve_tokens"] = max( + kwargs["kv_reserve_tokens"], 2 * kwargs["ladder_step_size"] + ) + if kwargs["model_source"] == "modelscope": model_path = kwargs["model_path"] if not os.path.isdir(model_path): From e9dddf46d9c806365eb739057a491f430f5408d7 Mon Sep 17 00:00:00 2001 From: Aswinkumar Date: Sun, 30 Aug 2026 15:23:28 -0700 Subject: [PATCH 2/4] Implement KV ladder policy in scheduler Add KV ladder policy implementation and management --- python/freetoken/scheduler/scheduler.py | 225 ++++++++++++++++++++---- 1 file changed, 187 insertions(+), 38 deletions(-) diff --git a/python/freetoken/scheduler/scheduler.py b/python/freetoken/scheduler/scheduler.py index 48923e3b0..d6215b7d2 100644 --- a/python/freetoken/scheduler/scheduler.py +++ b/python/freetoken/scheduler/scheduler.py @@ -133,6 +133,10 @@ def __init__(self, config: SchedulerConfig): min(config.max_extend_tokens, _chunk_cap) if _chunk_cap else config.max_extend_tokens ) self.config = config + self._kv_ladder_waiting: list[UserMsg] = [] + self._kv_ladder = ( + self._make_kv_ladder_policy() if getattr(config, "enable_kv_ladder", False) else None + ) self.status_reporter = SchedulerStatusReporter( log=logger.info_rank0, decode_log_interval=config.decode_log_interval, @@ -141,6 +145,54 @@ def __init__(self, config: SchedulerConfig): # Initialize the I/O mixin super().__init__(config, self.engine.tp_cpu_group) + def _make_kv_ladder_policy(self): + """Bind the pure ladder arithmetic to this engine's measured cache costs.""" + from freetoken.engine.cache_budget import expert_bytes_per_slot, net_cache_budget_bytes + from freetoken.kvcache.linear_state_pool import state_pool_bytes + + from .kv_ladder import KVLadderPolicy + + if self.config.tp_info.size != 1: + raise ValueError("--enable-kv-ladder currently requires TP=1") + if self.config.max_running_req != 1: + raise ValueError("--enable-kv-ladder requires max_running_req=1") + if not self.cache_manager.supports_runtime_rebuild: + raise ValueError("--enable-kv-ladder is unsupported by this model's KV cache") + if getattr(self.config.model_config, "dsv4_args", None) is not None: + raise ValueError("--enable-kv-ladder does not yet support DSV4's owned KV tiers") + moe = self.engine.moe_offload_cache + if moe is None: + raise ValueError("--enable-kv-ladder requires an offloaded MoE slot cache") + + cache_per_page, fixed_cache_size, _, _ = type(self.engine.kv_cache).kv_cost(self.config) + physical_mamba_slots = ( + self.engine.linear_state_pool.num_slots + if self.engine.linear_state_pool is not None else None + ) + fixed_cache_size += state_pool_bytes(self.config, physical_mamba_slots) + pool_budget = net_cache_budget_bytes( + self.config.memory_ratio, + self.engine._baseline_free, + self.engine._weights_bytes, + fixed_cache_size, + ) + policy = KVLadderPolicy( + step_tokens=self.config.ladder_step_size, + max_context_tokens=self.config.max_seq_len, + page_size=self.config.page_size, + pool_budget_bytes=pool_budget, + kv_bytes_per_page=cache_per_page, + moe_bytes_per_slot=expert_bytes_per_slot(moe.bank_sources), + min_moe_slots=self.config.model_config.num_experts, + ) + logger.info_rank0( + "KV ladder enabled: step=%d tokens, current=%d tokens, ceiling=%d tokens", + policy.step_tokens, + self.engine.num_pages * self.config.page_size, + policy.max_context_tokens, + ) + return policy + def run_when_idle(self) -> None: """Called when the scheduler is idle to perform background tasks.""" logger.info_rank0("Scheduler is idle, waiting for new reqs...") @@ -210,6 +262,7 @@ def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None: or self.prefill_manager.runnable or self.decode_manager.runnable or self._pending_rebuild is not None # a queued rebuild to drain toward + execute + or getattr(self, "_kv_ladder_waiting", None) ) for msg in self.receive_msg(blocking=blocking): self._process_one_msg(msg) @@ -222,6 +275,12 @@ def overlap_loop(self, last_data: ForwardData | None) -> ForwardData | None: ): self._execute_pending_rebuild() + if last_data is None and not ( + self.prefill_manager.runnable or self.decode_manager.runnable + or self._pending_rebuild is not None + ): + self._drain_kv_ladder_waiting() + # Order this iteration's host->device token_pool copies (issued on ``self.stream`` # during scheduling) after the previous batch's sampled-token writes (issued on the # engine stream in ``_forward``). Without this, a request that reuses a just-freed @@ -255,6 +314,7 @@ def normal_loop(self) -> None: self.prefill_manager.runnable or self.decode_manager.runnable or self._pending_rebuild is not None # a queued rebuild to execute at idle + or getattr(self, "_kv_ladder_waiting", None) ) for msg in self.receive_msg(blocking=blocking): self._process_one_msg(msg) @@ -267,6 +327,12 @@ def normal_loop(self) -> None: ): self._execute_pending_rebuild() + if not ( + self.prefill_manager.runnable or self.decode_manager.runnable + or self._pending_rebuild is not None + ): + self._drain_kv_ladder_waiting() + forward_input = self._schedule_next_batch() ongoing_data = None if forward_input is not None: @@ -474,6 +540,115 @@ def _gpu_mem_bytes(self) -> int: return 0 return torch.cuda.memory_reserved(self.device) + def _kv_ladder_plan(self, msg: UserMsg): + policy = getattr(self, "_kv_ladder", None) + if policy is None: + return None + moe = self.engine.moe_offload_cache + assert moe is not None + return policy.plan( + current_pages=self.engine.num_pages, + current_moe_slots=moe.cache_size, + input_tokens=len(msg.input_ids), + max_output_tokens=msg.sampling_params.max_tokens, + ) + + def _queue_for_kv_ladder(self, msg: UserMsg) -> bool: + """Hold a request until the next idle point if its possible length reaches this rung.""" + from .kv_ladder import KVLadderCapacityError + + # Preserve arrival order if an earlier request is already waiting for a rung change. + # max_running_req=1 means these drain one at a time and are re-planned against the + # geometry left by their predecessor. + if self._kv_ladder_waiting: + self._kv_ladder_waiting.append(msg) + return True + try: + plan = self._kv_ladder_plan(msg) + except KVLadderCapacityError as exc: + # Keep serving with the current geometry. Normal admission below will either clip + # max_tokens or return context_length_exceeded, rather than stranding the request. + logger.warning_rank0("KV ladder cannot grow for request %d: %s", msg.uid, exc) + return False + if plan is None: + return False + self._kv_ladder_waiting.append(msg) + logger.info_rank0( + "KV ladder queued request %d: possible length %d reaches %d-token rung; " + "next target %d tokens", + msg.uid, plan.required_tokens, plan.current_tokens, plan.target_tokens, + ) + return True + + def _drain_kv_ladder_waiting(self) -> None: + """At an idle safe point, grow the caches and then admit one held request.""" + waiting = getattr(self, "_kv_ladder_waiting", None) + if not waiting: + return + assert not self.prefill_manager.runnable and not self.decode_manager.runnable + msg = waiting.pop(0) + + from .kv_ladder import KVLadderCapacityError + + try: + plan = self._kv_ladder_plan(msg) + except KVLadderCapacityError as exc: + logger.warning_rank0("KV ladder cannot grow for request %d: %s", msg.uid, exc) + plan = None + if plan is not None: + self._pending_rebuild = CacheRebuildBackendMsg( + request_id=f"auto-kv-ladder:{msg.uid}:{plan.target_pages}", + moe_cache_size=plan.target_moe_slots, + num_pages=plan.target_pages, + ) + status = self._execute_pending_rebuild() + if status == "ok": + logger.info_rank0( + "KV ladder advanced to %d tokens; MoE cache %d -> %d slots", + plan.target_tokens, plan.current_moe_slots, plan.target_moe_slots, + ) + else: + logger.warning_rank0( + "KV ladder rebuild for request %d ended with status=%s; " + "admitting against the retained geometry", + msg.uid, status, + ) + self._admit_user_msg(msg) + + def _admit_user_msg(self, msg: UserMsg) -> None: + input_len, max_seq_len = len(msg.input_ids), self.engine.max_seq_len + max_output_len = max_seq_len - input_len + if max_output_len <= 0: + logger.warning_rank0( + f"Input sequence length {input_len} exceeds {max_seq_len}, " + f"request {msg.uid} is dropped." + ) + # Tell the client instead of dropping silently — otherwise its wait_for_ack + # never sees a `finished` reply and hangs until the request times out. + self.send_result( + [ + ErrorReplyMsg( + uid=msg.uid, + # "prompt is too long: N tokens > M" is the phrasing Claude Code and + # OpenClaw match on; the Anthropic wire has no error code to read. + error=( + f"prompt is too long: {input_len} tokens > {max_seq_len} maximum " + f"(prompt + generation); shorten the prompt or increase the KV " + f"cache budget" + ), + # OpenAI's standard class for this, for clients that read a code. + code="context_length_exceeded", + ) + ] + ) + return + if msg.sampling_params.max_tokens > max_output_len: + msg.sampling_params.max_tokens = max_output_len + logger.warning_rank0( + f"Adjust max_tokens to {max_output_len} for request {msg.uid}." + ) + self.prefill_manager.add_one_req(msg) + def _process_one_msg(self, msg: BaseBackendMsg) -> None: if isinstance(msg, BatchBackendMsg): for msg in msg.data: @@ -489,44 +664,17 @@ def _process_one_msg(self, msg: BaseBackendMsg) -> None: "Dropping request %d because its abort arrived before admission", msg.uid ) return - input_len, max_seq_len = len(msg.input_ids), self.engine.max_seq_len - max_output_len = max_seq_len - input_len - if max_output_len <= 0: - logger.warning_rank0( - f"Input sequence length {input_len} exceeds {max_seq_len}, " - f"request {msg.uid} is dropped." - ) - # Tell the client instead of dropping silently — otherwise its wait_for_ack - # never sees a `finished` reply and hangs until the request times out. - self.send_result( - [ - ErrorReplyMsg( - uid=msg.uid, - # "prompt is too long: N tokens > M" is the phrasing Claude Code and - # OpenClaw match on; the Anthropic wire has no error code to read. - error=( - f"prompt is too long: {input_len} tokens > {max_seq_len} maximum " - f"(prompt + generation); shorten the prompt or increase the KV " - f"cache budget" - ), - # OpenAI's standard class for this, for clients that read a code. - code="context_length_exceeded", - ) - ] - ) - return - if msg.sampling_params.max_tokens > max_output_len: - msg.sampling_params.max_tokens = max_output_len - logger.warning_rank0( - f"Adjust max_tokens to {max_output_len} for request {msg.uid}." - ) - self.prefill_manager.add_one_req(msg) + if not self._queue_for_kv_ladder(msg): + self._admit_user_msg(msg) elif isinstance(msg, AbortBackendMsg): logger.debug_rank0("Aborting request %d", msg.uid) tombstones = getattr(self, "_abort_tombstones", None) if tombstones is None: tombstones = self._abort_tombstones = {} tombstones[msg.uid] = None + waiting = getattr(self, "_kv_ladder_waiting", None) + if waiting: + waiting[:] = [req for req in waiting if req.uid != msg.uid] # Unknown aborts normally consume their tombstone when the cross-worker UserMsg # catches up. Bound hostile/no-followup abort traffic without affecting realistic # in-flight concurrency. @@ -627,7 +775,7 @@ def _reply_rebuild(self, request_id: str, status: str, error: str | None = None) ] ) - def _execute_pending_rebuild(self) -> None: + def _execute_pending_rebuild(self) -> str: from freetoken.engine.engine import CacheRebuildRejected msg = self._pending_rebuild @@ -653,21 +801,21 @@ def _execute_pending_rebuild(self) -> None: # Rejected before any destructive free — old cache intact, keep serving. logger.warning(f"cache rebuild rejected: {e}") self._reply_rebuild(msg.request_id, "rejected", error=str(e)) - return + return "rejected" except Exception as e: # noqa: BLE001 if not getattr(self.engine, "rebuild_teardown_started", True): # Failed before the destructive phase began: graphs and pools are untouched and # the engine is still serving. A destructive rollback would only add risk. logger.error(f"cache rebuild failed before teardown: {e!r} — old cache intact") self._reply_rebuild(msg.request_id, "rejected", error=repr(e)) - return + return "rejected" if self.config.tp_info.size > 1: # A lone-rank failure cannot be rolled back symmetrically: rebuild_cache runs TP # barriers, and ranks that succeeded will not re-enter them — a solo rollback # would desync the group. Keep the latch-failed behavior for tp>1. logger.error(f"cache rebuild failed: {e!r} — tp>1, latching failed") self._reply_rebuild(msg.request_id, "failed", error=repr(e)) - return + return "failed" # The destructive phase failed — typically a CUDA OOM while reallocating a pool or # recapturing graphs. The graphs/pools are already torn down, so the engine cannot # serve as-is. Rather than latch "failed" (which forces a full process restart), @@ -687,17 +835,18 @@ def _execute_pending_rebuild(self) -> None: "failed", error=f"{e!r}; rollback to the prior geometry also failed: {e2!r}", ) - return + return "failed" logger.warning("cache rebuild rolled back to the previous geometry — still serving") self._log_cache_geometry("Cache rolled back") self._reply_rebuild( msg.request_id, "rejected", error=f"rebuild failed and was rolled back: {e!r}" ) - return + return "rejected" # Outside the try: an ack/send failure after a fully-applied rebuild must not be # mistaken for a rebuild failure and roll back the geometry the engine now serves. self._log_cache_geometry("Cache rebuilt") self._reply_rebuild(msg.request_id, "ok") + return "ok" def _current_cache_geometry(self) -> dict: """The pools' current (serving) sizes as rebuild_cache kwargs — the rollback snapshot and From fc4f9161ade3f67c59b80540575d4bdb711986b1 Mon Sep 17 00:00:00 2001 From: Aswinkumar Date: Sun, 30 Aug 2026 15:24:10 -0700 Subject: [PATCH 3/4] Implement KV ladder policy for request handling This module implements a policy for managing key-value (KV) storage growth at request boundaries by optimizing the use of MoE (Mixture of Experts) cache slots. It includes classes for capacity errors and planning growth strategies based on current and target token requirements. --- python/freetoken/scheduler/kv_ladder.py | 108 ++++++++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 python/freetoken/scheduler/kv_ladder.py diff --git a/python/freetoken/scheduler/kv_ladder.py b/python/freetoken/scheduler/kv_ladder.py new file mode 100644 index 000000000..52ae6ec8c --- /dev/null +++ b/python/freetoken/scheduler/kv_ladder.py @@ -0,0 +1,108 @@ +"""Pure policy for growing KV at request boundaries by trading MoE cache slots. + +The scheduler owns the idle/drain mechanics; this module only decides whether a request +reaches the current rung and, if so, computes a budget-respecting next geometry. Keeping the +arithmetic CUDA-free makes the policy cheap to test and aligned with the engine's preflight. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from freetoken.utils import div_ceil + + +DEFAULT_KV_LADDER_STEP_TOKENS = 32_768 + + +class KVLadderCapacityError(ValueError): + """The requested rung cannot fit even after shrinking MoE to its safe floor.""" + + +@dataclass(frozen=True) +class KVLadderPlan: + required_tokens: int + current_tokens: int + target_tokens: int + target_pages: int + current_moe_slots: int + target_moe_slots: int + + +@dataclass(frozen=True) +class KVLadderPolicy: + step_tokens: int + max_context_tokens: int + page_size: int + pool_budget_bytes: int + kv_bytes_per_page: int + moe_bytes_per_slot: int + min_moe_slots: int + + def __post_init__(self) -> None: + for name in ( + "step_tokens", "max_context_tokens", "page_size", "pool_budget_bytes", + "kv_bytes_per_page", "moe_bytes_per_slot", "min_moe_slots", + ): + if getattr(self, name) <= 0: + raise ValueError(f"{name} must be positive") + + @property + def initial_tokens(self) -> int: + """The intended first rung, capped by the checkpoint's real context window.""" + return min(self.max_context_tokens, 2 * self.step_tokens) + + def plan( + self, + *, + current_pages: int, + current_moe_slots: int, + input_tokens: int, + max_output_tokens: int, + ) -> KVLadderPlan | None: + """Return a growth plan when ``prompt + possible output`` reaches the current rung. + + A long request may skip several rungs in one rebuild; the target remains aligned to the + same ``step_tokens`` ladder. At the model ceiling the ordinary admission path clips an + over-large output budget, exactly as it does without the ladder. + """ + current_tokens = current_pages * self.page_size + required_tokens = input_tokens + max_output_tokens + if input_tokens >= self.max_context_tokens: + return None + if required_tokens < current_tokens or current_tokens >= self.max_context_tokens: + return None + + # Always move at least one step. The +1 preserves one token of breathing room when the + # request's theoretical maximum lands exactly on a rung ("could hit the bound"). + required_rung = div_ceil(required_tokens + 1, self.step_tokens) * self.step_tokens + target_tokens = min( + self.max_context_tokens, + max(current_tokens + self.step_tokens, required_rung), + ) + target_pages = div_ceil(target_tokens, self.page_size) + if target_pages <= current_pages: + return None + + bytes_after_kv = self.pool_budget_bytes - target_pages * self.kv_bytes_per_page + affordable_moe = bytes_after_kv // self.moe_bytes_per_slot + target_moe_slots = min(current_moe_slots, affordable_moe) + if target_moe_slots < self.min_moe_slots: + max_kv_pages = ( + self.pool_budget_bytes - self.min_moe_slots * self.moe_bytes_per_slot + ) // self.kv_bytes_per_page + max_kv_tokens = max(0, max_kv_pages * self.page_size) + raise KVLadderCapacityError( + f"KV ladder rung {target_tokens} tokens cannot fit while retaining the " + f"minimum {self.min_moe_slots} MoE slots (budget permits at most " + f"{max_kv_tokens} KV tokens)" + ) + + return KVLadderPlan( + required_tokens=required_tokens, + current_tokens=current_tokens, + target_tokens=min(target_pages * self.page_size, self.max_context_tokens), + target_pages=target_pages, + current_moe_slots=current_moe_slots, + target_moe_slots=target_moe_slots, + ) From 1b90e9fc97a77b317313aff9cfcd24b86985e9e4 Mon Sep 17 00:00:00 2001 From: Aswinkumar Date: Sun, 30 Aug 2026 15:25:19 -0700 Subject: [PATCH 4/4] Update CLI documentation with new options Added new CLI options for KV ladder and MoE configurations. --- docs/cli.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/cli.md b/docs/cli.md index ff4af382d..bffb180c3 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -43,6 +43,7 @@ parsers all resolve automatically from the checkpoint and the GPU. | `--gpu` | GPU 0 | GPU to run on: a UUID from `nvidia-smi -L` or an `nvidia-smi` index; see [below](#choosing-a-gpu) | | `--max-running-requests` | 4 | Max concurrently running requests | | `--max-output-tokens` | 32768 | Default output budget for requests that omit one | +| `--ladder-step-size` | 32768 | KV tokens per automatic ladder rung; with `--enable-kv-ladder`, startup KV is at least twice this value | | `--max-seq-len-override` | from checkpoint | Max sequence length | | `--max-prefill-length` | 8192 | Chunked-prefill chunk size in tokens | | `--cuda-graph-max-bs`, `--graph` | = max running requests | Max batch size captured as CUDA graphs | @@ -82,6 +83,7 @@ See [models.md](models.md#moe-backends) for what each backend does. | `--moe-backend` | auto | `fused`/`offload`/`cpu`/`hybrid`; auto → offload, or hybrid with a `ft bench bw` profile | | `--moe-cache-size` / `--moe-cache-rate` / `--moe-cache-auto` | auto | GPU expert-cache size as slots / fraction of all experts / sized from free VRAM (mutually exclusive; auto is enabled by default for offload-family backends) | | `--kv-reserve-tokens` | 8192 | KV token floor reserved before `--moe-cache-auto` fills experts | +| `--enable-kv-ladder` | off | With `--moe-cache-auto --max-running-requests 1`: start KV at 2x `--ladder-step-size`, then grow in one-step rungs before requests that could reach the current bound; MoE slots shrink to keep the same VRAM budget, up to the model context limit | | `--moe-cpu-threads` | physical cores | CPU worker threads for the cpu/hybrid executor | | `--moe-cpu-layers` | all on GPU | With `offload`: which MoE layers decode on CPU (`3,7,11`, a count, or a fraction) | | `--moe-hybrid-max-fetch` | auto | With `hybrid`: max experts fetched over PCIe per layer per step; rest computed on CPU | @@ -172,4 +174,3 @@ profile that `ft serve --moe-backend auto` and `--moe-hybrid-max-fetch -1` then - What to measure: `--dtype`, `--model`, `--formats`, `--isa`. - `--threshold` (default 2.0) sets the call: recommend hybrid when CPU bandwidth beats PCIe by that factor. -