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. - 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, + ) 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 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):