From b94fdfdf4142c733d10ad3adb8a8df8a5eea83ca Mon Sep 17 00:00:00 2001 From: Hector Isaza <76170012+hecisaza@users.noreply.github.com> Date: Wed, 30 Sep 2026 15:36:41 +0000 Subject: [PATCH 1/8] scheduler: integrate #246, #221 and #180 as decode-floor v7 (#283) Replace the fair-prefill decode-floor overlay with the integrated v7 patcher qualified in #283: - a selected WAITING request refused KV allocation no longer stalls a runnable prefill; it keeps being retried and recovers once memory frees (#246) - explicit priorities keep ranking fair-prefill candidates under --scheduling-policy priority; FCFS is unchanged (#221) - installer migration follows #180's fail-closed contract: pristine, v1, v2, v5 and v5+priority installs migrate to the fresh-install v7 bytes; v3/v4, drifted, duplicated, marker-only and unmarked helpers are refused unwritten patch_mamba_align_chunking.py accepts decode-floor v7. #180's concurrency canary, its docs and its CPU progress regressions land with it. The decode-floor installation test is ported to the v7 contract with deployed producer fixtures, and the mamba A4 fixture downgrades v7 as well as v5. On image ef9f5013 these overlays install scheduler sha256 97c90a18fd0df6f9aeb1452c39801455d90756d315c6846ca0642e934fe4e269, the identity qualified live in #283 (baseline -> candidate -> baseline, TP2). Co-authored-by: RandomLlama <108938230+beastllama@users.noreply.github.com> Co-authored-by: Surlebeat <64695525+surlebeat@users.noreply.github.com> Co-authored-by: Mauvis Ledford <20738+krunkosaurus@users.noreply.github.com> --- CHANGELOG.md | 34 + docs/concurrent-agents.md | 104 ++ overlay/patch_mamba_align_chunking.py | 7 +- overlay/patch_scheduler_decode_floor.py | 616 ++++++--- tests/check_concurrent_agents.py | 189 +++ .../v5_main_d5713ec6.py | 1123 ++++++++++++++++ .../v5_priority_763bd30e.py | 1152 +++++++++++++++++ tests/fixtures/legacy_scheduler_helpers.py | 904 +++++++++++++ tests/test_mamba_align_chunking.py | 9 +- tests/test_prefill_concurrency.py | 153 +++ tests/test_scheduler_decode_floor.py | 136 +- 11 files changed, 4189 insertions(+), 238 deletions(-) create mode 100644 docs/concurrent-agents.md create mode 100644 tests/check_concurrent_agents.py create mode 100644 tests/fixtures/decode_floor_producers/v5_main_d5713ec6.py create mode 100644 tests/fixtures/decode_floor_producers/v5_priority_763bd30e.py create mode 100644 tests/fixtures/legacy_scheduler_helpers.py create mode 100644 tests/test_prefill_concurrency.py diff --git a/CHANGELOG.md b/CHANGELOG.md index bb926b27..90943ac8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -247,6 +247,40 @@ There were no git tags for 1.0.0–1.4.0; 1.5.0 is the first cut named as a rele ### Fixed +- Fair-prefill scheduler integration (#283): `overlay/patch_scheduler_decode_floor.py` + moves to decode-floor **v7**, combining the three scheduler directions of + #246 (@beastllama), #221 (@surlebeat) and #180 (@krunkosaurus) in one patcher. + - A selected WAITING request that vLLM refuses KV allocation no longer stalls + a runnable prefill. It keeps being retried and recovers once + memory frees (#246). + - Explicit request priorities still rank fair-prefill candidates under + `--scheduling-policy priority`. FCFS is unchanged (#221). + - Installer migration follows #180's fail-closed contract. Pristine, v1, v2, + v5 and v5+priority (the #221 layout on `main` until now) installations + migrate to the same bytes as a fresh v7 install, and a re-apply is a + verified no-op. v3/v4, drifted, duplicated, marker-only and unmarked + helpers are refused without writing. + - `overlay/patch_mamba_align_chunking.py` accepts decode-floor v7 as well + as v5. + - #180's concurrency canary (`tests/check_concurrent_agents.py`, + `docs/concurrent-agents.md`) and its CPU progress regressions + (`tests/test_prefill_concurrency.py`) land with it. + + Qualified live on a TP2 pair with a pre-registered baseline → candidate → + baseline run (installed scheduler `97c90a18…`, 2026-09-30): + - PASS: KV-blocked admission progress, the #180 contention canary (5/5 per + arm), TheGrill routine prefill, and the CPU migration matrix. + - Inconclusive, because the two baseline runs differ from each other by more + than the frozen tolerance: decode, concurrency and mixed TheGrill cells, + and temperature-0 output identity. No TheGrill gate regressed. + - Known limitation under `--scheduling-policy priority`: a short + high-priority prompt can wait about 85 s behind a larger KV-blocked + request of the same priority. This is tracked separately. + - `start-tp3.sh` and `start-tp4.sh` load the same two overlays, but only TP2 + was qualified live. + - The proposed TP=2 `CHUNK=0` default is **not** part of this change; + launcher defaults are unchanged. + - `start-tp4.sh` forwards `EXL3_FAT_GROUPED` to all four ranks (#286) and defaults it to `1` (E3) with the coupled `EXL3_TEMP_ROWS_FUSED` default (32 with E3, 256 with E2), matching `start.sh` and `start-tp3.sh`. Before, the diff --git a/docs/concurrent-agents.md b/docs/concurrent-agents.md new file mode 100644 index 00000000..7e5b56c1 --- /dev/null +++ b/docs/concurrent-agents.md @@ -0,0 +1,104 @@ +# Concurrent agents + +Shipped mixed-prefill defaults: `fair` on the two-node `start.sh` (TP=2), +`0` on `start-tp3.sh` (TP=3) and `skip` on `start-tp4.sh` (TP=4). On a launcher +that defaults to `skip`, `GLM53_MIXED_PREFILL_CHUNK=0` is one explicit option +for concurrent interactive agents — not a recommendation. It admits newcomers +immediately at the cost of the incumbent's decode rate, and the TP=2 `CHUNK=0` +default is a gated candidate whose contention qualification is still open. +`fair` is the measured policy for shared TP=2 use. + +Setting `0` restores vLLM's stock chunked-prefill scheduler: a decoding peer no +longer explicitly prevents another prompt from using the remaining token +budget in the same step. `MAX_NUM_SEQS` still limits active sequences, and +`MAX_NUM_BATCHED_TOKENS` limits the whole step. A context still needs its full +prefill before it can emit a first token; this setting does not impose a +wall-clock latency guarantee or bypass a full KV cache. + +The total token budget can still be exhausted by an already-running long +prefill. If a second prompt waits on capacity while cache and sequence slots +are available, inspect the per-request allocation as well as the mixed-prefill +rule. vLLM's existing `--long-prefill-token-threshold` option can cap that +allocation for both running and waiting prompts. For example, with +`MAX_NUM_BATCHED_TOKENS=1024`, a threshold of `512` leaves budget for a second +request (speculative decoding also consumes budget). Add it to `EXTRA_ARGS` +without removing any existing arguments, then recreate and requalify the +service. A smaller cap can reduce single-request prefill throughput; test a +short peer arriving during a cold long prefill and measure both latencies. + +`MAX_MODEL_LEN` is a per-request limit, not a promise that every sequence slot +can hold a context of that size. Compare the engine's startup KV-cache token +capacity and maximum-concurrency estimate with the combined active contexts. +Two million-token windows require capacity for both; enabling interleaving +does not add memory. Agent-side context compression can bound the resident +history while preserving a session across many hours and requests. + +The `skip` policy (still the TP=4 default) gives a new prompt zero prefill tokens +whenever a peer is decoding. A long generation can therefore starve another +agent for minutes, even with spare sequence slots and KV memory. A client idle +timeout then cancels and retries that waiting request. Requests already admitted +to the running list can also have their remaining prefill suppressed. + +`skip` and `-1` remain explicit options for protecting decode throughput at the +expense of new-prompt latency. A positive number caps mixed prefill tokens. +Interleaving can reduce per-agent generation speed, especially for long +contexts on the older sparse-MLA kernels. Choose that tradeoff explicitly. + +Existing `.env` files are not rewritten during updates. Set the policy you have +chosen for that launcher (`fair`, `skip`, a positive cap, or `0`) there, then +recreate the service through your approved deployment procedure. A plain `docker +restart` reuses the old container environment and does not apply this setting. +In-flight requests should be drained first. + +Before measuring, record the deployed source commit, image digest, model and +drafter revisions, launch arguments, context/sequence/token budgets, and +relevant environment. Compare with that exact versioned recipe. Isolate the +mixed-prefill setting; an upgrade of kernels, weights, or memory geometry +requires separate qualification. + +Run the CPU regressions with an installed scheduler file. These operate on +temporary copies and do not import the model or modify the running scheduler: + +```bash +GLM53_SCHEDULER_PY_SRC=/path/to/scheduler.py python3 tests/test_prefill_concurrency.py +GLM53_SCHEDULER_PY_SRC=/path/to/scheduler.py python3 tests/test_scheduler_decode_floor.py +``` + +Reapplication accepts only a complete known patch: the legacy helper site is +validated before anything is removed (v1/v2/v5 against the published helper text +by sha256, the current version against the installer's own text), one exact byte +range is removed, and the whole file is round-trip checked — re-adding that span +and re-applying the version's frozen gate sites must reproduce the input +byte-for-byte. A marker alone is never trusted. A partial patch, changed helper +(including decorators), duplicate helper, or altered gate fails without +modifying the scheduler, and an unversioned variant whose baked default differs +from the published one is refused rather than silently migrated. v3 and v4 are +refused as unsupported without touching the source: no authenticated producer of +their intermediate helper bodies was recovered from public history, so no +canonical v3/v4 image is invented for them. A validated legacy image (v1/v2/v5) +is migrated to the current installer version without changing the operator's +`GLM53_MIXED_PREFILL_CHUNK` value, and +the current version accepts the opt-in gate forms (`GLM53_MIXED_PREFILL_WARM_TOKENS` +/ `_MAX_WAIT_MS`) so a supported gate-v2 scheduler is not rejected at startup. +These CPU checks prove patch integrity and policy behavior, not live request +latency or KV capacity. + +After the deployment and recipe are approved, run one live canary: + +```bash +python3 tests/check_concurrent_agents.py --base-url http://127.0.0.1:8888/v1 --model GLM-5.3-Flash-EXL3 --out logs/concurrent-agents-canary.json +``` + +The canary starts a long, deterministic counting response with thinking +disabled. After A emits content, it submits B with a distinct 8k-word prompt +and asks for `OK`. It requires valid completed streams, the expected short +answer, and B's first content before A's last content. Two HTTP 200s alone do +not pass. The default overall deadline is 240 seconds; on failure it stops +its own requests, returns nonzero, and does not retry or run a benchmark +matrix. If authentication is enabled, supply `OPENAI_API_KEY` through the +environment. Receipts do not include that credential. + +Review the canary's output and latency before wider concurrency testing. +Confirm the model still handles realistic tool calls and long prompts in +your agent client. Increasing a client's idle timeout can accommodate long +prefills but does not repair scheduler starvation. diff --git a/overlay/patch_mamba_align_chunking.py b/overlay/patch_mamba_align_chunking.py index dca291b3..4c635fa9 100755 --- a/overlay/patch_mamba_align_chunking.py +++ b/overlay/patch_mamba_align_chunking.py @@ -9,7 +9,8 @@ The decode-floor overlay owns per-request caps; this overlay owns alignment, positive-grant progress and checkpoint backoff. Their anchors do not overlap, so either application order is supported. Decode-floor versions before v5 are -rejected because they do not preserve per-request alignment caps. +rejected because they do not preserve per-request alignment caps; v5 and +v7 (#283) are accepted. Pristine sources and exact retained v1 installations migrate atomically. Partial or drifted installations fail before writing. @@ -31,6 +32,8 @@ LEGACY_MARK = "# [glm53-mamba-align-chunking-v1]" DECODE_FLOOR_MARK = "# [glm53-decode-floor" DECODE_FLOOR_V5 = "# [glm53-decode-floor:v5]" +# v7 (#283) keeps v5's per-request alignment-cap hooks byte-identical apart from the marker. +DECODE_FLOOR_V7 = "# [glm53-decode-floor:v7]" IMPORT_OLD = """from vllm.v1.kv_cache_interface import KVCacheConfig """ @@ -167,7 +170,7 @@ def main() -> int: if not P.is_file(): raise SystemExit(f"missing {P}") original = text = P.read_text() - if DECODE_FLOOR_MARK in text and DECODE_FLOOR_V5 not in text: + if DECODE_FLOOR_MARK in text and DECODE_FLOOR_V5 not in text and DECODE_FLOOR_V7 not in text: raise SystemExit( f"{P}: patch_scheduler_decode_floor.py older than v5 present; its " "per-request cap is not applied to the Mamba alignment" diff --git a/overlay/patch_scheduler_decode_floor.py b/overlay/patch_scheduler_decode_floor.py index 1019e530..4ae9e8e9 100644 --- a/overlay/patch_scheduler_decode_floor.py +++ b/overlay/patch_scheduler_decode_floor.py @@ -32,9 +32,24 @@ afterwards a 2s age override may borrow one such chunk, only after all shared debt is repaid. In-flight async prefill blocks the next mixed turn. Prefills are selected by last - completed positive prefill service plus round-robin within each - request-priority tier when scheduling-policy=priority. Decode + completed positive prefill service plus round-robin. Decode token/input budget is allocated first by the base scheduler. + v7: under scheduling-policy=priority, lower numeric priority + ranks first (FCFS ignores priority). Selection is + allocation-aware: a runnable prefill already passed when a + later selected candidate is natively refused KV allocation is + carried to the next step within its tier. A carried lower-tier + request is pre-selected ahead of a higher tier only when every + strictly-higher-tier candidate was natively refused in the + previous step (one step per refusal); new or unrefused + higher-tier work always ranks first. DROP-5d: preselection is + refusal-licensed (a candidate refused in the previous step is + skipped once); carry is a trace annotation only, created on a + native refusal and discarded on every release. DROP-5f: under + the priority policy a candidate whose chosen rung exceeds the + credit takes the largest affordable rung when a lower tier is + present, and a lower tier never borrows in a step where a + higher tier was credit/gap-deferred. Solo prefill retains the base scheduler's limits. Timing is a host busy-time proxy. @@ -45,11 +60,17 @@ GLM53_FAIR_PREFILL_MAX_STEP_MS default 2000 (estimated mixed-step limit) GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 -Versioned installer: `# [glm53-decode-floor:v5]`. v1 (no version), v2, v3 -and v4 images are unpatched then re-patched. Fail closed if anchors drift. +Versioned installer: `# [glm53-decode-floor:v7]`. Accepted inputs are +authenticated by exact helper bytes (sha256 + length) and exact hook anchors: +pristine, v1, v2, historical v5, current-main v5, priority v5 (#221), warm/ +deadline v6 (#180/#80) and carry v6 (#246). v3/v4 and any unknown, drifted, +duplicated or marker-only input are refused without writing. +GLM53_FAIR_TRACE= or SCHED_FAIR_TRACE= (default off) appends JSONL +scheduler events; the latter passes through the launcher's GLM53_EXTRA_ENV. """ from __future__ import annotations +import hashlib import inspect import os import sys @@ -67,23 +88,8 @@ MARK_V3 = "# [glm53-decode-floor:v3]" MARK_V4 = "# [glm53-decode-floor:v4]" MARK_V5 = "# [glm53-decode-floor:v5]" -PRIORITY_MARK = "# [glm53-prefill-priority:v1]" -LEGACY_RANK_PREFIX = """ def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - self.last_service.get(r.request_id, 0.0), -""" -PRIORITY_RANK_PREFIX = """ def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - # [glm53-prefill-priority:v1] Preserve fairness within priority tiers. - int(getattr(r, "priority", 0) or 0) if getattr(self, "priority_mode", False) else 0, - self.last_service.get(r.request_id, 0.0), -""" -LEGACY_CANDIDATES = """ self._candidates = self._rank_prefills(prefills) -""" -PRIORITY_CANDIDATES = """ policy = getattr(sched, "policy", None) - self.priority_mode = getattr(policy, "value", policy) == "priority" - self._candidates = self._rank_prefills(prefills) -""" +MARK_V6 = "# [glm53-decode-floor:v6]" +MARK_V7 = "# [glm53-decode-floor:v7]" IMPORT_OLD = """import itertools import time @@ -270,7 +276,7 @@ break """ -class _Glm53MixedPrefill: # [glm53-decode-floor:v5] +class _Glm53MixedPrefill: # [glm53-decode-floor:v7] """Bound contention using completion feedback, without synchronizing GPUs.""" LADDER = (128, 256, 512, 768, 1024, 1536, 2048) @@ -303,10 +309,30 @@ def __init__(self, now=None): self.selected = set() self._candidates = [] self._tried = set() + self._passed = set() + self._carry = set() + self._refused = set() + self._refused_prev = set() + self._deferred_tiers = set() + self.priority_mode = False + self.refusals = 0 self.step_mode = "solo" self.defer_reason = "none" self.missed_prefill = 0 self._model_cache = None + self._incarn = {} + self._trace = None + self.trace_seq = 0 + self.trace_dropped = 0 + # SCHED_FAIR_TRACE is the launcher-forwardable alias: start.sh reserves + # every GLM53_* name in GLM53_EXTRA_ENV. + self.trace_path = (os.environ.get("GLM53_FAIR_TRACE", "").strip() + or os.environ.get("SCHED_FAIR_TRACE", "").strip()) + if self.trace_path: + try: + self._trace = open(self.trace_path, "a", buffering=1 << 16) + except OSError: + self.trace_dropped += 1 def _e(self, name, default): v = os.environ.get(name) @@ -359,7 +385,7 @@ def _parse(self) -> None: self.max_step_s = 2.0 if self.mode == "fair" and not self.logged_boot: print( - f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"[glm53-decode-floor] fair v7 probe_chunk={self.chunk} " f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " f"interval_s={self.interval_s} max_step_s={self.max_step_s} " f"max_chunks={self.max_chunks}", @@ -410,6 +436,10 @@ def _prune(self, live): dead = [k for k in store if k not in live] for k in dead: store.pop(k, None) + for bag in (self._carry, self._refused, self._refused_prev, self._passed): + bag.intersection_update(live) + for rid in [k for k in self._incarn if k not in live]: + self._incarn.pop(rid, None) @property @@ -485,25 +515,96 @@ def _target(self, remaining, room): def _credit_limit(self): return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + def _ev(self, kind, step=None, **fields): + if not self.trace_path: + return + import json # local: the scheduler module does not import json + self.trace_seq += 1 + if step is None: + step = (self._open_rec or {}).get("step_id") + fields.update(ev=kind, seq=self.trace_seq, t=round(self._now(), 6), step=step, + dropped=self.trace_dropped) + try: + if self._trace is None: + raise OSError("trace not open") + self._trace.write(json.dumps(fields, separators=(",", ":"), default=str) + "\n") + except (OSError, ValueError, TypeError): + # Never silently disable: count the loss; seq gaps expose it too. + self.trace_dropped += 1 + + def _trace_flush(self): + if self._trace is not None: + try: + self._trace.flush() + except (OSError, ValueError): + self.trace_dropped += 1 + + def _forget(self, rid): + for bag in (self._carry, self._refused, self._refused_prev, self._passed): + bag.discard(rid) + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + store.pop(rid, None) + + def _track_incarnations(self, reqs): + # A request id re-bound to a different request object is a new request: + # it must not inherit carry, refusal, service age or arrival. + for r in reqs: + rid, obj = r.request_id, id(r) + old = self._incarn.get(rid) + if old is not None and old != obj: + self._forget(rid) + self._ev("incarnation_reset", rid=rid) + self._incarn[rid] = obj + + def _tier(self, r): + if not self.priority_mode: + return 0 + try: + return int(getattr(r, "priority", 0) or 0) + except (TypeError, ValueError): + return 0 + def _rank_prefills(self, prefills): + # [glm53-decode-floor:v7] priority tier first (priority policy only), + # then completed-service age and round-robin within the tier. return sorted(prefills, key=lambda r: ( - # [glm53-prefill-priority:v1] Preserve fairness within priority tiers. - int(getattr(r, "priority", 0) or 0) if getattr(self, "priority_mode", False) else 0, + self._tier(r), self.last_service.get(r.request_id, 0.0), self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) - def _promote_next(self): + def _promote_next(self, prefer_passed=False): grants = (self._open_rec or {}).get("grants", {}) - for r in self._candidates: + candidates = self._candidates + if prefer_passed: + # DROP-5e: after a native refusal, prefer a candidate whose native position + # has NOT passed yet (it can still be scheduled this step), but ONLY within + # its priority tier: a lower tier never backfills past a higher-tier + # candidate, even one whose position already passed (that one takes the + # next step through the refusal licence). Passed ones are annotated as + # carry for the trace only. + order = {r.request_id: i for i, r in enumerate(candidates)} + candidates = sorted(candidates, key=lambda r: ( + self._tier(r), r.request_id in self._passed, order[r.request_id])) + for r in candidates: if len(self.selected | set(grants)) >= self.max_chunks: break if r.request_id not in self._tried: self.selected.add(r.request_id) + if (prefer_passed and r.request_id in self._passed + and r.request_id not in self._refused + and r.request_id not in self._refused_prev): + # Its scheduler position already passed this step: carry it. + self._carry.add(r.request_id) + self._ev("carry", rid=r.request_id, tier=self._tier(r)) def _release(self, rid, reason="allocation"): rec = self._open_rec if rec is None: return + # A zero reported for an unselected request is a policy pass, not an + # admission attempt; keep it eligible for carry. + if rid not in self.selected and rid not in rec["grants"]: + return grant = rec["grants"].pop(rid, None) if grant: self.credit += grant[1] @@ -511,14 +612,37 @@ def _release(self, rid, reason="allocation"): rec["borrowed"] = False self.selected.discard(rid) self._tried.add(rid) + # DROP-5b: a released request (credit, gap, zero-progress or allocation) never + # keeps carry, so a deferral cannot renew its own pre-selection next step. + self._carry.discard(rid) + if reason in ("zero_progress", "allocation"): + # A selected request that could not progress rotates behind its + # tier peers. Credit/gap deferrals keep their fair rank. + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n self.defer_reason = reason - self._promote_next() + self._ev("release", rid=rid, reason=reason, refunded=bool(grant)) + # Only a native allocation refusal may create carry. Credit, gap-budget and + # zero/alignment deferrals re-rank without carrying anyone. + self._promote_next(prefer_passed=(reason == "allocation")) def note_scheduled(self, request, num_new_tokens): # Called after alignment/encoder caps; final allocation is sealed below. if num_new_tokens <= 0: self._release(request.request_id, "zero_progress") + def note_alloc_failed(self, request): + # Native KV allocation refusal (new_blocks is None), distinct from a + # policy/alignment zero. Only a selected or granted request counts. + rid = request.request_id + rec = self._open_rec + if rec is not None and (rid in self.selected or rid in rec["grants"]): + self._refused.add(rid) + self._carry.discard(rid) # a refused request cannot hold carry + self.refusals += 1 + self._ev("refusal", rid=rid, tier=self._tier(request)) + self._release(rid, "allocation") + def protect_decode(self, request): return (self.mode == "fair" and self._open_rec is not None and self._open_rec["had_decode"] @@ -533,6 +657,9 @@ def begin_step(self, sched): if self._sched_id is not None and self._sched_id != sid: self.inflight.clear() self._open_rec = None + self._carry.clear() + self._refused.clear() + self._refused_prev = set() self.credit = 0.0 self.in_contention = False self.last_account_mono = None @@ -553,6 +680,7 @@ def begin_step(self, sched): if self.mode == "fair": sched.running[:] = decodes + [r for r in running if self.needs_prefill_compute(r)] + self._track_incarnations(running + waiting) for r in running + waiting: self.arrival.setdefault(r.request_id, now) policy = getattr(sched, "policy", None) @@ -560,6 +688,18 @@ def begin_step(self, sched): self._candidates = self._rank_prefills(prefills) self._tried = set() self.selected = set() + self._passed = set() + self._deferred_tiers = set() + live_prefills = {r.request_id for r in prefills} + self._carry.intersection_update(live_prefills) + # DROP-5d: carry is a one-step trace annotation only; it never grants + # preselection and never survives into the next step. + self._carry.clear() + # Refusals are evidence for exactly one following step. + self._refused_prev = self._refused & live_prefills + self._refused = set() + # Refused requests keep ordinary ranked retry but may not be carry-preselected. + self._carry.difference_update(self._refused_prev) sched._glm53_align_prefill_limit = None self._open_rec = { "step_id": int(sched.current_step), "t_submit": now, @@ -571,6 +711,7 @@ def begin_step(self, sched): if not decodes and not any(r["had_decode"] for r in self.inflight.values()): self.in_contention = False self.credit = 0.0 + self._carry.clear() self.step_mode = "legacy" if self.mode != "fair" else "solo" self.defer_reason = "none" if self.mode != "fair" or not prefills or not decodes: @@ -583,7 +724,24 @@ def begin_step(self, sched): self.defer_reason = "async_inflight" else: self.step_mode = "prefill_turn" + # DROP-5d refusal-licensed preselection: walk the ranking and skip each + # candidate natively refused in the previous step, so the next-ranked one gets + # the turn. Carry never grants preselection. A lower tier can precede a higher + # tier only when every higher-ranked candidate was refused in the previous step + # (the bounded exception); a new or unrefused candidate always keeps its rank. + # The skipped request is retried in this same step if the selected one defers, + # and ranks normally next step. + for r in self._candidates: + if len(self.selected) >= self.max_chunks: + break + if r.request_id in self._refused_prev: + continue + self.selected.add(r.request_id) self._promote_next() + self._ev("step", mode=self.step_mode, prio=self.priority_mode, + cand=[(r.request_id, self._tier(r)) for r in self._candidates], + sel=sorted(self.selected), carry=sorted(self._carry), + refused_prev=sorted(self._refused_prev)) self._maybe_log() def cap_for(self, sched, request, computed=None): @@ -606,13 +764,16 @@ def cap_for(self, sched, request, computed=None): rec = self._open_rec if rid in rec["grants"]: return rec["grants"][rid][0] + self._passed.add(rid) if self.step_mode != "prefill_turn" or rid not in self.selected: return 0 reserved = sum(g[1] for g in rec["grants"].values()) gap_room = max(0.0, self.max_step_s - reserved) age = self._now() - self.last_service.get(rid, self.arrival[rid]) pick = self._target(remaining, gap_room) + my_tier = self._tier(request) if pick is None: + self._deferred_tiers.add(my_tier) # DROP-5f: a deferred tier blocks lower-tier borrowing self._release(rid, "gap_budget") if age >= self.interval_s: self.missed_prefill += 1 @@ -629,10 +790,26 @@ def cap_for(self, sched, request, computed=None): # after all shared debt is repaid: queue churn or several aged # newcomers cannot repeatedly overdraw the decoder. due = rid not in self.last_service or age >= self.interval_s + # DROP-5f: under the priority policy a lower tier may not borrow in a step + # where a strictly higher-tier candidate was credit- or gap-deferred. + higher_deferred = any(d < my_tier for d in self._deferred_tiers) + # DROP-5f affordable-rung admission: when a strictly lower tier is present, + # a candidate whose chosen (largest step-fitting) rung does not fit the + # credit takes the largest rung that DOES fit it instead of deferring, so + # the lower tier cannot borrow past an affordable higher-tier admission. + # With a single tier (FCFS) the credit-saving behaviour is unchanged. + lower_present = self.priority_mode and any( + self._tier(c) > my_tier for c in self._candidates) + alt = (self._target(remaining, min(gap_room, max(0.0, self.credit))) + if lower_present else None) if (due and self.credit >= -1e-9 and not rec["grants"] - and not rec["borrowed"]): + and not rec["borrowed"] and not higher_deferred): borrowed = True + elif alt is not None: + cap, cost = alt + self._ev("affordable_rung", rid=rid, cap=cap) else: + self._deferred_tiers.add(my_tier) self._release(rid, "credit") if age >= self.interval_s: self.missed_prefill += 1 @@ -640,6 +817,7 @@ def cap_for(self, sched, request, computed=None): return 0 self.credit -= cost rec["grants"][rid] = (cap, cost, borrowed) + self._ev("grant", rid=rid, cap=cap, borrowed=borrowed) rec["borrowed"] |= borrowed self._tried.add(rid) sched._glm53_align_prefill_limit = cap @@ -651,6 +829,7 @@ def finish_step(self, sched, scheduler_output): self._open_rec = None if rec is None: return + self._trace_flush() scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() if n > 0} requests = getattr(sched, "requests", {}) @@ -672,7 +851,8 @@ def finish_step(self, sched, scheduler_output): # Empty schedules do not necessarily have a completion callback. return rec.update(output=scheduler_output, scheduled=scheduled, - prefill_tokens=prefill, reserved=reserved) + prefill_tokens=prefill, reserved=reserved, + incarn={rid: self._incarn.get(rid) for rid in prefill}) del rec["grants"] if self.mode == "fair": self.inflight[id(scheduler_output)] = rec @@ -694,11 +874,20 @@ def observe_output(self, sched, scheduler_output): served = {rid: min(n, max(0, int(actual.get(rid, 0)))) for rid, n in rec["prefill_tokens"].items() if int(actual.get(rid, 0)) > 0} + stale = {rid for rid in served + if rec.get("incarn", {}).get(rid) != self._incarn.get(rid)} + for rid in stale: + self._ev("stale_completion", step=rec["step_id"], rid=rid) for rid, n in served.items(): + if rid in stale: + continue # a replaced request's late completion earns the new one nothing self.last_service[rid] = now self.rr_n += 1 self.rr_seq[rid] = self.rr_n self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + self._carry.discard(rid) + self._ev("served", step=rec["step_id"], rid=rid, tokens=n, + contention=int(rec["had_decode"])) n = sum(served.values()) if rec["had_decode"]: # Reservation was already debited; settle once, including overruns. @@ -719,6 +908,7 @@ def observe_output(self, sched, scheduler_output): f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) self._prune(self._live_ids(sched)) + self._trace_flush() def _maybe_log(self): if self.hist_every <= 0 or self.steps % self.hist_every != 1: @@ -760,8 +950,8 @@ def _helper_text() -> str: return ( "\n" + body - + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5]\n\n" - + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5]\n" + + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v7]\n\n" + + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v7]\n" + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" ) @@ -975,176 +1165,250 @@ def replace_once(text: str, old: str, new: str, label: str) -> str: return text.replace(old, new, 1) -def _strip_helper(text: str, label: str, *, expected: str | None = None) -> str: - start = text.find("class _Glm53MixedPrefill:") - if start < 0: - raise SystemExit(f"{P}: {label} helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: {label} helper end not found") - end = min(candidates) - if expected is not None and text[start:end].strip() != expected.strip(): - raise SystemExit(f"{P}: {label} helper drifted") - return text[:start] + text[end:] - - -def _strip_v1_helper(text: str) -> str: - start = text.find(V1_HELPER_START) - if start < 0: - raise SystemExit(f"{P}: v1 helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: v1 helper end not found") - end = min(candidates) - return text[:start] + "\n" + text[end:] - - -def unpatch_v1(text: str) -> str: - if V1_RUNNING_NEW in text: - text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 running insertion drifted") - if V1_WAITING_NEW in text: - text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 waiting insertion drifted") - if V1_HELPER_START in text: - text = _strip_v1_helper(text) - leftover = [ - "_glm53_mixed_prefill_policy(self.running, request)", - V1_HELPER_START, - ] - for s in leftover: - if s in text: - raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") - return text +V1_PAIRS = ( + (V1_RUNNING_NEW, RUNNING_OLD, 'v1-running'), + (V1_WAITING_NEW, WAITING_OLD, 'v1-waiting'), +) +V2_PAIRS = ( + (V2_BEGIN_NEW, BEGIN_OLD, 'v2-begin'), + (V2_OBS_NEW, OBS_OLD, 'v2-obs'), + (V2_RUNNING_NEW, RUNNING_OLD, 'v2-running'), + (V2_WAITING_NEW, WAITING_OLD, 'v2-waiting'), + (V2_ALIGN_NEW, ALIGN_OLD, 'v2-align'), + (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'v2-running-mamba'), + (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'v2-waiting-mamba'), +) +# v5/v6 hook text is the frozen v4 insertion with the marker advanced; every +# known v5/v6 producer shares it byte-for-byte, so identities differ only in +# their helper bytes. +V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) +V6_PAIRS = tuple((new.replace(MARK_V4, MARK_V6), old, label) for new, old, label in V4_PAIRS) + +_ALLOC_OLD_CALL = "_GLM53_MIXED.note_scheduled(request, 0) " + MARK_V7 +_ALLOC_NEW_CALL = "_GLM53_MIXED.note_alloc_failed(request) " + MARK_V7 + + +def _v7_pair(new, old, label): + new = new.replace(MARK_V4, MARK_V7) + if label in ("running_alloc", "waiting_alloc"): + # v7 reports native KV allocation refusal distinctly from a zero cap. + if new.count(_ALLOC_OLD_CALL) != 1: + raise SystemExit(f"{P}: v7 {label} template drifted") + new = new.replace(_ALLOC_OLD_CALL, _ALLOC_NEW_CALL, 1) + return new, old, label + + +V7_PAIRS = tuple(_v7_pair(*p) for p in V4_PAIRS) + +NEEDLE = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" +ADAPTIVE_K_HEAD = "class _Glm53AdaptiveK:" +CLASS_HEAD = "class _Glm53MixedPrefill:" + +# --------------------------------------------------------------------------- +# Authenticated migration identities. A helper *site* is the exact text from the +# newline before its class/def head up to the next CUDAGraphStat import or +# _Glm53AdaptiveK class; it is compared by sha256 AND length, then the version's +# exact hook insertions must invert once each, then the removal must round-trip +# byte-identically. Marker presence alone authenticates nothing. +# +# provenance (MiaAI-Lab/GLM-5.3-Flash-EXL3-2x-DGX-Sparks public history): +# v1 f3043c95bbf9 HELPER (as registered by #180 4ffef088) +# v2 9a23b2d8f061 HELPER (as registered by #180 4ffef088) +# v5-historical a9bbbadc4c72 _helper_text() +# v5-main 943912cdcda2 _helper_text() (published TP2/3/4 fair default) +# v5-priority #221 9dbd6873ab45 _helper_text() (in-place priority migration) +# v6-warm-deadline #180 4ffef088edb1 == #80 ff30124ed456 _helper_text() +# v6-carry #246 dc702b4fca1c _helper_text() +# v1-image-d9758a6 d9758a6 == 9a557cf HELPER, the unversioned 2026-09-13/14 +# variant baked into the published :exl3-instanttensor image +# (ImageID ef9f5013c41a). #180 refuses it as a *retained* v1 because +# its baked "0" default would land a held default change. v7 +# accepts it ONLY as a migration source: the helper is removed +# and replaced by v7, whose mode comes from the launcher's +# GLM53_MIXED_PREFILL_CHUNK, so that default does not land. +# Without this entry v7 cannot boot on the published image. +# v7-drop2 #283 DROP-2 patcher f49d85d8 _helper_text() (installed on .123 Sep 27) +# v7-drop4 #283 DROP-4 patcher 19e25753 _helper_text() (installed on .123 Sep 27-28) +# Earlier v7 helper bodies are migration inputs only: a v7 marker is +# "verified" solely by this installer's exact current helper bytes. +# v3/v4 no public producer body: refused, never migrated. +# --------------------------------------------------------------------------- +IDENTITIES = ( + # name, marker, head token, pairs, sha256, length + ("v1", MARK, V1_HELPER_START, V1_PAIRS, + "c0c10c6385bd7e75bfc0f724378960b9d9cb943cf18c386a5a780cf03cb4c48d", 784), + ("v1-image-d9758a6", MARK, V1_HELPER_START, V1_PAIRS, + "9ad0c79991ae3f3f8fa382810dee92c58316a03a27c81774ecd889ad986af29e", 781), + ("v2", MARK_V2, CLASS_HEAD, V2_PAIRS, + "e170f2bf6d0c0fe0493a91ff54bd719036f6c4b5ea9fcf2ee1346968a4e3fbed", 13378), + ("v5-historical", MARK_V5, CLASS_HEAD, V5_PAIRS, + "c05769276df6da0a24b3b8c21e251f39ee7aa2a74ee92c6d6ce3d85c38151a53", 20866), + ("v5-main", MARK_V5, CLASS_HEAD, V5_PAIRS, + "1dbbb945a3414bd12d78b53070424f432a2f07d0092898f9751814f1f9dc40f9", 20866), + ("v5-priority", MARK_V5, CLASS_HEAD, V5_PAIRS, + "b8bf9c9cfcb5e0e62bb3e842c3679c9adc0911c2ce66d2705131ac481a129b55", 21170), + ("v6-warm-deadline", MARK_V6, CLASS_HEAD, V6_PAIRS, + "4eac19b23a39ffa58c765191fd8d120b2d9cd400dcdbfe9019fcea4dee88e03a", 23415), + ("v6-carry", MARK_V6, CLASS_HEAD, V6_PAIRS, + "278b37ad9eea23de047726b054fd7b6929f6bdcfe76e68e21f230e55528bdc15", 22688), + ("v7-drop2", MARK_V7, CLASS_HEAD, V7_PAIRS, + "9ed55941e86f51c82845fb32d207acb58183f4bb55642f79d7babfcd18d2aac5", 27780), + ("v7-drop4", MARK_V7, CLASS_HEAD, V7_PAIRS, + "45e3786bebb7a0e207989130bfa26feeeae2280467531321c7da9ffa6b013d3d", 27988), +) -def unpatch_v2(text: str) -> str: - pairs = ( - (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), - (V2_OBS_NEW, OBS_OLD, "v2-obs"), - (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), - (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), - (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), - (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), - (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), - ) - for new, old, label in pairs: - if new in text: - text = replace_once(text, new, old, label) - elif MARK_V2 in new: - # Some insertions may already have been removed; fail if marker remains. - pass - if "class _Glm53MixedPrefill:" in text: - text = _strip_helper(text, "v2") - if MARK_V2 in text: - raise SystemExit(f"{P}: v2 leftover after unpatch") - return text +def _refuse(reason: str): + raise SystemExit(f"{P}: refusing to rewrite the scheduler: {reason}") -def unpatch_v3(text: str) -> str: - pairs = ( - (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), - (V3_OBS_NEW, OBS_OLD, "v3-obs"), - (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), - (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), - (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), - (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), - (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), - (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), - ) - for new, old, label in pairs: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v3") - if MARK_V3 in text: - raise SystemExit(f"{P}: v3 leftover after unpatch") - return text +def _site(text: str, head_token: str, label: str): + """Return the unique helper site start (the newline before its head), or None.""" + found = text.count(head_token) + if found == 0: + return None + if found != 1: + _refuse(f"{found} {head_token!r} definitions found; {label} requires exactly one") + head = text.find(head_token) + if head < 2 or text[head - 1] != "\n" or text[head - 2] != "\n": + _refuse(f"{label} helper is not preceded by the installer's blank-line boundary") + return head - 1 -def unpatch_v4(text: str) -> str: - for new, old, label in V4_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v4") - if MARK_V4 in text: - raise SystemExit(f"{P}: v4 leftover after unpatch") - return text +def _span_matches(span: str, sha: str, length: int) -> bool: + return len(span) == length and hashlib.sha256(span.encode("utf-8")).hexdigest() == sha -# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 -# insertions above stay frozen so a v4 image can be unpatched exactly. -V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) +def _followed_by_anchor(text: str, end: int) -> bool: + """After the exact helper bytes only blank lines may precede the next known anchor + (the CUDAGraphStat import, or _Glm53AdaptiveK which patch_adaptive_k inserts).""" + # Exactly the established boundaries: the helper directly followed by the + # CUDAGraphStat import, or by patch_adaptive_k's class after its single + # separator newline. Nothing else is tolerated. + rest = text[end:] + return (rest.startswith(NEEDLE) or rest.startswith(ADAPTIVE_K_HEAD) + or rest.startswith("\n" + ADAPTIVE_K_HEAD)) -def unpatch_v5(text: str) -> str: - for new, old, label in V5_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v5", expected=_helper_text()) - if MARK_V5 in text: - raise SystemExit(f"{P}: v5 leftover after unpatch") - return text +def _unpatch(text: str, name, marker, head_token, pairs, sha=None, length=None, exact=None): + """Invert one identity: exact hook insertions, then its authenticated helper site. + + The site is authenticated as exactly `length` bytes (or `exact`) from its start, + by sha256; it must then be followed only by blank lines and a known anchor. + """ + for new, old, label in pairs: + if text.count(new) != 1: + _refuse(f"{name}: {label} insertion missing, drifted or duplicated") + text = text.replace(new, old, 1) + start = _site(text, head_token, name) + if start is None: + _refuse(f"{name}: helper site missing") + n = len(exact) if exact is not None else length + span = text[start:start + n] + ok = (span == exact) if exact is not None else _span_matches(span, sha, length) + if not ok: + _refuse(f"{name}: helper site not authenticated (got sha256 " + f"{hashlib.sha256(span.encode()).hexdigest()[:16]} over {len(span)} bytes)") + if not _followed_by_anchor(text, start + n): + _refuse(f"{name}: unexpected text between the helper and its anchor") + clean = text[:start] + text[start + n:] + if marker in clean or head_token in clean: + _refuse(f"{name}: marker or helper definition left after unpatch") + return clean, start, span + + +def _round_trip(original: str, clean: str, start: int, span: str, pairs, name: str) -> None: + again = clean[:start] + span + clean[start:] + for new, old, label in pairs: + again = replace_once(again, old, new, label) + if again != original: + _refuse(f"{name}: round trip is not byte-identical; not a canonical {name} install") -def apply_v5(text: str) -> str: +def apply_v7(text: str, at: int | None = None) -> str: if "import os\n" not in text.split("import time\n", 1)[0]: text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") - needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" - text = replace_once(text, needle, _helper_text() + needle, "helper") - for new, old, label in V5_PAIRS: + if text.count(CLASS_HEAD) or text.count(V1_HELPER_START): + _refuse("a helper definition is still present before v7 install") + helper = _helper_text() + if at is None: + text = replace_once(text, NEEDLE, helper + NEEDLE, "helper") + else: + # Reinstall at the removed legacy site so later overlays keep their place. + text = text[:at] + helper + text[at:] + for new, old, label in V7_PAIRS: text = replace_once(text, old, new, label) compile(text, str(P), "exec") return text +def _identify(text: str): + """Authenticate the installed legacy identity (or None for a pristine scheduler).""" + for unsupported in (MARK_V3, MARK_V4): + if unsupported in text: + _refuse(f"{unsupported} is unsupported: no authenticated producer body exists; " + "restore an authenticated/pristine scheduler first") + if MARK_V7 in text: + marker = MARK_V7 + elif MARK_V6 in text: + marker = MARK_V6 + elif MARK_V5 in text: + marker = MARK_V5 + elif MARK_V2 in text: + marker = MARK_V2 + elif MARK in text or V1_HELPER_START in text: + marker = MARK + else: + if CLASS_HEAD in text: + _refuse("unmarked _Glm53MixedPrefill helper present") + return None + reasons = [] + for name, mark, head, pairs, sha, length in IDENTITIES: + if mark != marker: + continue + try: + clean, start, span = _unpatch(text, name, mark, head, pairs, sha=sha, length=length) + except SystemExit as exc: + reasons.append(str(exc).split("scheduler: ", 1)[-1]) + continue + _round_trip(text, clean, start, span, pairs, name) + return name, clean, start + _refuse(f"no authenticated identity for {marker}: " + " | ".join(reasons)) + + def main() -> int: if not P.is_file(): raise SystemExit(f"missing {P}") text = P.read_text() original = text - if MARK_V5 in text: - if PRIORITY_MARK not in text: - text = replace_once(text, LEGACY_RANK_PREFIX, PRIORITY_RANK_PREFIX, "prefill priority ranking") - text = replace_once(text, LEGACY_CANDIDATES, PRIORITY_CANDIDATES, "prefill scheduler policy") - # Validate existing anchors/helper instead of trusting the marker alone. - # unpatch_v5 checks every v5 insertion occurs exactly once, strips the - # helper, and rejects leftover markers. The result is discarded: a - # later overlay may legitimately sit between the helper and the - # cuda_graph import anchor, so re-applying at that fixed anchor would - # relocate the helper and fail a byte-compare on a healthy file. - unpatch_v5(text) - # The import edit is part of the applied state; the byte-compare used - # to cover it implicitly. - if "import os\n" not in text.split("import time\n", 1)[0]: - raise SystemExit(f"{P}: v5 import drifted") - compile(text, str(P), "exec") - if text != original: - P.write_text(text) - print(f"{P.name}: migrated prefill-priority ranking ({PRIORITY_MARK})") - else: - print(f"{P.name}: {MARK_V5} already present — verified") - return 0 - if MARK_V4 in text: - text = unpatch_v4(text) - elif MARK_V3 in text: - text = unpatch_v3(text) - elif MARK_V2 in text: - text = unpatch_v2(text) - elif MARK in text or V1_HELPER_START in text: - text = unpatch_v1(text) - text = apply_v5(text) - if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text: - raise SystemExit(f"{P}: older marker left after migration") + if MARK_V7 in text: + # Position-independent verification of the installed v7: exact hooks, + # exact CURRENT helper bytes, byte-identical round trip. No write. + try: + clean, start, span = _unpatch(text, "v7", MARK_V7, CLASS_HEAD, V7_PAIRS, + exact=_helper_text()) + except SystemExit: + clean = None # not the current helper: only an authenticated legacy v7 may migrate + if clean is not None: + _round_trip(text, clean, start, span, V7_PAIRS, "v7") + if "import os\n" not in text.split("import time\n", 1)[0]: + _refuse("v7 import drifted") + compile(text, str(P), "exec") + print(f"{P.name}: {MARK_V7} already present - verified") + return 0 + found = _identify(text) + if found is None: + text = apply_v7(text) + source = "pristine" + else: + source, clean, start = found + text = apply_v7(clean, at=start) + for old in (MARK_V2, MARK_V3, MARK_V4, MARK_V5, MARK_V6): + if old in text: + raise SystemExit(f"{P}: older marker left after migration") if text != original: P.write_text(text) - print(f"patched {P.name} ({MARK_V5})") + print(f"patched {P.name} ({MARK_V7}) from {source}") return 0 diff --git a/tests/check_concurrent_agents.py b/tests/check_concurrent_agents.py new file mode 100644 index 00000000..4beee3df --- /dev/null +++ b/tests/check_concurrent_agents.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""One live concurrency canary: B must answer while A is still generating. + +No sweep or automatic retry. Run after approving the pinned deployment recipe. +Uses only stdlib; OPENAI_API_KEY is optional and is never included in receipts. +""" +from __future__ import annotations + +import argparse +import http.client +import json +import os +from pathlib import Path +import socket +import threading +import time +from urllib.parse import urlsplit +import uuid + + +class Request(threading.Thread): + def __init__(self, base_url, body, timeout): + super().__init__(daemon=True) + self.url, self.body, self.timeout = urlsplit(base_url), body, timeout + self.first, self.done, self.cancelled = threading.Event(), threading.Event(), threading.Event() + self.sock = None + self.started_at = self.first_at = self.last_at = self.finished_at = None + self.status = self.error = self.finish_reason = self.usage = None + self.text = "" + + def cancel(self): + self.cancelled.set() + if self.sock is not None: + try: + self.sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + + def run(self): + self.started_at = time.monotonic() + cls = http.client.HTTPSConnection if self.url.scheme == "https" else http.client.HTTPConnection + conn = cls(self.url.hostname, self.url.port, timeout=min(self.timeout, 30)) + try: + conn.connect() + self.sock = conn.sock + self.sock.settimeout(self.timeout) + if self.cancelled.is_set(): + return + headers = {"Content-Type": "application/json"} + key = os.environ.get("OPENAI_API_KEY") + if key: + headers["Authorization"] = "Bearer " + key + conn.request("POST", self.url.path.rstrip("/") + "/chat/completions", + json.dumps(self.body), headers) + response = conn.getresponse() + self.status = response.status + if response.status != 200: + raise RuntimeError(f"HTTP {response.status}") + while not self.cancelled.is_set(): + line = response.readline() + if not line: + break + if not line.startswith(b"data:"): + continue + payload = line[5:].strip() + if payload == b"[DONE]": + break + chunk = json.loads(payload) + if chunk.get("error"): + raise RuntimeError("provider returned an SSE error") + if chunk.get("usage"): + self.usage = chunk["usage"] + for choice in chunk.get("choices", []): + delta = choice.get("delta") or {} + content = delta.get("content") or "" + if content: + now = time.monotonic() + if self.first_at is None: + self.first_at = now + self.first.set() + self.last_at = now + self.text += content + if choice.get("finish_reason"): + self.finish_reason = choice["finish_reason"] + except Exception as exc: + if not self.cancelled.is_set(): + # Exception type suffices; don't echo a provider body or secrets. + self.error = type(exc).__name__ + finally: + conn.close() + self.finished_at = time.monotonic() + self.done.set() + + def receipt(self): + def elapsed(t): + return None if t is None or self.started_at is None else round(t - self.started_at, 3) + return {"http": self.status, "ttft_s": elapsed(self.first_at), + "wall_s": elapsed(self.finished_at), "finish_reason": self.finish_reason, + "usage": self.usage, "error": self.error, "cancelled": self.cancelled.is_set()} + + +def wait_for_first(req, deadline): + while time.monotonic() < deadline: + if req.first.is_set(): + return True + if req.done.is_set(): + return False + req.first.wait(min(0.1, max(0, deadline - time.monotonic()))) + return False + + +def assess(a, b): + if a.error or b.error or a.cancelled.is_set() or b.cancelled.is_set(): + return False, "a request failed, timed out, or was cancelled" + if a.status != 200 or b.status != 200 or not a.finish_reason or not b.finish_reason: + return False, "both requests must complete valid streams" + if b.text.strip() != "OK": + return False, "short request did not return the expected OK" + if not a.text.lstrip().startswith("1"): + return False, "long request did not begin the requested count" + if a.first_at is None or a.last_at is None or b.started_at is None or b.first_at is None: + return False, "missing content timing" + if not a.first_at <= b.started_at <= b.first_at < a.last_at: + return False, "B must start after A's first token and answer before A's last token" + return True, "B answered and A continued generating afterward" + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default="http://127.0.0.1:8888/v1") + parser.add_argument("--model", default="GLM-5.3-Flash-EXL3") + parser.add_argument("--filler-words", type=int, default=8000) + parser.add_argument("--long-max-tokens", type=int, default=1024) + parser.add_argument("--timeout", type=float, default=240) + parser.add_argument("--out", type=Path) + args = parser.parse_args() + url = urlsplit(args.base_url) + if url.scheme not in ("http", "https") or not url.hostname or url.username or url.query or url.fragment: + parser.error("base URL must be HTTP(S) without credentials, query, or fragment") + if args.filler_words < 1 or args.long_max_tokens < 128 or args.timeout <= 0: + parser.error("positive filler/timeout and at least 128 long-request tokens required") + common = {"model": args.model, "stream": True, "stream_options": {"include_usage": True}, + "temperature": 0, "top_p": 1, "chat_template_kwargs": {"enable_thinking": False}} + unique = uuid.uuid4().hex + a_body = {**common, "max_tokens": args.long_max_tokens, "messages": [{"role": "user", "content": + f"Run {unique}-A. Count from 1 to 100000. Output only the numbers separated by spaces. Do not stop early."}]} + b_body = {**common, "max_tokens": 16, "messages": [{"role": "user", "content": + f"Run {unique}-B. Context: " + "the " * args.filler_words + "\nReply with exactly OK."}]} + a, b = Request(args.base_url, a_body, args.timeout), Request(args.base_url, b_body, args.timeout) + deadline = time.monotonic() + args.timeout + reason = None + try: + a.start() + if not wait_for_first(a, deadline): + reason = "A failed to produce content; B was not submitted" + else: + print("A is generating; starting B with a distinct prompt.", flush=True) + b.start() + while not (a.done.is_set() and b.done.is_set()) and time.monotonic() < deadline: + if a.error or b.error or (b.first_at is None and a.done.is_set()): + reason = "request failure or A finished before B produced content" + break + time.sleep(min(0.1, max(0, deadline - time.monotonic()))) + if not (a.done.is_set() and b.done.is_set()) and reason is None: + reason = "canary deadline exceeded" + except KeyboardInterrupt: + reason = "canary interrupted" + finally: + for req in (a, b): + if req.ident is not None and not req.done.is_set(): + req.cancel() + for req in (a, b): + if req.ident is not None: + req.join(timeout=2) + passed, assessed_reason = assess(a, b) + receipt = {"passed": passed and reason is None, "reason": reason or assessed_reason, + "model": args.model, "filler_words": args.filler_words, + "long_max_tokens": args.long_max_tokens, "a": a.receipt(), "b": b.receipt(), + "b_first_before_a_last_s": None if b.first_at is None or a.last_at is None + else round(a.last_at - b.first_at, 3)} + print(json.dumps(receipt, indent=2)) + if args.out: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(receipt, indent=2) + "\n") + return 0 if receipt["passed"] else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/fixtures/decode_floor_producers/v5_main_d5713ec6.py b/tests/fixtures/decode_floor_producers/v5_main_d5713ec6.py new file mode 100644 index 00000000..78c53778 --- /dev/null +++ b/tests/fixtures/decode_floor_producers/v5_main_d5713ec6.py @@ -0,0 +1,1123 @@ +#!/usr/bin/env python3 +"""Mixed-prefill policy: skip / cap / off / fair (issue #6). + +A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover +MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill +chunk. Mixed execution leaves the uniform decode FULL-graph path, and a +128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond +budget. + +GLM53_MIXED_PREFILL_CHUNK: + skip / -1 — do not mix prefill with decode. Starves every + waiting/running prefill while any peer is decoding; there is + no age limit. Solo prefill is unchanged. + N>0 — cap mixed prefill chunks to N tokens while a peer decodes. + The cap is fed into hybrid Mamba alignment so N < block_size + still makes sub-block progress. 128 still stalls ~10 tok/s. + 0 / off — disable the extra isolation policy. + fair — service-time mixing (default on TP=2/3/4 since 2026-09-15, + v5). Decode-only + steps between prefill turns; at most + GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). + Only prefill that contends with a decoder is charged (solo + prefill is cost-sampled, not debt). Credit accrues at SHARE + of accounted engine time. v5 fits a fixed-plus-per-token step + cost from solo and mixed samples, targets the largest ladder + rung (128..2048) whose estimated step fits + GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that + rung instead of spending it on small chunks (v4 priced 1024 + tokens off 128-token samples linearly, ~3x too high, then + stalled at 128-token steps: ~70 tok/s at 20% share). A + never-served newcomer gets one prompt step-bounded probe; + afterwards a 2s age override may borrow one such chunk, only + after all shared debt is repaid. In-flight async prefill + blocks the next mixed turn. Prefills are selected by last + completed positive prefill service plus round-robin. Decode + token/input budget is allocated first by the base scheduler. + Solo prefill retains the base scheduler's limits. Timing is a + host busy-time proxy. + +Fair knobs (read at runtime; identical on every rank): + GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) + GLM53_FAIR_PREFILL_SHARE default 0.30 (credit accrual fraction) + GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 + GLM53_FAIR_PREFILL_MAX_STEP_MS default 2000 (estimated mixed-step limit) + GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 + +Versioned installer: `# [glm53-decode-floor:v5]`. v1 (no version), v2, v3 +and v4 images are unpatched then re-patched. Fail closed if anchors drift. +""" +from __future__ import annotations + +import inspect +import os +import sys +import time +from pathlib import Path + +P = Path( + os.environ.get( + "GLM53_SCHEDULER_PY", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", + ) +) +MARK = "# [glm53-decode-floor]" +MARK_V2 = "# [glm53-decode-floor:v2]" +MARK_V3 = "# [glm53-decode-floor:v3]" +MARK_V4 = "# [glm53-decode-floor:v4]" +MARK_V5 = "# [glm53-decode-floor:v5]" + +IMPORT_OLD = """import itertools +import time +""" +IMPORT_NEW = """import itertools +import os +import time +""" + +# v1 helper + insertions (recipe f906ee9 / this overlay before v2). +V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" +V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. +V2_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] + # NOTE(woosuk) on the scheduling algorithm: +""" +V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] + pooler_outputs = model_runner_output.pooler_output +""" +V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" +V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] +""" +V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] + if num_new_tokens == 0: + break +""" + + + +# Frozen v3 anchors for migration from the reviewed implementation. +V3_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v3] + # NOTE(woosuk) on the scheduling algorithm: +""" +V3_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v3] + pooler_outputs = model_runner_output.pooler_output +""" +V3_FIN_OLD = """ # Check if the scheduling constraints are satisfied. + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) +""" +V3_FIN_NEW = """ # Check if the scheduling constraints are satisfied. + _GLM53_MIXED.note_schedule_output(self, num_scheduled_tokens) # [glm53-decode-floor:v3] + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) +""" +V3_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( + num_new_tokens, + token_budget, + int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), + _GLM53_MIXED.needs_prefill_compute(request), + ) + + # Make sure the input position does not exceed the max model len. +""" +V3_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( + num_new_tokens, + token_budget, + int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), + _GLM53_MIXED.needs_prefill_compute(request), + ) + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if num_new_tokens <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + # chunked prefill has to be enabled explicitly to allow +""" +V3_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v3] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V3_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] +""" +V3_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] + if num_new_tokens == 0: + break +""" + +class _Glm53MixedPrefill: # [glm53-decode-floor:v5] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) + except ValueError: + self.share = 0.30 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) + except ValueError: + self.max_step_s = 2.0 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self): + grants = (self._open_rec or {}).get("grants", {}) + for r in self._candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + self.defer_reason = reason + self._promote_next() + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + self._promote_next() + self._maybe_log() + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + return 0 if peer_decode else None + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + + + +def _helper_text() -> str: + body = inspect.getsource(_Glm53MixedPrefill) + return ( + "\n" + + body + + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5]\n\n" + + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5]\n" + + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" + ) + + +HELPER = None # filled at apply time so tests can call _helper_text() + + +BEGIN_OLD = """ self.current_step += 1 + # NOTE(woosuk) on the scheduling algorithm: +""" +BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] + # NOTE(woosuk) on the scheduling algorithm: +""" + +OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output +""" +OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] + pooler_outputs = model_runner_output.pooler_output +""" + +RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + + # Make sure the input position does not exceed the max model len. +""" +RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" + +WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + + # chunked prefill has to be enabled explicitly to allow +""" +WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" + +RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) +""" +RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] +""" + +WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if num_new_tokens == 0: + break +""" +WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] + if num_new_tokens == 0: + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" +FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] + with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" + +RUNNING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled because one of the following +""" +RUNNING_ZERO_NEW = """ if num_new_tokens == 0: + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + # The request cannot be scheduled because one of the following +""" + +WAITING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled. + break +""" +WAITING_ZERO_NEW = """ if num_new_tokens == 0: + # The request cannot be scheduled. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. + # Preempt the lowest-priority request. +""" +PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. + if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] + break + # Preempt the lowest-priority request. +""" + +RUNNING_ALLOC_OLD = """ if new_blocks is None: + # Cannot schedule this request. + break +""" +RUNNING_ALLOC_NEW = """ if new_blocks is None: + # Cannot schedule this request. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + req_index += 1 + continue + break +""" + +WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + break + + # KVTransfer:""" +WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break + + # KVTransfer:""" + +V4_PAIRS = ( + (BEGIN_NEW, BEGIN_OLD, 'begin'), + (OBS_NEW, OBS_OLD, 'obs'), + (RUNNING_NEW, RUNNING_OLD, 'running'), + (WAITING_NEW, WAITING_OLD, 'waiting'), + (ALIGN_NEW, ALIGN_OLD, 'align'), + (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), + (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), + (FIN_NEW, FIN_OLD, 'fin'), + (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), + (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), + (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), + (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), + (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), +) + + +def replace_once(text: str, old: str, new: str, label: str) -> str: + n = text.count(old) + if n != 1: + raise SystemExit(f"{P}: expected one {label} target, found {n}") + return text.replace(old, new, 1) + + +def _strip_helper(text: str, label: str, *, expected: str | None = None) -> str: + start = text.find("class _Glm53MixedPrefill:") + if start < 0: + raise SystemExit(f"{P}: {label} helper start not found") + if start > 0 and text[start - 1] == "\n": + start -= 1 + end_ak = text.find("class _Glm53AdaptiveK:", start) + end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) + candidates = [i for i in (end_ak, end_cg) if i > start] + if not candidates: + raise SystemExit(f"{P}: {label} helper end not found") + end = min(candidates) + if expected is not None and text[start:end].strip() != expected.strip(): + raise SystemExit(f"{P}: {label} helper drifted") + return text[:start] + text[end:] + + +def _strip_v1_helper(text: str) -> str: + start = text.find(V1_HELPER_START) + if start < 0: + raise SystemExit(f"{P}: v1 helper start not found") + if start > 0 and text[start - 1] == "\n": + start -= 1 + end_ak = text.find("class _Glm53AdaptiveK:", start) + end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) + candidates = [i for i in (end_ak, end_cg) if i > start] + if not candidates: + raise SystemExit(f"{P}: v1 helper end not found") + end = min(candidates) + return text[:start] + "\n" + text[end:] + + +def unpatch_v1(text: str) -> str: + if V1_RUNNING_NEW in text: + text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") + elif "_glm53_mixed_prefill_policy(self.running, request)" in text: + raise SystemExit(f"{P}: v1 running insertion drifted") + if V1_WAITING_NEW in text: + text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") + elif "_glm53_mixed_prefill_policy(self.running, request)" in text: + raise SystemExit(f"{P}: v1 waiting insertion drifted") + if V1_HELPER_START in text: + text = _strip_v1_helper(text) + leftover = [ + "_glm53_mixed_prefill_policy(self.running, request)", + V1_HELPER_START, + ] + for s in leftover: + if s in text: + raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") + return text + + +def unpatch_v2(text: str) -> str: + pairs = ( + (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), + (V2_OBS_NEW, OBS_OLD, "v2-obs"), + (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), + (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), + (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), + (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), + (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), + ) + for new, old, label in pairs: + if new in text: + text = replace_once(text, new, old, label) + elif MARK_V2 in new: + # Some insertions may already have been removed; fail if marker remains. + pass + if "class _Glm53MixedPrefill:" in text: + text = _strip_helper(text, "v2") + if MARK_V2 in text: + raise SystemExit(f"{P}: v2 leftover after unpatch") + return text + + +def unpatch_v3(text: str) -> str: + pairs = ( + (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), + (V3_OBS_NEW, OBS_OLD, "v3-obs"), + (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), + (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), + (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), + (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), + (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), + (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), + ) + for new, old, label in pairs: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v3") + if MARK_V3 in text: + raise SystemExit(f"{P}: v3 leftover after unpatch") + return text + + +def unpatch_v4(text: str) -> str: + for new, old, label in V4_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v4") + if MARK_V4 in text: + raise SystemExit(f"{P}: v4 leftover after unpatch") + return text + + +# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 +# insertions above stay frozen so a v4 image can be unpatched exactly. +V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) + + +def unpatch_v5(text: str) -> str: + for new, old, label in V5_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v5", expected=_helper_text()) + if MARK_V5 in text: + raise SystemExit(f"{P}: v5 leftover after unpatch") + return text + + +def apply_v5(text: str) -> str: + if "import os\n" not in text.split("import time\n", 1)[0]: + text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") + needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" + text = replace_once(text, needle, _helper_text() + needle, "helper") + for new, old, label in V5_PAIRS: + text = replace_once(text, old, new, label) + compile(text, str(P), "exec") + return text + + +def main() -> int: + if not P.is_file(): + raise SystemExit(f"missing {P}") + text = P.read_text() + original = text + if MARK_V5 in text: + # Validate existing anchors/helper instead of trusting the marker alone. + # unpatch_v5 checks every v5 insertion occurs exactly once, strips the + # helper, and rejects leftover markers. The result is discarded: a + # later overlay may legitimately sit between the helper and the + # cuda_graph import anchor, so re-applying at that fixed anchor would + # relocate the helper and fail a byte-compare on a healthy file. + unpatch_v5(text) + # The import edit is part of the applied state; the byte-compare used + # to cover it implicitly. + if "import os\n" not in text.split("import time\n", 1)[0]: + raise SystemExit(f"{P}: v5 import drifted") + compile(text, str(P), "exec") + print(f"{P.name}: {MARK_V5} already present — verified") + return 0 + if MARK_V4 in text: + text = unpatch_v4(text) + elif MARK_V3 in text: + text = unpatch_v3(text) + elif MARK_V2 in text: + text = unpatch_v2(text) + elif MARK in text or V1_HELPER_START in text: + text = unpatch_v1(text) + text = apply_v5(text) + if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text: + raise SystemExit(f"{P}: older marker left after migration") + if text != original: + P.write_text(text) + print(f"patched {P.name} ({MARK_V5})") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/fixtures/decode_floor_producers/v5_priority_763bd30e.py b/tests/fixtures/decode_floor_producers/v5_priority_763bd30e.py new file mode 100644 index 00000000..1019e530 --- /dev/null +++ b/tests/fixtures/decode_floor_producers/v5_priority_763bd30e.py @@ -0,0 +1,1152 @@ +#!/usr/bin/env python3 +"""Mixed-prefill policy: skip / cap / off / fair (issue #6). + +A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover +MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill +chunk. Mixed execution leaves the uniform decode FULL-graph path, and a +128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond +budget. + +GLM53_MIXED_PREFILL_CHUNK: + skip / -1 — do not mix prefill with decode. Starves every + waiting/running prefill while any peer is decoding; there is + no age limit. Solo prefill is unchanged. + N>0 — cap mixed prefill chunks to N tokens while a peer decodes. + The cap is fed into hybrid Mamba alignment so N < block_size + still makes sub-block progress. 128 still stalls ~10 tok/s. + 0 / off — disable the extra isolation policy. + fair — service-time mixing (default on TP=2/3/4 since 2026-09-15, + v5). Decode-only + steps between prefill turns; at most + GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). + Only prefill that contends with a decoder is charged (solo + prefill is cost-sampled, not debt). Credit accrues at SHARE + of accounted engine time. v5 fits a fixed-plus-per-token step + cost from solo and mixed samples, targets the largest ladder + rung (128..2048) whose estimated step fits + GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that + rung instead of spending it on small chunks (v4 priced 1024 + tokens off 128-token samples linearly, ~3x too high, then + stalled at 128-token steps: ~70 tok/s at 20% share). A + never-served newcomer gets one prompt step-bounded probe; + afterwards a 2s age override may borrow one such chunk, only + after all shared debt is repaid. In-flight async prefill + blocks the next mixed turn. Prefills are selected by last + completed positive prefill service plus round-robin within each + request-priority tier when scheduling-policy=priority. Decode + token/input budget is allocated first by the base scheduler. + Solo prefill retains the base scheduler's limits. Timing is a + host busy-time proxy. + +Fair knobs (read at runtime; identical on every rank): + GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) + GLM53_FAIR_PREFILL_SHARE default 0.30 (credit accrual fraction) + GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 + GLM53_FAIR_PREFILL_MAX_STEP_MS default 2000 (estimated mixed-step limit) + GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 + +Versioned installer: `# [glm53-decode-floor:v5]`. v1 (no version), v2, v3 +and v4 images are unpatched then re-patched. Fail closed if anchors drift. +""" +from __future__ import annotations + +import inspect +import os +import sys +import time +from pathlib import Path + +P = Path( + os.environ.get( + "GLM53_SCHEDULER_PY", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", + ) +) +MARK = "# [glm53-decode-floor]" +MARK_V2 = "# [glm53-decode-floor:v2]" +MARK_V3 = "# [glm53-decode-floor:v3]" +MARK_V4 = "# [glm53-decode-floor:v4]" +MARK_V5 = "# [glm53-decode-floor:v5]" +PRIORITY_MARK = "# [glm53-prefill-priority:v1]" +LEGACY_RANK_PREFIX = """ def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), +""" +PRIORITY_RANK_PREFIX = """ def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + # [glm53-prefill-priority:v1] Preserve fairness within priority tiers. + int(getattr(r, "priority", 0) or 0) if getattr(self, "priority_mode", False) else 0, + self.last_service.get(r.request_id, 0.0), +""" +LEGACY_CANDIDATES = """ self._candidates = self._rank_prefills(prefills) +""" +PRIORITY_CANDIDATES = """ policy = getattr(sched, "policy", None) + self.priority_mode = getattr(policy, "value", policy) == "priority" + self._candidates = self._rank_prefills(prefills) +""" + +IMPORT_OLD = """import itertools +import time +""" +IMPORT_NEW = """import itertools +import os +import time +""" + +# v1 helper + insertions (recipe f906ee9 / this overlay before v2). +V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" +V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. +V2_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] + # NOTE(woosuk) on the scheduling algorithm: +""" +V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] + pooler_outputs = model_runner_output.pooler_output +""" +V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" +V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] +""" +V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] + if num_new_tokens == 0: + break +""" + + + +# Frozen v3 anchors for migration from the reviewed implementation. +V3_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v3] + # NOTE(woosuk) on the scheduling algorithm: +""" +V3_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v3] + pooler_outputs = model_runner_output.pooler_output +""" +V3_FIN_OLD = """ # Check if the scheduling constraints are satisfied. + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) +""" +V3_FIN_NEW = """ # Check if the scheduling constraints are satisfied. + _GLM53_MIXED.note_schedule_output(self, num_scheduled_tokens) # [glm53-decode-floor:v3] + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) +""" +V3_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( + num_new_tokens, + token_budget, + int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), + _GLM53_MIXED.needs_prefill_compute(request), + ) + + # Make sure the input position does not exceed the max model len. +""" +V3_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( + num_new_tokens, + token_budget, + int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), + _GLM53_MIXED.needs_prefill_compute(request), + ) + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if num_new_tokens <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + # chunked prefill has to be enabled explicitly to allow +""" +V3_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v3] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V3_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] +""" +V3_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] + if num_new_tokens == 0: + break +""" + +class _Glm53MixedPrefill: # [glm53-decode-floor:v5] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) + except ValueError: + self.share = 0.30 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) + except ValueError: + self.max_step_s = 2.0 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + # [glm53-prefill-priority:v1] Preserve fairness within priority tiers. + int(getattr(r, "priority", 0) or 0) if getattr(self, "priority_mode", False) else 0, + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self): + grants = (self._open_rec or {}).get("grants", {}) + for r in self._candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + self.defer_reason = reason + self._promote_next() + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + policy = getattr(sched, "policy", None) + self.priority_mode = getattr(policy, "value", policy) == "priority" + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + self._promote_next() + self._maybe_log() + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + return 0 if peer_decode else None + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + + + +def _helper_text() -> str: + body = inspect.getsource(_Glm53MixedPrefill) + return ( + "\n" + + body + + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5]\n\n" + + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5]\n" + + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" + ) + + +HELPER = None # filled at apply time so tests can call _helper_text() + + +BEGIN_OLD = """ self.current_step += 1 + # NOTE(woosuk) on the scheduling algorithm: +""" +BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] + # NOTE(woosuk) on the scheduling algorithm: +""" + +OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output +""" +OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] + pooler_outputs = model_runner_output.pooler_output +""" + +RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + + # Make sure the input position does not exceed the max model len. +""" +RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" + +WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + + # chunked prefill has to be enabled explicitly to allow +""" +WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" + +RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) +""" +RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] +""" + +WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if num_new_tokens == 0: + break +""" +WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] + if num_new_tokens == 0: + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" +FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] + with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" + +RUNNING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled because one of the following +""" +RUNNING_ZERO_NEW = """ if num_new_tokens == 0: + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + # The request cannot be scheduled because one of the following +""" + +WAITING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled. + break +""" +WAITING_ZERO_NEW = """ if num_new_tokens == 0: + # The request cannot be scheduled. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. + # Preempt the lowest-priority request. +""" +PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. + if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] + break + # Preempt the lowest-priority request. +""" + +RUNNING_ALLOC_OLD = """ if new_blocks is None: + # Cannot schedule this request. + break +""" +RUNNING_ALLOC_NEW = """ if new_blocks is None: + # Cannot schedule this request. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + req_index += 1 + continue + break +""" + +WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + break + + # KVTransfer:""" +WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break + + # KVTransfer:""" + +V4_PAIRS = ( + (BEGIN_NEW, BEGIN_OLD, 'begin'), + (OBS_NEW, OBS_OLD, 'obs'), + (RUNNING_NEW, RUNNING_OLD, 'running'), + (WAITING_NEW, WAITING_OLD, 'waiting'), + (ALIGN_NEW, ALIGN_OLD, 'align'), + (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), + (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), + (FIN_NEW, FIN_OLD, 'fin'), + (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), + (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), + (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), + (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), + (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), +) + + +def replace_once(text: str, old: str, new: str, label: str) -> str: + n = text.count(old) + if n != 1: + raise SystemExit(f"{P}: expected one {label} target, found {n}") + return text.replace(old, new, 1) + + +def _strip_helper(text: str, label: str, *, expected: str | None = None) -> str: + start = text.find("class _Glm53MixedPrefill:") + if start < 0: + raise SystemExit(f"{P}: {label} helper start not found") + if start > 0 and text[start - 1] == "\n": + start -= 1 + end_ak = text.find("class _Glm53AdaptiveK:", start) + end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) + candidates = [i for i in (end_ak, end_cg) if i > start] + if not candidates: + raise SystemExit(f"{P}: {label} helper end not found") + end = min(candidates) + if expected is not None and text[start:end].strip() != expected.strip(): + raise SystemExit(f"{P}: {label} helper drifted") + return text[:start] + text[end:] + + +def _strip_v1_helper(text: str) -> str: + start = text.find(V1_HELPER_START) + if start < 0: + raise SystemExit(f"{P}: v1 helper start not found") + if start > 0 and text[start - 1] == "\n": + start -= 1 + end_ak = text.find("class _Glm53AdaptiveK:", start) + end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) + candidates = [i for i in (end_ak, end_cg) if i > start] + if not candidates: + raise SystemExit(f"{P}: v1 helper end not found") + end = min(candidates) + return text[:start] + "\n" + text[end:] + + +def unpatch_v1(text: str) -> str: + if V1_RUNNING_NEW in text: + text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") + elif "_glm53_mixed_prefill_policy(self.running, request)" in text: + raise SystemExit(f"{P}: v1 running insertion drifted") + if V1_WAITING_NEW in text: + text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") + elif "_glm53_mixed_prefill_policy(self.running, request)" in text: + raise SystemExit(f"{P}: v1 waiting insertion drifted") + if V1_HELPER_START in text: + text = _strip_v1_helper(text) + leftover = [ + "_glm53_mixed_prefill_policy(self.running, request)", + V1_HELPER_START, + ] + for s in leftover: + if s in text: + raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") + return text + + +def unpatch_v2(text: str) -> str: + pairs = ( + (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), + (V2_OBS_NEW, OBS_OLD, "v2-obs"), + (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), + (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), + (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), + (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), + (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), + ) + for new, old, label in pairs: + if new in text: + text = replace_once(text, new, old, label) + elif MARK_V2 in new: + # Some insertions may already have been removed; fail if marker remains. + pass + if "class _Glm53MixedPrefill:" in text: + text = _strip_helper(text, "v2") + if MARK_V2 in text: + raise SystemExit(f"{P}: v2 leftover after unpatch") + return text + + +def unpatch_v3(text: str) -> str: + pairs = ( + (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), + (V3_OBS_NEW, OBS_OLD, "v3-obs"), + (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), + (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), + (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), + (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), + (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), + (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), + ) + for new, old, label in pairs: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v3") + if MARK_V3 in text: + raise SystemExit(f"{P}: v3 leftover after unpatch") + return text + + +def unpatch_v4(text: str) -> str: + for new, old, label in V4_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v4") + if MARK_V4 in text: + raise SystemExit(f"{P}: v4 leftover after unpatch") + return text + + +# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 +# insertions above stay frozen so a v4 image can be unpatched exactly. +V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) + + +def unpatch_v5(text: str) -> str: + for new, old, label in V5_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v5", expected=_helper_text()) + if MARK_V5 in text: + raise SystemExit(f"{P}: v5 leftover after unpatch") + return text + + +def apply_v5(text: str) -> str: + if "import os\n" not in text.split("import time\n", 1)[0]: + text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") + needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" + text = replace_once(text, needle, _helper_text() + needle, "helper") + for new, old, label in V5_PAIRS: + text = replace_once(text, old, new, label) + compile(text, str(P), "exec") + return text + + +def main() -> int: + if not P.is_file(): + raise SystemExit(f"missing {P}") + text = P.read_text() + original = text + if MARK_V5 in text: + if PRIORITY_MARK not in text: + text = replace_once(text, LEGACY_RANK_PREFIX, PRIORITY_RANK_PREFIX, "prefill priority ranking") + text = replace_once(text, LEGACY_CANDIDATES, PRIORITY_CANDIDATES, "prefill scheduler policy") + # Validate existing anchors/helper instead of trusting the marker alone. + # unpatch_v5 checks every v5 insertion occurs exactly once, strips the + # helper, and rejects leftover markers. The result is discarded: a + # later overlay may legitimately sit between the helper and the + # cuda_graph import anchor, so re-applying at that fixed anchor would + # relocate the helper and fail a byte-compare on a healthy file. + unpatch_v5(text) + # The import edit is part of the applied state; the byte-compare used + # to cover it implicitly. + if "import os\n" not in text.split("import time\n", 1)[0]: + raise SystemExit(f"{P}: v5 import drifted") + compile(text, str(P), "exec") + if text != original: + P.write_text(text) + print(f"{P.name}: migrated prefill-priority ranking ({PRIORITY_MARK})") + else: + print(f"{P.name}: {MARK_V5} already present — verified") + return 0 + if MARK_V4 in text: + text = unpatch_v4(text) + elif MARK_V3 in text: + text = unpatch_v3(text) + elif MARK_V2 in text: + text = unpatch_v2(text) + elif MARK in text or V1_HELPER_START in text: + text = unpatch_v1(text) + text = apply_v5(text) + if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text: + raise SystemExit(f"{P}: older marker left after migration") + if text != original: + P.write_text(text) + print(f"patched {P.name} ({MARK_V5})") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/fixtures/legacy_scheduler_helpers.py b/tests/fixtures/legacy_scheduler_helpers.py new file mode 100644 index 00000000..558097f9 --- /dev/null +++ b/tests/fixtures/legacy_scheduler_helpers.py @@ -0,0 +1,904 @@ +#!/usr/bin/env python3 +"""Published helper texts of the legacy mixed-prefill installer versions. + +Byte-exact artefacts, recovered from the public history of +MiaAI-Lab/GLM-5.3-Flash-EXL3-2x-DGX-Sparks and used as migration fixtures. +Do not reformat: the installer's registry keys on the sha256 of these strings. + + v1 f3043c95bbf9 overlay/patch_scheduler_decode_floor.py:HELPER + cross-attested byte-identical by cad980f7263:HELPER_V1, which documents + itself as "byte-identical to the shipped v1 overlay" + v2 9a23b2d8f061 overlay/patch_scheduler_decode_floor.py:HELPER + v5 a9bbbadc4c72 overlay/patch_scheduler_decode_floor.py:_helper_text() + +v3 and v4 are absent on purpose: 180725a5ce33 introduced their markers as +migration targets for intermediate builds whose helper bodies were never +recovered from public history, so the installer refuses those markers without +touching the source instead of migrating on a guessed body. +""" +from __future__ import annotations + +HELPERS = {} + +HELPERS[1] = r''' +def _glm53_mixed_prefill_policy(running, current): + """Mixed-step prefill policy when a peer in `running` is decoding. + + None = no extra policy. 0 = skip this prefill this step. N>0 = cap. + """ + raw = os.environ.get("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + if raw in ("0", "off", "no"): + return None + if raw in ("skip", "-1"): + cap = 0 + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + return None + cur_id = getattr(current, "request_id", None) + for r in running: + if r is current or getattr(r, "request_id", None) == cur_id: + continue + if r.num_computed_tokens >= r.num_prompt_tokens: + return cap + return None + + +''' + +HELPERS[2] = r''' +class _Glm53MixedPrefill: # [glm53-decode-floor:v2] + """Skip / cap / off / fair mixed-prefill policy. CPU-only; no GPU sync.""" + + def __init__(self, now=None) -> None: + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.chunk = 256 + self.share = 0.20 + self.interval_s = 2.0 + self.max_chunks = 1 + self.hist_every = 50 + self.logged_boot = False + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.rr_n = 0 + self.step_prefill_ids = set() + self.inflight = [] + self.inflight_mixed = 0 + self.history = [] + self.last_schedule_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self.step_mode = "solo" + self.selected = set() + self.defer_reason = "none" + self.granted_this_step = 0 + self.steps = 0 + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.20")) + except ValueError: + self.share = 0.20 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair chunk={self.chunk} share={self.share} " + f"interval_s={self.interval_s} max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + @staticmethod + def prefill_remaining(request) -> int: + """Tokens still requiring prefill/replay compute (not the decode row).""" + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + num_tokens = int(getattr(request, "num_tokens", prompt) or prompt) + if computed < prompt: + return prompt - computed + prefill_end = max(prompt, num_tokens - 1 if num_tokens > 0 else prompt) + return max(0, prefill_end - computed) + + def needs_prefill_compute(self, request) -> bool: + return self.prefill_remaining(request) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + def _recent_share(self, now): + window = 8.0 + while self.history and now - self.history[0][0] > window: + self.history.pop(0) + prefill_s = sum(dt for _, dt, mix in self.history if mix) + decode_s = sum(dt for _, dt, mix in self.history if not mix) + total = prefill_s + decode_s + if total <= 0: + return 0.0, 0.0, 0.0 + return prefill_s / total, prefill_s, decode_s + + def _oldest_wait(self, now, prefills): + oldest = 0.0 + for r in prefills: + rid = getattr(r, "request_id", None) + if rid is None: + continue + served = self.last_service.get(rid) + if served is None: + oldest = max(oldest, now - self.arrival.get(rid, now)) + else: + oldest = max(oldest, now - served) + return oldest + + def _select(self, prefills): + ranked = [] + for r in prefills: + rid = getattr(r, "request_id", None) + if rid is None: + continue + ranked.append((self.last_service.get(rid, 0.0), self.rr_seq.get(rid, 0), rid, r)) + ranked.sort(key=lambda x: (x[0], x[1], x[2])) + out = [] + for _, _, rid, r in ranked[: self.max_chunks]: + out.append(r) + return out + + def _flush_prior_schedule(self) -> None: + if self.last_schedule_mono is None: + return + ids = set(self.step_prefill_ids) + self.inflight.append(ids) + if ids: + self.inflight_mixed += 1 + self.step_prefill_ids = set() + self.last_schedule_mono = None + + def begin_step(self, sched) -> None: + now = self._now() + sid = id(sched) + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self.inflight_mixed = 0 + self.step_prefill_ids = set() + self.last_schedule_mono = None + self._sched_id = sid + tag = (sid, int(getattr(sched, "current_step", 0) or 0)) + if tag == self.step_tag: + return + self._flush_prior_schedule() + self.step_tag = tag + self.steps += 1 + self.granted_this_step = 0 + self.selected = set() + self.step_prefill_ids = set() + self.last_schedule_mono = now + sched._glm53_align_prefill_limit = None + live = self._live_ids(sched) + self._prune(live) + running = list(getattr(sched, "running", None) or []) + waiting = list(self._iter_waiting(sched)) + for r in running + waiting: + rid = getattr(r, "request_id", None) + if rid is not None and rid not in self.arrival: + self.arrival[rid] = now + prefills = [r for r in running + waiting if self.needs_prefill_compute(r)] + decodes = [r for r in running if not self.needs_prefill_compute(r)] + if self.mode != "fair": + self.step_mode = "legacy" + self.defer_reason = "none" + return + if not prefills: + self.step_mode = "solo" + self.defer_reason = "none" + return + if not decodes: + self.step_mode = "solo" + self.defer_reason = "none" + return + if self.inflight_mixed > 0: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + self._maybe_log(now, len(prefills), len(decodes)) + return + share, _, _ = self._recent_share(now) + since = now - self.last_prefill_turn_mono if self.last_prefill_turn_mono else self.interval_s + starved = self._oldest_wait(now, prefills) >= self.interval_s + want = ( + self.last_prefill_turn_mono == 0.0 + or share < self.share + or since >= self.interval_s + or starved + ) + if not want: + self.step_mode = "decode_only" + self.defer_reason = "share" + self._maybe_log(now, len(prefills), len(decodes)) + return + chosen = self._select(prefills) + self.selected = {getattr(r, "request_id", None) for r in chosen} + self.selected.discard(None) + self.step_mode = "prefill_turn" + self.defer_reason = "none" + self._maybe_log(now, len(prefills), len(decodes)) + + def _maybe_log(self, now, n_prefill, n_decode) -> None: + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + share, pre_s, dec_s = self._recent_share(now) + print( + f"[glm53-decode-floor] step={self.steps} mode={self.step_mode} " + f"defer={self.defer_reason} inflight={self.inflight_mixed} " + f"prefill={n_prefill} decode={n_decode} selected={len(self.selected)} " + f"share={share:.3f} pre_s={pre_s:.3f} dec_s={dec_s:.3f}", + flush=True, + ) + + def cap_for(self, sched, request): + self.begin_step(sched) + if not self.needs_prefill_compute(request): + sched._glm53_align_prefill_limit = None + return None + if self.mode == "off": + sched._glm53_align_prefill_limit = None + return None + running = list(getattr(sched, "running", None) or []) + rid = getattr(request, "request_id", None) + peer_decode = False + for r in running: + if r is request or getattr(r, "request_id", None) == rid: + continue + if not self.needs_prefill_compute(r): + peer_decode = True + break + if self.mode == "skip": + cap = 0 if peer_decode else None + elif self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + else: + # fair + if self.step_mode == "solo": + cap = None + elif self.step_mode != "prefill_turn": + cap = 0 + elif rid not in self.selected: + cap = 0 + elif self.granted_this_step >= self.max_chunks: + cap = 0 + else: + cap = self.chunk + if cap is not None and cap > 0: + sched._glm53_align_prefill_limit = cap + self.step_prefill_ids.add(rid) + self.granted_this_step += 1 + else: + sched._glm53_align_prefill_limit = None + return cap + + def note_scheduled(self, request, num_new_tokens: int) -> None: + """Keep step_prefill_ids as actual post-alignment prefill tokens only.""" + rid = getattr(request, "request_id", None) + if rid is None: + return + if num_new_tokens <= 0 or not self.needs_prefill_compute(request): + self.step_prefill_ids.discard(rid) + return + self.step_prefill_ids.add(rid) + + def observe_output(self, sched, scheduler_output) -> None: + now = self._now() + num_scheduled = getattr(scheduler_output, "num_scheduled_tokens", None) or {} + if self.inflight: + ids = self.inflight.pop(0) + if ids: + self.inflight_mixed = max(0, self.inflight_mixed - 1) + elif self.last_schedule_mono is not None: + ids = set(self.step_prefill_ids) + self.step_prefill_ids = set() + self.last_schedule_mono = None + else: + ids = set() + served = set() + for rid in ids: + n = int(num_scheduled.get(rid, 0) or 0) + if n > 0: + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + served.add(rid) + mixed = bool(served) + start = getattr(self, "_observe_prev", None) + dt = 0.0 if start is None else max(0.0, now - start) + self._observe_prev = now + if dt > 0: + self.history.append((now, dt, mixed)) + if mixed: + self.last_prefill_turn_mono = now + live = self._live_ids(sched) + self._prune(live) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + + +_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v2] + + +def _glm53_mixed_prefill_policy(sched, request): # [glm53-decode-floor:v2] + return _GLM53_MIXED.cap_for(sched, request) + + +''' + +HELPERS[5] = r''' +class _Glm53MixedPrefill: # [glm53-decode-floor:v5] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.20")) + except ValueError: + self.share = 0.20 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "1000")) / 1000.0) + except ValueError: + self.max_step_s = 1.0 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self): + grants = (self._open_rec or {}).get("grants", {}) + for r in self._candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + self.defer_reason = reason + self._promote_next() + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + self._promote_next() + self._maybe_log() + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + return 0 if peer_decode else None + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + +_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5] + +def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5] + return _GLM53_MIXED.cap_for(sched, request, computed) + + +''' diff --git a/tests/test_mamba_align_chunking.py b/tests/test_mamba_align_chunking.py index fb2ba7db..ee296fce 100755 --- a/tests/test_mamba_align_chunking.py +++ b/tests/test_mamba_align_chunking.py @@ -134,10 +134,17 @@ def part_a(src: Path) -> None: if MARK not in original: with tempfile.TemporaryDirectory() as tmp: old = Path(tmp) / "scheduler.py" - unsupported = original.replace("# [glm53-decode-floor:v5]", "# [glm53-decode-floor:v4]") + # The decode-floor overlay in the source is v5 or v7 (#283); downgrade whichever is present. + unsupported = original.replace("# [glm53-decode-floor:v5]", "# [glm53-decode-floor:v4]").replace( + "# [glm53-decode-floor:v7]", "# [glm53-decode-floor:v4]") old.write_text(unsupported) r3 = apply(old) check(r3.returncode != 0 and old.read_text() == unsupported, "A4 an older decode-floor version is rejected without writing") + if "# [glm53-decode-floor:v7]" in original: + v7 = Path(tmp) / "scheduler-v7.py" + v7.write_text(original) + r4 = apply(v7) + check(r4.returncode == 0 and MARK in v7.read_text(), "A5 a decode-floor v7 (#283) source is accepted") else: print(" note: source already carries this overlay; A4 requires an unpatched scheduler") patch = runpy.run_path(str(PATCH)) diff --git a/tests/test_prefill_concurrency.py b/tests/test_prefill_concurrency.py new file mode 100644 index 00000000..0f2c426d --- /dev/null +++ b/tests/test_prefill_concurrency.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 +"""CPU regressions for prompt progress while a peer keeps generating. + +These cases exercise the policy text the overlay actually installs +(``_helper_text()``), so a drifted copy cannot pass. Patch integrity, version +migration, idempotence and refusal live in ``tests/test_scheduler_decode_floor.py``; +set ``GLM53_SCHEDULER_PY_SRC`` to an installed ``scheduler.py`` to run that +matrix against the deployed source as well. +""" +from __future__ import annotations + +import contextlib +import io +import os +from pathlib import Path +import importlib.util +import sys +from types import SimpleNamespace +import unittest +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +PATCH = ROOT / "overlay" / "patch_scheduler_decode_floor.py" +spec = importlib.util.spec_from_file_location("glm53_prefill_v6", PATCH) +OVERLAY = importlib.util.module_from_spec(spec) +sys.modules[spec.name] = OVERLAY +spec.loader.exec_module(OVERLAY) +HELPER_NS: dict = {"os": os, "time": __import__("time")} +exec(OVERLAY._helper_text(), HELPER_NS) +POLICY = HELPER_NS["_Glm53MixedPrefill"] +MIXED_KEYS = ("GLM53_MIXED_PREFILL_CHUNK", "GLM53_MIXED_PREFILL_WARM_TOKENS", + "GLM53_MIXED_PREFILL_MAX_WAIT_MS", "GLM53_MIXED_PREFILL_LATE_CAP") + + +class Req: + def __init__(self, rid, prompt, computed=0, decode=False): + self.request_id = rid + self.num_prompt_tokens = prompt + self.num_computed_tokens = computed + self.num_tokens = prompt + int(decode) + self.spec_token_ids = list(range(7)) if decode else [] + self.num_output_placeholders = 0 + self.has_encoder_inputs = False + self.is_prefill_chunk = not decode + + @property + def num_tokens_with_spec(self): + return self.num_tokens + len(self.spec_token_ids) + + +class Sched: + def __init__(self, running, waiting): + self.running = list(running) + self.waiting = list(waiting) + self.skipped_waiting = [] + self.current_step = 1 + self.max_model_len = 850000 + self.num_sampled_tokens_per_step = 1 + self.need_mamba_block_aligned_split = False + self.scheduler_config = SimpleNamespace(long_prefill_token_threshold=3584) + self.requests = {r.request_id: r for r in self.running + self.waiting} + + +class PrefillProgressTests(unittest.TestCase): + def setUp(self): + self.env = patch.dict(os.environ) + self.env.start() + self.addCleanup(self.env.stop) + for key in MIXED_KEYS: + os.environ.pop(key, None) + self.peer = Req("long-generation", 512, 512, decode=True) + self.newcomer = Req("new-agent", 1200) + self.sched = Sched([self.peer], [self.newcomer]) + + def policy(self, **env): + with patch.dict(os.environ, env), contextlib.redirect_stdout(io.StringIO()): + p = POLICY(now=lambda: 1000.0) + p.hist_every = 0 + return p + + def run_steps(self, policy, steps=64, step_budget=1024, request=None): + """Offer the newcomer the remaining step budget, honouring the policy cap.""" + request = request or self.newcomer + for _ in range(steps): + cap = policy.cap_for(self.sched, request) + if cap is None or cap > 0: + budget = step_budget - 8 # the peer's own decode row + if cap is not None: + budget = min(budget, cap) + request.num_computed_tokens += min( + budget, request.num_prompt_tokens - request.num_computed_tokens) + self.sched.current_step += 1 + return request.num_computed_tokens + + def test_stock_modes_let_the_newcomer_finish_while_the_peer_decodes(self): + # 0 / off / no = no extra isolation: the base scheduler keeps its budget. + for mode in ("0", "off", "no"): + with self.subTest(mode=mode): + self.newcomer.num_computed_tokens = 0 + self.assertEqual(self.run_steps(self.policy(GLM53_MIXED_PREFILL_CHUNK=mode)), + self.newcomer.num_prompt_tokens) + + def test_skip_hold_defers_the_newcomer_until_the_peer_stops(self): + # `skip` gives the newcomer zero tokens while any peer decodes. + newcomer = self.newcomer + self.assertEqual(self.run_steps(self.policy(GLM53_MIXED_PREFILL_CHUNK="skip")), 0) + self.sched.running = [] + self.sched.requests = {newcomer.request_id: newcomer} + self.assertEqual(self.run_steps(self.policy(GLM53_MIXED_PREFILL_CHUNK="skip")), + newcomer.num_prompt_tokens) + + def test_bounded_chunk_grants_the_cap_every_step_under_the_peer(self): + # A positive cap is a per-step eligibility bound while a peer decodes, + # not a service-time guarantee: every step must offer the newcomer + # exactly min(cap, remaining) tokens (nothing here measures wall time), + # so 1200 tokens complete in ceil(1200/128) steps rather than waiting + # for the peer to stop. + p = self.policy(GLM53_MIXED_PREFILL_CHUNK="128") + newcomer = self.newcomer + granted = [] + for _ in range(10): + before = newcomer.num_computed_tokens + remaining = newcomer.num_prompt_tokens - before + self.assertGreater(remaining, 0) + self.run_steps(p, steps=1) + granted.append(newcomer.num_computed_tokens - before) + self.assertEqual(granted, [128] * 9 + [48]) + self.assertEqual(newcomer.num_computed_tokens, newcomer.num_prompt_tokens) + + def test_solo_prefill_is_never_capped(self): + solo = Req("solo", 900) + self.sched.running, self.sched.waiting = [solo], [] + self.sched.requests = {solo.request_id: solo} + for mode in ("skip", "128"): + with self.subTest(mode=mode): + solo.num_computed_tokens = 0 + self.assertEqual(self.run_steps(self.policy(GLM53_MIXED_PREFILL_CHUNK=mode), + steps=1, request=solo), solo.num_prompt_tokens) + + def test_request_cannot_count_itself_as_a_decoding_peer(self): + # Only a DIFFERENT request that needs no prefill is a peer: a request that + # appears in `running` must not be held by its own row. + self.sched.running, self.sched.waiting = [self.newcomer], [] + self.sched.requests = {self.newcomer.request_id: self.newcomer} + self.newcomer.num_computed_tokens = 600 + p = self.policy(GLM53_MIXED_PREFILL_CHUNK="skip") + self.assertIsNone(p.cap_for(self.sched, self.newcomer)) + self.assertEqual(self.run_steps(p, steps=1), self.newcomer.num_prompt_tokens) + + +if __name__ == "__main__": + sys.exit(0 if unittest.main(exit=False).result.wasSuccessful() else 1) diff --git a/tests/test_scheduler_decode_floor.py b/tests/test_scheduler_decode_floor.py index 51de851f..401251d2 100644 --- a/tests/test_scheduler_decode_floor.py +++ b/tests/test_scheduler_decode_floor.py @@ -447,71 +447,89 @@ def test_finished_decode_does_not_hold_phantom_input_reservation(self): def installation_tests(): + """Installer contract for decode-floor v7 (#283; integrates #246/#221/#180). + + Deployed and fixture producers migrate to the same v7 bytes as a fresh install, + a second run is a verified no-op, and unsupported (v3/v4), drifted, duplicated, + marker-only or unmarked inputs are refused without writing (#180's fail-closed + contract). This is a self-contained subset of the 13-producer / 14-refusal + matrix run for #283 (summarized in the #283 integration PR). + """ src = next((p for p in [Path(os.environ.get('GLM53_SCHEDULER_PY_SRC', '/missing')), Path('/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py'), Path('/tmp/sched-live.py')] if p.is_file()), None) if src is None: raise SystemExit('Set GLM53_SCHEDULER_PY_SRC to the pinned scheduler source') - clean = src.read_text() - for marker, fn in [(mod.MARK_V5, mod.unpatch_v5), (mod.MARK_V4, mod.unpatch_v4), (mod.MARK_V3, mod.unpatch_v3), (mod.MARK_V2, mod.unpatch_v2)]: - if marker in clean: - clean = fn(clean) - if mod.V1_HELPER_START in clean: - clean = mod.unpatch_v1(clean) - with tempfile.TemporaryDirectory() as temp: - for version in (0, 1, 2, 3, 4): - text = clean - if version: - marker = mod.MARK if version == 1 else getattr(mod, f'MARK_V{version}') - helper = ('\ndef _glm53_mixed_prefill_policy(running, current):\n return 0\n\n' if version == 1 else - f'\nclass _Glm53MixedPrefill: {marker}\n pass\n\n') - needle = 'from vllm.compilation.cuda_graph import CUDAGraphStat\n' - text = text.replace(needle, helper + needle, 1) - if version == 4: - for new, old, label in mod.V4_PAIRS: - text = mod.replace_once(text, old, new, label) - else: - names = ['RUNNING', 'WAITING'] if version == 1 else ['BEGIN', 'OBS', 'RUNNING', 'WAITING', 'ALIGN', 'RUNNING_MAMBA', 'WAITING_MAMBA'] - if version == 3: - names.append('FIN') - for name in names: - old = mod.V3_FIN_OLD if name == 'FIN' else getattr(mod, name + '_OLD') - text = mod.replace_once(text, old, getattr(mod, f'V{version}_{name}_NEW'), name) - target = Path(temp) / f'scheduler_v{version}.py' - target.write_text(text) - env = {**os.environ, 'GLM53_SCHEDULER_PY': str(target), 'GLM53_MIXED_PREFILL_CHUNK': 'skip'} - subprocess.run([sys.executable, str(PATCH)], env=env, check=True, capture_output=True) - installed = target.read_text() - compile(installed, str(target), 'exec') - assert mod.MARK_V5 in installed and mod.MARK_V4 not in installed and mod.MARK_V3 not in installed and mod.MARK_V2 not in installed - subprocess.run([sys.executable, str(PATCH)], env=env, check=True, capture_output=True) - assert target.read_text() == installed - # Marker alone must not suppress validation or overwrite source drift. - drifted = installed.replace('_GLM53_MIXED.finish_step(self, scheduler_output)', '_GLM53_MIXED.finish_step_changed(self, scheduler_output)', 1) - target.write_text(drifted) - result = subprocess.run([sys.executable, str(PATCH)], env=env, capture_output=True) - assert result.returncode != 0 and target.read_text() == drifted - # A v5 helper installed before the prefill-priority change must migrate - # in place to the current layout, byte-identical to a fresh install. - legacy = (installed - .replace(mod.PRIORITY_RANK_PREFIX, mod.LEGACY_RANK_PREFIX, 1) - .replace(mod.PRIORITY_CANDIDATES, mod.LEGACY_CANDIDATES, 1)) - assert mod.PRIORITY_MARK not in legacy and mod.MARK_V5 in legacy - target = Path(temp) / 'scheduler_v5_legacy.py' - target.write_text(legacy) + image = src.read_text() + found = mod._identify(image) + clean = image if found is None else found[1] + assert mod.MARK_V7 not in clean and mod.CLASS_HEAD not in clean + + fixtures = ROOT / 'tests' / 'fixtures' + + def load_fixture(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module # inspect.getsource needs the module registered + spec.loader.exec_module(module) + return module + + legacy = load_fixture('legacy_scheduler_helpers', fixtures / 'legacy_scheduler_helpers.py') + + def legacy_producer(version): + text = clean.replace(mod.NEEDLE, legacy.HELPERS[version] + mod.NEEDLE, 1) + for new, old, label in (mod.V1_PAIRS if version == 1 else mod.V2_PAIRS): + text = mod.replace_once(text, old, new, label) + return text + + producers = {'pristine': clean, 'v1': legacy_producer(1), 'v2': legacy_producer(2)} + if found is not None: + producers[found[0]] = image + for ident, fname in (('v5-main', 'v5_main_d5713ec6.py'), ('v5-priority', 'v5_priority_763bd30e.py')): + producers[ident] = load_fixture('producer_' + ident.replace('-', '_'), + fixtures / 'decode_floor_producers' / fname).apply_v5(clean) + + def run(text, temp): + target = Path(temp) / 'scheduler.py' + target.write_text(text) env = {**os.environ, 'GLM53_SCHEDULER_PY': str(target), 'GLM53_MIXED_PREFILL_CHUNK': 'skip'} - subprocess.run([sys.executable, str(PATCH)], env=env, check=True, capture_output=True) - assert target.read_text() == installed - subprocess.run([sys.executable, str(PATCH)], env=env, check=True, capture_output=True) - assert target.read_text() == installed - # A drifted migration anchor fails closed: no partial write even when - # the first of the two anchors still matches. - drifted = legacy.replace('self._candidates = self._rank_prefills(prefills)\n', - 'self._candidates = self._rank_prefills(prefills) # drift\n', 1) - target.write_text(drifted) - result = subprocess.run([sys.executable, str(PATCH)], env=env, capture_output=True) - assert result.returncode != 0 and target.read_text() == drifted - return installed + result = subprocess.run([sys.executable, str(PATCH)], env=env, capture_output=True, text=True) + return result, target.read_text() + + installed = None + with tempfile.TemporaryDirectory() as temp: + for ident, text in producers.items(): + result, after = run(text, temp) + assert result.returncode == 0, (ident, result.stderr) + assert f'from {ident}' in result.stdout, (ident, result.stdout) + compile(after, ident, 'exec') + assert mod.MARK_V7 in after, ident + for older in (mod.MARK_V2, mod.MARK_V3, mod.MARK_V4, mod.MARK_V5, mod.MARK_V6): + assert older not in after, (ident, older) + if installed is None: + installed = after + assert after == installed, f'{ident} did not migrate to the fresh-install bytes' + again, repeat = run(after, temp) + assert again.returncode == 0 and repeat == after and 'already present' in again.stdout, ident + + v5 = producers['v5-priority'] + refused = { + 'v3-marker': v5.replace(mod.MARK_V5, mod.MARK_V3), + 'v4-marker': v5.replace(mod.MARK_V5, mod.MARK_V4), + 'v5-hook-drift': v5.replace('_GLM53_MIXED.finish_step(self, scheduler_output)', + '_GLM53_MIXED.finish_step_changed(self, scheduler_output)', 1), + 'v5-duplicate-helper': v5.replace(mod.CLASS_HEAD, mod.CLASS_HEAD + ' pass\n\n' + mod.CLASS_HEAD, 1), + 'v5-marker-only': clean.replace(mod.NEEDLE, mod.MARK_V5 + '\n' + mod.NEEDLE, 1), + 'v7-hook-drift': installed.replace('_GLM53_MIXED.note_alloc_failed(request) # [glm53-decode-floor:v7]', + '_GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v7]', 1), + 'v7-duplicate-helper': installed.replace(mod.CLASS_HEAD, mod.CLASS_HEAD + ' pass\n\n' + mod.CLASS_HEAD, 1), + 'unmarked-helper': clean.replace(mod.NEEDLE, '\n\nclass _Glm53MixedPrefill:\n pass\n' + mod.NEEDLE, 1), + } + for case, text in refused.items(): + assert text != installed and text != v5, f'{case}: mutation did not apply' + result, after = run(text, temp) + assert result.returncode != 0 and after == text, f'{case} was not refused without writing' + return installed def main(): From 621dd36fa91493ead4aaf20a0332a7d43e75b301 Mon Sep 17 00:00:00 2001 From: Hector Isaza <76170012+hecisaza@users.noreply.github.com> Date: Thu, 1 Oct 2026 00:44:39 +0000 Subject: [PATCH 2/8] docs: decode-floor v7 references; drop stale concurrent-agents page (#283) REFERENCE.md: name decode-floor v7 where it said v5, and describe the v7 installer contract (v3/v4 and drifted installs are refused, not migrated). Drop docs/concurrent-agents.md from this integration: its launcher-default statements predate main's fair default on all launchers. --- CHANGELOG.md | 5 +- docs/REFERENCE.md | 10 ++-- docs/concurrent-agents.md | 104 -------------------------------------- 3 files changed, 7 insertions(+), 112 deletions(-) delete mode 100644 docs/concurrent-agents.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 90943ac8..06d9622f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -262,9 +262,8 @@ There were no git tags for 1.0.0–1.4.0; 1.5.0 is the first cut named as a rele helpers are refused without writing. - `overlay/patch_mamba_align_chunking.py` accepts decode-floor v7 as well as v5. - - #180's concurrency canary (`tests/check_concurrent_agents.py`, - `docs/concurrent-agents.md`) and its CPU progress regressions - (`tests/test_prefill_concurrency.py`) land with it. + - #180's concurrency canary (`tests/check_concurrent_agents.py`) and its CPU + progress regressions (`tests/test_prefill_concurrency.py`) land with it. Qualified live on a TP2 pair with a pre-registered baseline → candidate → baseline run (installed scheduler `97c90a18…`, 2026-09-30): diff --git a/docs/REFERENCE.md b/docs/REFERENCE.md index 64b22a09..07c63df8 100644 --- a/docs/REFERENCE.md +++ b/docs/REFERENCE.md @@ -854,7 +854,7 @@ uses the resolved scheduler LCM for shared checkpoints, not the drafter's smaller block, and preserves smaller Mamba private-state boundaries. EAGLE back-off applies only when a participating non-SWA group needs it. Alignment capacity is bounded by the actual grant so positive sub-block grants progress. -The launchers apply the chunking overlay after decode-floor v5. +The launchers apply the chunking overlay after decode-floor (v5 or v7). The GPU receipts below qualify the retained PR238 implementation. They do not automatically qualify the combined loader and fine-grained APC changes. @@ -868,7 +868,7 @@ one superseded state per concurrent batch. `patch_mamba_align_chunking.py` uses the Mamba group's actual block for checkpoint alignment, not the drafter's smaller block, and applies EAGLE back-off only when the full-attention group needs it. Sub-block token caps still make progress. The launchers apply -the chunking overlay after decode-floor v5. +the chunking overlay after decode-floor (v5 or v7). **Installer compatibility.** The public InstantTensor image carries a legacy `glm53-hybrid-apc` coordinator without the current v3 verification form. @@ -2001,8 +2001,8 @@ After CUDA compile, Python overlay edits (`overlay/exl3.py`, tests) are a cheap | `tests/test_draft_kv_compact.py` | CPU geometry/alignment checks and pinned-source allocator, backend rejection, idempotence, two-file preflight, and prefix-cache lookup (coordinator + managers) tests | | `overlay/patch_glm_video_placeholders.py` | align video timestamp blocks to encoder `grid_t` | | `overlay/patch_suppress_stops_in_reasoning.py` | fail-closed detokenizer guard: client `stop` dormant until `` | -| `overlay/patch_scheduler_decode_floor.py` | skip / cap / off / `fair` mixed-prefill; v5 fixed-cost step fit + largest step-fitting chunk + bounded contention credit, decode first; versioned installer | -| `tests/test_scheduler_decode_floor.py` | v5 migration from v1/v2/v3/v4; cost fit, ladder climb-back, prompt probe, async accounting, bounded credit, alignment and actual scheduler budget regressions | +| `overlay/patch_scheduler_decode_floor.py` | skip / cap / off / `fair` mixed-prefill; v7 (#283): v5 fixed-cost step fit + largest step-fitting chunk + bounded contention credit, decode first, plus KV-blocked admission progress (#246) and priority tiers (#221); authenticated installer migration (#180) | +| `tests/test_scheduler_decode_floor.py` | v7 migration from pristine/v1/v2/v5/v5+priority, refusal of v3/v4 and drifted installs; cost fit, ladder climb-back, prompt probe, async accounting, bounded credit, alignment and actual scheduler budget regressions | | `overlay/patch_xgrammar_termination.py` | source-exact vLLM #52805/#53046 backports; stop at termination and validate post-reasoning speculative drafts before FSM advance | | `tests/test_xgrammar_termination.py` | exact two-file patch, idempotence, cross-file fail-closed drift, termination/rollback/reset and post-reasoning draft behavior, launcher wiring | | `overlay/patch_cache_reset.py` | mount only the upstream cache-reset dev router (`/reset_prefix_cache` et al., #31) when `GLM53_EXPOSE_CACHE_RESET=1`; runtime-mounted by `start.sh` (`CACHE_RESET_PATCH_HOST`) | @@ -2128,7 +2128,7 @@ then-unmerged overlay stack per-group retention + fine-grained hits + gate v2 main carries per-KV-cache-group retention via merged #130, where an empty `GLM53_APC_RETENTION_INTERVAL_SWA` inherits the global retention interval instead of applying #83's automatic rule. #84's overlay patch and the #80 gate-v2 knobs are not -in main. Current TP=2 main defaults to `GLM53_MIXED_PREFILL_CHUNK=fair` (v5); +in main. Current TP=2 main defaults to `GLM53_MIXED_PREFILL_CHUNK=fair` (decode-floor v7); the historical gate-v2 measurements below do not qualify that policy. | ctx 50K per lane (distinct prefixes, verified warm) | ×1 | ×2 | ×4 | ×8 | ×16 | diff --git a/docs/concurrent-agents.md b/docs/concurrent-agents.md deleted file mode 100644 index 7e5b56c1..00000000 --- a/docs/concurrent-agents.md +++ /dev/null @@ -1,104 +0,0 @@ -# Concurrent agents - -Shipped mixed-prefill defaults: `fair` on the two-node `start.sh` (TP=2), -`0` on `start-tp3.sh` (TP=3) and `skip` on `start-tp4.sh` (TP=4). On a launcher -that defaults to `skip`, `GLM53_MIXED_PREFILL_CHUNK=0` is one explicit option -for concurrent interactive agents — not a recommendation. It admits newcomers -immediately at the cost of the incumbent's decode rate, and the TP=2 `CHUNK=0` -default is a gated candidate whose contention qualification is still open. -`fair` is the measured policy for shared TP=2 use. - -Setting `0` restores vLLM's stock chunked-prefill scheduler: a decoding peer no -longer explicitly prevents another prompt from using the remaining token -budget in the same step. `MAX_NUM_SEQS` still limits active sequences, and -`MAX_NUM_BATCHED_TOKENS` limits the whole step. A context still needs its full -prefill before it can emit a first token; this setting does not impose a -wall-clock latency guarantee or bypass a full KV cache. - -The total token budget can still be exhausted by an already-running long -prefill. If a second prompt waits on capacity while cache and sequence slots -are available, inspect the per-request allocation as well as the mixed-prefill -rule. vLLM's existing `--long-prefill-token-threshold` option can cap that -allocation for both running and waiting prompts. For example, with -`MAX_NUM_BATCHED_TOKENS=1024`, a threshold of `512` leaves budget for a second -request (speculative decoding also consumes budget). Add it to `EXTRA_ARGS` -without removing any existing arguments, then recreate and requalify the -service. A smaller cap can reduce single-request prefill throughput; test a -short peer arriving during a cold long prefill and measure both latencies. - -`MAX_MODEL_LEN` is a per-request limit, not a promise that every sequence slot -can hold a context of that size. Compare the engine's startup KV-cache token -capacity and maximum-concurrency estimate with the combined active contexts. -Two million-token windows require capacity for both; enabling interleaving -does not add memory. Agent-side context compression can bound the resident -history while preserving a session across many hours and requests. - -The `skip` policy (still the TP=4 default) gives a new prompt zero prefill tokens -whenever a peer is decoding. A long generation can therefore starve another -agent for minutes, even with spare sequence slots and KV memory. A client idle -timeout then cancels and retries that waiting request. Requests already admitted -to the running list can also have their remaining prefill suppressed. - -`skip` and `-1` remain explicit options for protecting decode throughput at the -expense of new-prompt latency. A positive number caps mixed prefill tokens. -Interleaving can reduce per-agent generation speed, especially for long -contexts on the older sparse-MLA kernels. Choose that tradeoff explicitly. - -Existing `.env` files are not rewritten during updates. Set the policy you have -chosen for that launcher (`fair`, `skip`, a positive cap, or `0`) there, then -recreate the service through your approved deployment procedure. A plain `docker -restart` reuses the old container environment and does not apply this setting. -In-flight requests should be drained first. - -Before measuring, record the deployed source commit, image digest, model and -drafter revisions, launch arguments, context/sequence/token budgets, and -relevant environment. Compare with that exact versioned recipe. Isolate the -mixed-prefill setting; an upgrade of kernels, weights, or memory geometry -requires separate qualification. - -Run the CPU regressions with an installed scheduler file. These operate on -temporary copies and do not import the model or modify the running scheduler: - -```bash -GLM53_SCHEDULER_PY_SRC=/path/to/scheduler.py python3 tests/test_prefill_concurrency.py -GLM53_SCHEDULER_PY_SRC=/path/to/scheduler.py python3 tests/test_scheduler_decode_floor.py -``` - -Reapplication accepts only a complete known patch: the legacy helper site is -validated before anything is removed (v1/v2/v5 against the published helper text -by sha256, the current version against the installer's own text), one exact byte -range is removed, and the whole file is round-trip checked — re-adding that span -and re-applying the version's frozen gate sites must reproduce the input -byte-for-byte. A marker alone is never trusted. A partial patch, changed helper -(including decorators), duplicate helper, or altered gate fails without -modifying the scheduler, and an unversioned variant whose baked default differs -from the published one is refused rather than silently migrated. v3 and v4 are -refused as unsupported without touching the source: no authenticated producer of -their intermediate helper bodies was recovered from public history, so no -canonical v3/v4 image is invented for them. A validated legacy image (v1/v2/v5) -is migrated to the current installer version without changing the operator's -`GLM53_MIXED_PREFILL_CHUNK` value, and -the current version accepts the opt-in gate forms (`GLM53_MIXED_PREFILL_WARM_TOKENS` -/ `_MAX_WAIT_MS`) so a supported gate-v2 scheduler is not rejected at startup. -These CPU checks prove patch integrity and policy behavior, not live request -latency or KV capacity. - -After the deployment and recipe are approved, run one live canary: - -```bash -python3 tests/check_concurrent_agents.py --base-url http://127.0.0.1:8888/v1 --model GLM-5.3-Flash-EXL3 --out logs/concurrent-agents-canary.json -``` - -The canary starts a long, deterministic counting response with thinking -disabled. After A emits content, it submits B with a distinct 8k-word prompt -and asks for `OK`. It requires valid completed streams, the expected short -answer, and B's first content before A's last content. Two HTTP 200s alone do -not pass. The default overall deadline is 240 seconds; on failure it stops -its own requests, returns nonzero, and does not retry or run a benchmark -matrix. If authentication is enabled, supply `OPENAI_API_KEY` through the -environment. Receipts do not include that credential. - -Review the canary's output and latency before wider concurrency testing. -Confirm the model still handles realistic tool calls and long prompts in -your agent client. Increasing a client's idle timeout can accommodate long -prefills but does not repair scheduler starvation. From 300187d71266fa34eb0d9cde33e07a8fea7ade43 Mon Sep 17 00:00:00 2001 From: Hector Isaza <76170012+hecisaza@users.noreply.github.com> Date: Thu, 1 Oct 2026 01:00:13 +0000 Subject: [PATCH 3/8] tests: cover v7 re-verification, native KV refusal recovery, all public producers Address review on #309: - installation_tests() verifies an already-current v7 scheduler exactly (as the patcher's main() does) before unpatching, then checks it equals a fresh install and that re-applying is a verified no-op. - New multi-step regression: a selected waiter natively refused KV lets the runnable peer progress in the same step, is not preselected for one step, stays a candidate and recovers once memory frees, while the peer keeps progressing (#246). Fails if the refusal is ignored or the refused waiter keeps its slot. - Byte-exact producer fixtures for every accepted public identity: v1-image-d9758a6 (helper rebuilds image ef9f5013 byte-for-byte), v5-historical (e329e55), v6-warm-deadline (#180, f89e240) and v6-carry (#246, dc702b4), each migrated and re-applied. --- .../v1_image_d9758a6_helper.py | 6 + .../v5_historical_8e2a32b2.py | 1113 ++++++++++++++++ .../v6_carry_0870fab9.py | 1183 +++++++++++++++++ .../v6_warm_deadline_1f74558d.py | 1143 ++++++++++++++++ tests/test_scheduler_decode_floor.py | 61 +- 5 files changed, 3501 insertions(+), 5 deletions(-) create mode 100644 tests/fixtures/decode_floor_producers/v1_image_d9758a6_helper.py create mode 100644 tests/fixtures/decode_floor_producers/v5_historical_8e2a32b2.py create mode 100644 tests/fixtures/decode_floor_producers/v6_carry_0870fab9.py create mode 100644 tests/fixtures/decode_floor_producers/v6_warm_deadline_1f74558d.py diff --git a/tests/fixtures/decode_floor_producers/v1_image_d9758a6_helper.py b/tests/fixtures/decode_floor_producers/v1_image_d9758a6_helper.py new file mode 100644 index 00000000..60367a08 --- /dev/null +++ b/tests/fixtures/decode_floor_producers/v1_image_d9758a6_helper.py @@ -0,0 +1,6 @@ +"""The decode-floor v1 helper baked into image ef9f5013 (identity v1-image-d9758a6). + +Inserted at the patcher's NEEDLE with the v1 hook pairs, it rebuilds that image's +scheduler.py byte-for-byte from the pristine source. +""" +HELPER = '\ndef _glm53_mixed_prefill_policy(running, current):\n """Mixed-step prefill policy when a peer in `running` is decoding.\n\n None = no extra policy. 0 = skip this prefill this step. N>0 = cap.\n """\n raw = os.environ.get("GLM53_MIXED_PREFILL_CHUNK", "0").strip().lower()\n if raw in ("0", "off", "no"):\n return None\n if raw in ("skip", "-1"):\n cap = 0\n else:\n try:\n cap = int(raw)\n except ValueError:\n cap = 0\n if cap <= 0:\n return None\n cur_id = getattr(current, "request_id", None)\n for r in running:\n if r is current or getattr(r, "request_id", None) == cur_id:\n continue\n if r.num_computed_tokens >= r.num_prompt_tokens:\n return cap\n return None\n\n\n' diff --git a/tests/fixtures/decode_floor_producers/v5_historical_8e2a32b2.py b/tests/fixtures/decode_floor_producers/v5_historical_8e2a32b2.py new file mode 100644 index 00000000..ddfe522a --- /dev/null +++ b/tests/fixtures/decode_floor_producers/v5_historical_8e2a32b2.py @@ -0,0 +1,1113 @@ +#!/usr/bin/env python3 +"""Mixed-prefill policy: skip / cap / off / fair (issue #6). + +A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover +MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill +chunk. Mixed execution leaves the uniform decode FULL-graph path, and a +128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond +budget. + +GLM53_MIXED_PREFILL_CHUNK: + skip / -1 — do not mix prefill with decode. Starves every + waiting/running prefill while any peer is decoding; there is + no age limit. Solo prefill is unchanged. + N>0 — cap mixed prefill chunks to N tokens while a peer decodes. + The cap is fed into hybrid Mamba alignment so N < block_size + still makes sub-block progress. 128 still stalls ~10 tok/s. + 0 / off — disable the extra isolation policy. + fair — service-time mixing (TP=2 default since 2026-09-15, v5; + opt-in on TP=3/TP=4). Decode-only + steps between prefill turns; at most + GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). + Only prefill that contends with a decoder is charged (solo + prefill is cost-sampled, not debt). Credit accrues at SHARE + of accounted engine time. v5 fits a fixed-plus-per-token step + cost from solo and mixed samples, targets the largest ladder + rung (128..2048) whose estimated step fits + GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that + rung instead of spending it on small chunks (v4 priced 1024 + tokens off 128-token samples linearly, ~3x too high, then + stalled at 128-token steps: ~70 tok/s at 20% share). A + never-served newcomer gets one prompt step-bounded probe; + afterwards a 2s age override may borrow one such chunk, only + after all shared debt is repaid. In-flight async prefill + blocks the next mixed turn. Prefills are selected by last + completed positive prefill service plus round-robin. Decode + token/input budget is allocated first by the base scheduler. + Solo prefill retains the base scheduler's limits. Timing is a + host busy-time proxy. + +Fair knobs (read at runtime; identical on every rank): + GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) + GLM53_FAIR_PREFILL_SHARE default 0.20 (credit accrual fraction) + GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 + GLM53_FAIR_PREFILL_MAX_STEP_MS default 1000 (estimated mixed-step limit) + GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 + +Versioned installer: `# [glm53-decode-floor:v5]`. v1 (no version), v2, v3 +and v4 images are unpatched then re-patched. Fail closed if anchors drift. +""" +from __future__ import annotations + +import inspect +import os +import sys +import time +from pathlib import Path + +P = Path( + os.environ.get( + "GLM53_SCHEDULER_PY", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", + ) +) +MARK = "# [glm53-decode-floor]" +MARK_V2 = "# [glm53-decode-floor:v2]" +MARK_V3 = "# [glm53-decode-floor:v3]" +MARK_V4 = "# [glm53-decode-floor:v4]" +MARK_V5 = "# [glm53-decode-floor:v5]" + +IMPORT_OLD = """import itertools +import time +""" +IMPORT_NEW = """import itertools +import os +import time +""" + +# v1 helper + insertions (recipe f906ee9 / this overlay before v2). +V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" +V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. +V2_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] + # NOTE(woosuk) on the scheduling algorithm: +""" +V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] + pooler_outputs = model_runner_output.pooler_output +""" +V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" +V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] +""" +V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] + if num_new_tokens == 0: + break +""" + + + +# Frozen v3 anchors for migration from the reviewed implementation. +V3_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v3] + # NOTE(woosuk) on the scheduling algorithm: +""" +V3_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v3] + pooler_outputs = model_runner_output.pooler_output +""" +V3_FIN_OLD = """ # Check if the scheduling constraints are satisfied. + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) +""" +V3_FIN_NEW = """ # Check if the scheduling constraints are satisfied. + _GLM53_MIXED.note_schedule_output(self, num_scheduled_tokens) # [glm53-decode-floor:v3] + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) +""" +V3_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( + num_new_tokens, + token_budget, + int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), + _GLM53_MIXED.needs_prefill_compute(request), + ) + + # Make sure the input position does not exceed the max model len. +""" +V3_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( + num_new_tokens, + token_budget, + int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), + _GLM53_MIXED.needs_prefill_compute(request), + ) + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if num_new_tokens <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + # chunked prefill has to be enabled explicitly to allow +""" +V3_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v3] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V3_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] +""" +V3_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] + if num_new_tokens == 0: + break +""" + +class _Glm53MixedPrefill: # [glm53-decode-floor:v5] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.20")) + except ValueError: + self.share = 0.20 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "1000")) / 1000.0) + except ValueError: + self.max_step_s = 1.0 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self): + grants = (self._open_rec or {}).get("grants", {}) + for r in self._candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + self.defer_reason = reason + self._promote_next() + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + self._promote_next() + self._maybe_log() + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + return 0 if peer_decode else None + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + + + +def _helper_text() -> str: + body = inspect.getsource(_Glm53MixedPrefill) + return ( + "\n" + + body + + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5]\n\n" + + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5]\n" + + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" + ) + + +HELPER = None # filled at apply time so tests can call _helper_text() + + +BEGIN_OLD = """ self.current_step += 1 + # NOTE(woosuk) on the scheduling algorithm: +""" +BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] + # NOTE(woosuk) on the scheduling algorithm: +""" + +OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output +""" +OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] + pooler_outputs = model_runner_output.pooler_output +""" + +RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + + # Make sure the input position does not exceed the max model len. +""" +RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" + +WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + + # chunked prefill has to be enabled explicitly to allow +""" +WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" + +RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) +""" +RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] +""" + +WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if num_new_tokens == 0: + break +""" +WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] + if num_new_tokens == 0: + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" +FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] + with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" + +RUNNING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled because one of the following +""" +RUNNING_ZERO_NEW = """ if num_new_tokens == 0: + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + # The request cannot be scheduled because one of the following +""" + +WAITING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled. + break +""" +WAITING_ZERO_NEW = """ if num_new_tokens == 0: + # The request cannot be scheduled. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. + # Preempt the lowest-priority request. +""" +PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. + if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] + break + # Preempt the lowest-priority request. +""" + +RUNNING_ALLOC_OLD = """ if new_blocks is None: + # Cannot schedule this request. + break +""" +RUNNING_ALLOC_NEW = """ if new_blocks is None: + # Cannot schedule this request. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + req_index += 1 + continue + break +""" + +WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + break + + # KVTransfer:""" +WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break + + # KVTransfer:""" + +V4_PAIRS = ( + (BEGIN_NEW, BEGIN_OLD, 'begin'), + (OBS_NEW, OBS_OLD, 'obs'), + (RUNNING_NEW, RUNNING_OLD, 'running'), + (WAITING_NEW, WAITING_OLD, 'waiting'), + (ALIGN_NEW, ALIGN_OLD, 'align'), + (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), + (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), + (FIN_NEW, FIN_OLD, 'fin'), + (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), + (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), + (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), + (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), + (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), +) + + +def replace_once(text: str, old: str, new: str, label: str) -> str: + n = text.count(old) + if n != 1: + raise SystemExit(f"{P}: expected one {label} target, found {n}") + return text.replace(old, new, 1) + + +def _strip_helper(text: str, label: str) -> str: + start = text.find("class _Glm53MixedPrefill:") + if start < 0: + raise SystemExit(f"{P}: {label} helper start not found") + if start > 0 and text[start - 1] == "\n": + start -= 1 + end_ak = text.find("class _Glm53AdaptiveK:", start) + end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) + candidates = [i for i in (end_ak, end_cg) if i > start] + if not candidates: + raise SystemExit(f"{P}: {label} helper end not found") + end = min(candidates) + return text[:start] + text[end:] + + +def _strip_v1_helper(text: str) -> str: + start = text.find(V1_HELPER_START) + if start < 0: + raise SystemExit(f"{P}: v1 helper start not found") + if start > 0 and text[start - 1] == "\n": + start -= 1 + end_ak = text.find("class _Glm53AdaptiveK:", start) + end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) + candidates = [i for i in (end_ak, end_cg) if i > start] + if not candidates: + raise SystemExit(f"{P}: v1 helper end not found") + end = min(candidates) + return text[:start] + "\n" + text[end:] + + +def unpatch_v1(text: str) -> str: + if V1_RUNNING_NEW in text: + text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") + elif "_glm53_mixed_prefill_policy(self.running, request)" in text: + raise SystemExit(f"{P}: v1 running insertion drifted") + if V1_WAITING_NEW in text: + text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") + elif "_glm53_mixed_prefill_policy(self.running, request)" in text: + raise SystemExit(f"{P}: v1 waiting insertion drifted") + if V1_HELPER_START in text: + text = _strip_v1_helper(text) + leftover = [ + "_glm53_mixed_prefill_policy(self.running, request)", + V1_HELPER_START, + ] + for s in leftover: + if s in text: + raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") + return text + + +def unpatch_v2(text: str) -> str: + pairs = ( + (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), + (V2_OBS_NEW, OBS_OLD, "v2-obs"), + (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), + (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), + (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), + (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), + (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), + ) + for new, old, label in pairs: + if new in text: + text = replace_once(text, new, old, label) + elif MARK_V2 in new: + # Some insertions may already have been removed; fail if marker remains. + pass + if "class _Glm53MixedPrefill:" in text: + text = _strip_helper(text, "v2") + if MARK_V2 in text: + raise SystemExit(f"{P}: v2 leftover after unpatch") + return text + + +def unpatch_v3(text: str) -> str: + pairs = ( + (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), + (V3_OBS_NEW, OBS_OLD, "v3-obs"), + (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), + (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), + (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), + (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), + (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), + (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), + ) + for new, old, label in pairs: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v3") + if MARK_V3 in text: + raise SystemExit(f"{P}: v3 leftover after unpatch") + return text + + +def unpatch_v4(text: str) -> str: + for new, old, label in V4_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v4") + if MARK_V4 in text: + raise SystemExit(f"{P}: v4 leftover after unpatch") + return text + + +# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 +# insertions above stay frozen so a v4 image can be unpatched exactly. +V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) + + +def unpatch_v5(text: str) -> str: + for new, old, label in V5_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v5") + if MARK_V5 in text: + raise SystemExit(f"{P}: v5 leftover after unpatch") + return text + + +def apply_v5(text: str) -> str: + if "import os\n" not in text.split("import time\n", 1)[0]: + text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") + needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" + text = replace_once(text, needle, _helper_text() + needle, "helper") + for new, old, label in V5_PAIRS: + text = replace_once(text, old, new, label) + compile(text, str(P), "exec") + return text + + +def main() -> int: + if not P.is_file(): + raise SystemExit(f"missing {P}") + text = P.read_text() + original = text + if MARK_V5 in text: + # Validate existing anchors/helper instead of trusting the marker alone. + clean = unpatch_v5(text) + if apply_v5(clean) != text: + raise SystemExit(f"{P}: v5 helper drifted") + print(f"{P.name}: {MARK_V5} already present — verified") + return 0 + if MARK_V4 in text: + text = unpatch_v4(text) + elif MARK_V3 in text: + text = unpatch_v3(text) + elif MARK_V2 in text: + text = unpatch_v2(text) + elif MARK in text or V1_HELPER_START in text: + text = unpatch_v1(text) + text = apply_v5(text) + if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text: + raise SystemExit(f"{P}: older marker left after migration") + if text != original: + P.write_text(text) + print(f"patched {P.name} ({MARK_V5})") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/fixtures/decode_floor_producers/v6_carry_0870fab9.py b/tests/fixtures/decode_floor_producers/v6_carry_0870fab9.py new file mode 100644 index 00000000..0c16d294 --- /dev/null +++ b/tests/fixtures/decode_floor_producers/v6_carry_0870fab9.py @@ -0,0 +1,1183 @@ +#!/usr/bin/env python3 +"""Mixed-prefill policy: skip / cap / off / fair (issue #6). + +A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover +MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill +chunk. Mixed execution leaves the uniform decode FULL-graph path, and a +128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond +budget. + +GLM53_MIXED_PREFILL_CHUNK: + skip / -1 — do not mix prefill with decode. Starves every + waiting/running prefill while any peer is decoding; there is + no age limit. Solo prefill is unchanged. + N>0 — cap mixed prefill chunks to N tokens while a peer decodes. + The cap is fed into hybrid Mamba alignment so N < block_size + still makes sub-block progress. 128 still stalls ~10 tok/s. + 0 / off — disable the extra isolation policy. + fair — service-time mixing (default on TP=2/3/4 since 2026-09-15, + v6). Decode-only + steps between prefill turns; at most + GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). + Only prefill that contends with a decoder is charged (solo + prefill is cost-sampled, not debt). Credit accrues at SHARE + of accounted engine time. v5 fits a fixed-plus-per-token step + cost from solo and mixed samples, targets the largest ladder + rung (128..2048) whose estimated step fits + GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that + rung instead of spending it on small chunks (v4 priced 1024 + tokens off 128-token samples linearly, ~3x too high, then + stalled at 128-token steps: ~70 tok/s at 20% share). A + never-served newcomer gets one prompt step-bounded probe; + afterwards a 2s age override may borrow one such chunk, only + after all shared debt is repaid. In-flight async prefill + blocks the next mixed turn. Prefills are selected by last + completed positive prefill service plus round-robin. Fair + ordering is allocation-aware: a runnable candidate passed by + the scheduler is carried one step when a later candidate + fails KV allocation, preventing blocked-head reset starvation. + Decode token/input budget is allocated first by the base scheduler. + Solo prefill retains the base scheduler's limits. Timing is a + host busy-time proxy. + +Fair knobs (read at runtime; identical on every rank): + GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) + GLM53_FAIR_PREFILL_SHARE default 0.30 (credit accrual fraction) + GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 + GLM53_FAIR_PREFILL_MAX_STEP_MS default 2000 (estimated mixed-step limit) + GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 + +Versioned installer: `# [glm53-decode-floor:v6]`. v1 (no version), v2, v3, +v4 and v5 images are unpatched then re-patched. Fail closed if anchors drift. +""" +from __future__ import annotations + +import inspect +import hashlib +import os +import sys +import time +from pathlib import Path + +P = Path( + os.environ.get( + "GLM53_SCHEDULER_PY", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", + ) +) +MARK = "# [glm53-decode-floor]" +MARK_V2 = "# [glm53-decode-floor:v2]" +MARK_V3 = "# [glm53-decode-floor:v3]" +MARK_V4 = "# [glm53-decode-floor:v4]" +MARK_V5 = "# [glm53-decode-floor:v5]" +MARK_V6 = "# [glm53-decode-floor:v6]" + +IMPORT_OLD = """import itertools +import time +""" +IMPORT_NEW = """import itertools +import os +import time +""" + +# v1 helper + insertions (recipe f906ee9 / this overlay before v2). +V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" +V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. +V2_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] + # NOTE(woosuk) on the scheduling algorithm: +""" +V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] + pooler_outputs = model_runner_output.pooler_output +""" +V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" +V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] +""" +V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] + if num_new_tokens == 0: + break +""" + + + +# Frozen v3 anchors for migration from the reviewed implementation. +V3_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v3] + # NOTE(woosuk) on the scheduling algorithm: +""" +V3_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v3] + pooler_outputs = model_runner_output.pooler_output +""" +V3_FIN_OLD = """ # Check if the scheduling constraints are satisfied. + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) +""" +V3_FIN_NEW = """ # Check if the scheduling constraints are satisfied. + _GLM53_MIXED.note_schedule_output(self, num_scheduled_tokens) # [glm53-decode-floor:v3] + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) +""" +V3_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( + num_new_tokens, + token_budget, + int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), + _GLM53_MIXED.needs_prefill_compute(request), + ) + + # Make sure the input position does not exceed the max model len. +""" +V3_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( + num_new_tokens, + token_budget, + int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), + _GLM53_MIXED.needs_prefill_compute(request), + ) + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if num_new_tokens <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + + # chunked prefill has to be enabled explicitly to allow +""" +V3_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v3] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V3_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] +""" +V3_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] + if num_new_tokens == 0: + break +""" + +class _Glm53MixedPrefill: # [glm53-decode-floor:v6] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self._passed = set() + self._carry = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) + except ValueError: + self.share = 0.30 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) + except ValueError: + self.max_step_s = 2.0 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v6 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, + self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self, prefer_passed=False): + grants = (self._open_rec or {}).get("grants", {}) + candidates = self._candidates + if prefer_passed: + candidates = ([r for r in candidates if r.request_id in self._passed] + + [r for r in candidates if r.request_id not in self._passed]) + for r in candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + if prefer_passed and r.request_id in self._passed: + # Allocation failed after this request's scheduler position + # had already passed. Carry it into the next step instead of + # resetting to the same blocked head of the fair queue. + self._carry.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + # The running loop reports zero even when policy returned cap=0 for an + # unselected request. That request was passed, not attempted; keep it + # eligible for allocation-aware carry if a later selected request fails. + if rid not in self.selected and rid not in rec["grants"]: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + if reason == "zero_progress": + # A selected request that could not make progress rotates behind + # never-served peers. Credit/gap deferrals are not admissions and + # retain their fair rank. + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.defer_reason = reason + self._promote_next(prefer_passed=True) + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self._carry.clear() + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + self._passed = set() + live_prefills = {r.request_id for r in prefills} + self._carry.intersection_update(live_prefills) + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self._carry.clear() + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + for r in self._candidates: + if r.request_id in self._carry: + self.selected.add(r.request_id) + if len(self.selected) >= self.max_chunks: + break + self._promote_next() + self._maybe_log() + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + return 0 if peer_decode else None + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + self._passed.add(rid) + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + self._carry.discard(rid) + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + + + +def _helper_text() -> str: + body = inspect.getsource(_Glm53MixedPrefill) + return ( + "\n" + + body + + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v6]\n\n" + + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v6]\n" + + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" + ) + + +HELPER = None # filled at apply time so tests can call _helper_text() + + +BEGIN_OLD = """ self.current_step += 1 + # NOTE(woosuk) on the scheduling algorithm: +""" +BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] + # NOTE(woosuk) on the scheduling algorithm: +""" + +OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output +""" +OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] + pooler_outputs = model_runner_output.pooler_output +""" + +RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + + # Make sure the input position does not exceed the max model len. +""" +RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" + +WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + + # chunked prefill has to be enabled explicitly to allow +""" +WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" + +RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) +""" +RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] +""" + +WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if num_new_tokens == 0: + break +""" +WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] + if num_new_tokens == 0: + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" +FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] + with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" + +RUNNING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled because one of the following +""" +RUNNING_ZERO_NEW = """ if num_new_tokens == 0: + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + # The request cannot be scheduled because one of the following +""" + +WAITING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled. + break +""" +WAITING_ZERO_NEW = """ if num_new_tokens == 0: + # The request cannot be scheduled. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. + # Preempt the lowest-priority request. +""" +PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. + if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] + break + # Preempt the lowest-priority request. +""" + +RUNNING_ALLOC_OLD = """ if new_blocks is None: + # Cannot schedule this request. + break +""" +RUNNING_ALLOC_NEW = """ if new_blocks is None: + # Cannot schedule this request. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + req_index += 1 + continue + break +""" + +WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + break + + # KVTransfer:""" +WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break + + # KVTransfer:""" + +V4_PAIRS = ( + (BEGIN_NEW, BEGIN_OLD, 'begin'), + (OBS_NEW, OBS_OLD, 'obs'), + (RUNNING_NEW, RUNNING_OLD, 'running'), + (WAITING_NEW, WAITING_OLD, 'waiting'), + (ALIGN_NEW, ALIGN_OLD, 'align'), + (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), + (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), + (FIN_NEW, FIN_OLD, 'fin'), + (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), + (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), + (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), + (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), + (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), +) + + +def replace_once(text: str, old: str, new: str, label: str) -> str: + n = text.count(old) + if n != 1: + raise SystemExit(f"{P}: expected one {label} target, found {n}") + return text.replace(old, new, 1) + + +def _strip_helper(text: str, label: str, *, expected: str | None = None, + expected_sha256: str | None = None) -> str: + start = text.find("class _Glm53MixedPrefill:") + if start < 0: + raise SystemExit(f"{P}: {label} helper start not found") + if start > 0 and text[start - 1] == "\n": + start -= 1 + end_ak = text.find("class _Glm53AdaptiveK:", start) + end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) + candidates = [i for i in (end_ak, end_cg) if i > start] + if not candidates: + raise SystemExit(f"{P}: {label} helper end not found") + end = min(candidates) + actual = text[start:end].strip() + if expected is not None and actual != expected.strip(): + raise SystemExit(f"{P}: {label} helper drifted") + if (expected_sha256 is not None + and hashlib.sha256(actual.encode()).hexdigest() != expected_sha256): + raise SystemExit(f"{P}: {label} helper drifted") + return text[:start] + text[end:] + + +def _strip_v1_helper(text: str) -> str: + start = text.find(V1_HELPER_START) + if start < 0: + raise SystemExit(f"{P}: v1 helper start not found") + if start > 0 and text[start - 1] == "\n": + start -= 1 + end_ak = text.find("class _Glm53AdaptiveK:", start) + end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) + candidates = [i for i in (end_ak, end_cg) if i > start] + if not candidates: + raise SystemExit(f"{P}: v1 helper end not found") + end = min(candidates) + return text[:start] + "\n" + text[end:] + + +def unpatch_v1(text: str) -> str: + if V1_RUNNING_NEW in text: + text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") + elif "_glm53_mixed_prefill_policy(self.running, request)" in text: + raise SystemExit(f"{P}: v1 running insertion drifted") + if V1_WAITING_NEW in text: + text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") + elif "_glm53_mixed_prefill_policy(self.running, request)" in text: + raise SystemExit(f"{P}: v1 waiting insertion drifted") + if V1_HELPER_START in text: + text = _strip_v1_helper(text) + leftover = [ + "_glm53_mixed_prefill_policy(self.running, request)", + V1_HELPER_START, + ] + for s in leftover: + if s in text: + raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") + return text + + +def unpatch_v2(text: str) -> str: + pairs = ( + (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), + (V2_OBS_NEW, OBS_OLD, "v2-obs"), + (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), + (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), + (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), + (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), + (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), + ) + for new, old, label in pairs: + if new in text: + text = replace_once(text, new, old, label) + elif MARK_V2 in new: + # Some insertions may already have been removed; fail if marker remains. + pass + if "class _Glm53MixedPrefill:" in text: + text = _strip_helper(text, "v2") + if MARK_V2 in text: + raise SystemExit(f"{P}: v2 leftover after unpatch") + return text + + +def unpatch_v3(text: str) -> str: + pairs = ( + (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), + (V3_OBS_NEW, OBS_OLD, "v3-obs"), + (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), + (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), + (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), + (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), + (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), + (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), + ) + for new, old, label in pairs: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v3") + if MARK_V3 in text: + raise SystemExit(f"{P}: v3 leftover after unpatch") + return text + + +def unpatch_v4(text: str) -> str: + for new, old, label in V4_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v4") + if MARK_V4 in text: + raise SystemExit(f"{P}: v4 leftover after unpatch") + return text + + +# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 +# insertions above stay frozen so a v4 image can be unpatched exactly. +V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) +V5_HELPER_SHA256 = "d8009ae8ecdc9e731a8f2109f0d89a8f27140c83808674da54c0f6e30788cd5c" + + +def unpatch_v5(text: str) -> str: + for new, old, label in V5_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v5", expected_sha256=V5_HELPER_SHA256) + if MARK_V5 in text: + raise SystemExit(f"{P}: v5 leftover after unpatch") + return text + + +V6_PAIRS = tuple((new.replace(MARK_V4, MARK_V6), old, label) for new, old, label in V4_PAIRS) + + +def unpatch_v6(text: str) -> str: + for new, old, label in V6_PAIRS: + text = replace_once(text, new, old, label) + text = _strip_helper(text, "v6", expected=_helper_text()) + if MARK_V6 in text: + raise SystemExit(f"{P}: v6 leftover after unpatch") + return text + + +def apply_v6(text: str) -> str: + if "import os\n" not in text.split("import time\n", 1)[0]: + text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") + needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" + text = replace_once(text, needle, _helper_text() + needle, "helper") + for new, old, label in V6_PAIRS: + text = replace_once(text, old, new, label) + compile(text, str(P), "exec") + return text + + +def main() -> int: + if not P.is_file(): + raise SystemExit(f"missing {P}") + text = P.read_text() + original = text + if MARK_V6 in text: + # Validate existing anchors/helper instead of trusting the marker alone. + # unpatch_v6 checks every v6 insertion occurs exactly once, strips the + # helper, and rejects leftover markers. The result is discarded: a + # later overlay may legitimately sit between the helper and the + # cuda_graph import anchor, so re-applying at that fixed anchor would + # relocate the helper and fail a byte-compare on a healthy file. + unpatch_v6(text) + # The import edit is part of the applied state; the byte-compare used + # to cover it implicitly. + if "import os\n" not in text.split("import time\n", 1)[0]: + raise SystemExit(f"{P}: v6 import drifted") + compile(text, str(P), "exec") + print(f"{P.name}: {MARK_V6} already present - verified") + return 0 + if MARK_V5 in text: + text = unpatch_v5(text) + elif MARK_V4 in text: + text = unpatch_v4(text) + elif MARK_V3 in text: + text = unpatch_v3(text) + elif MARK_V2 in text: + text = unpatch_v2(text) + elif MARK in text or V1_HELPER_START in text: + text = unpatch_v1(text) + text = apply_v6(text) + if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text or MARK_V5 in text: + raise SystemExit(f"{P}: older marker left after migration") + if text != original: + P.write_text(text) + print(f"patched {P.name} ({MARK_V6})") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/fixtures/decode_floor_producers/v6_warm_deadline_1f74558d.py b/tests/fixtures/decode_floor_producers/v6_warm_deadline_1f74558d.py new file mode 100644 index 00000000..9b915c00 --- /dev/null +++ b/tests/fixtures/decode_floor_producers/v6_warm_deadline_1f74558d.py @@ -0,0 +1,1143 @@ +#!/usr/bin/env python3 +"""Mixed-prefill policy: skip / cap / off / fair (issue #6). + +A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover +MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill +chunk. Mixed execution leaves the uniform decode FULL-graph path, and a +128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond +budget. + +GLM53_MIXED_PREFILL_CHUNK: + skip / -1 — do not mix prefill with decode. Starves every + waiting/running prefill while any peer is decoding; there is + no age limit. Solo prefill is unchanged. + N>0 — cap mixed prefill chunks to N tokens while a peer decodes. + The cap is fed into hybrid Mamba alignment so N < block_size + still makes sub-block progress. 128 still stalls ~10 tok/s. + 0 / off — disable the extra isolation policy. + fair — service-time mixing (TP=2 default since 2026-09-15, v5; + opt-in on TP=3/TP=4). Decode-only + steps between prefill turns; at most + GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). + Only prefill that contends with a decoder is charged (solo + prefill is cost-sampled, not debt). Credit accrues at SHARE + of accounted engine time. v5 fits a fixed-plus-per-token step + cost from solo and mixed samples, targets the largest ladder + rung (128..2048) whose estimated step fits + GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that + rung instead of spending it on small chunks (v4 priced 1024 + tokens off 128-token samples linearly, ~3x too high, then + stalled at 128-token steps: ~70 tok/s at 20% share). A + never-served newcomer gets one prompt step-bounded probe; + afterwards a 2s age override may borrow one such chunk, only + after all shared debt is repaid. In-flight async prefill + blocks the next mixed turn. Prefills are selected by last + completed positive prefill service plus round-robin. Decode + token/input budget is allocated first by the base scheduler. + Solo prefill retains the base scheduler's limits. Timing is a + host busy-time proxy. + +Fair knobs (read at runtime; identical on every rank): + GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) + GLM53_FAIR_PREFILL_SHARE default 0.20 (credit accrual fraction) + GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 + GLM53_FAIR_PREFILL_MAX_STEP_MS default 1000 (estimated mixed-step limit) + GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 + +Versioned installer: `# [glm53-decode-floor:v6]`. v1 (no version), v2 and v5 +images are unpatched then re-patched. A v3 or v4 marker is refused without +touching the source: no authenticated producer of those intermediate helper +bodies was recovered from public history, so a canonical v3/v4 image cannot be +established, and no body is invented to fill the gap. Fail closed if anchors +drift. + +Fail-closed migration: the legacy helper site is validated *before* anything is +removed -- v1/v2/v5 against the published helper text (sha256), v6 against this +installer's own text. The frozen gate sites are inverted, one exact byte range +is removed, and the whole file is round-trip checked: re-adding the same span +and re-applying that version's frozen sites must reproduce the input +byte-for-byte. A drifted, duplicated, decorated, marker-only, v3/v4-marked or +otherwise unattested site is refused with no write. + +v6 (opt-in gate, both features OFF by default -- v5 behaviour is preserved): + GLM53_MIXED_PREFILL_WARM_TOKENS >0 admits a request whose uncached remainder + is <= this many tokens (typically one hybrid block, 3584) without waiting + for the decoders. 0 (default) disables the bypass. + GLM53_MIXED_PREFILL_MAX_WAIT_MS >0 releases a request the `skip` hold has + starved for this long under GLM53_MIXED_PREFILL_LATE_CAP (default 512) + tokens per step. 0 (default) waits forever, i.e. v1..v5 behaviour. The + release only makes the request eligible for a LATE_CAP step; it bounds + neither allocation nor compute service (capacity and the base scheduler's + own limits still decide), and the request then crawls like cap:N. + An explicit `cap` keeps its cap (only the warm bypass applies); `fair` is a + different mechanism and is not touched by either knob. +""" +from __future__ import annotations + +import hashlib +import inspect +import os +import sys +import time +from pathlib import Path + +P = Path( + os.environ.get( + "GLM53_SCHEDULER_PY", + "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", + ) +) +MARK = "# [glm53-decode-floor]" +MARK_V2 = "# [glm53-decode-floor:v2]" +MARK_V3 = "# [glm53-decode-floor:v3]" +MARK_V4 = "# [glm53-decode-floor:v4]" +MARK_V5 = "# [glm53-decode-floor:v5]" +MARK_V6 = "# [glm53-decode-floor:v6]" + +IMPORT_OLD = """import itertools +import time +""" +IMPORT_NEW = """import itertools +import os +import time +""" + +# v1 helper + insertions (recipe f906ee9 / this overlay before v2). +V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" +V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] + if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. +V2_BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] + # NOTE(woosuk) on the scheduling algorithm: +""" +V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] + pooler_outputs = model_runner_output.pooler_output +""" +V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" +V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" +V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] +""" +V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] + if num_new_tokens == 0: + break +""" + + +class _Glm53MixedPrefill: # [glm53-decode-floor:v5] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.logged_gate = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.20")) + except ValueError: + self.share = 0.20 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "1000")) / 1000.0) + except ValueError: + self.max_step_s = 1.0 + # v6 gate (warm bypass + deadline). Both features are OFF unless an + # operator sets the knobs: 0 disables each one and is the default, so a + # skip/cap deployment keeps exactly the v5 hold behaviour. + try: + self.warm_tokens = int(self._e("GLM53_MIXED_PREFILL_WARM_TOKENS", "0")) + except ValueError: + self.warm_tokens = 0 + if not 0 <= self.warm_tokens <= 1_000_000: + self.warm_tokens = 0 + try: + self.max_wait_ms = int(self._e("GLM53_MIXED_PREFILL_MAX_WAIT_MS", "0")) + except ValueError: + self.max_wait_ms = 0 + if not 0 <= self.max_wait_ms <= 600_000: + self.max_wait_ms = 0 + try: + self.late_cap = int(self._e("GLM53_MIXED_PREFILL_LATE_CAP", "512")) + except ValueError: + self.late_cap = 512 + if not 64 <= self.late_cap <= 8192: + self.late_cap = 512 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + if (self.warm_tokens or self.max_wait_ms) and not self.logged_gate: + print( + f"[glm53-decode-floor] gate v6 warm_tokens={self.warm_tokens} " + f"max_wait_ms={self.max_wait_ms} late_cap={self.late_cap}", + flush=True, + ) + self.logged_gate = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self): + grants = (self._open_rec or {}).get("grants", {}) + for r in self._candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + self.defer_reason = reason + self._promote_next() + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + self._promote_next() + self._maybe_log() + + def _warm_bypass(self, remaining): + """Admit a small uncached remainder at once (v6). Off while the knob is 0.""" + return self.warm_tokens > 0 and remaining <= self.warm_tokens + + def _deadline_cap(self, request, remaining): + """Late admission for a request the skip hold would starve (v6). + + Returns 0 while the deadline is disabled (0 ms = wait forever, v5 + behaviour) or not reached, else LATE_CAP tokens for this step. The + first-seen stamp is the per-request arrival the fair policy already + keeps, so chunking, preemption and requeue cannot reset the wait. + """ + if self.max_wait_ms <= 0: + return 0 + rid = request.request_id + first_seen = self.arrival.get(rid) + if first_seen is None: + first_seen = self._now() + self.arrival[rid] = first_seen + if (self._now() - first_seen) * 1000.0 >= self.max_wait_ms: + return max(1, min(self.late_cap, remaining)) + return 0 + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + if not peer_decode: + return None + if self._warm_bypass(remaining): + return None + return self._deadline_cap(request, remaining) + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + if cap is not None and self._warm_bypass(remaining): + cap = None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + + + +BEGIN_OLD = """ self.current_step += 1 + # NOTE(woosuk) on the scheduling algorithm: +""" +BEGIN_NEW = """ self.current_step += 1 + _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] + # NOTE(woosuk) on the scheduling algorithm: +""" + +OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output +""" +OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens + _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] + pooler_outputs = model_runner_output.pooler_output +""" + +RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + + # Make sure the input position does not exceed the max model len. +""" +RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: + num_new_tokens = self.scheduler_config.long_prefill_token_threshold + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) + mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): + num_new_tokens = min(num_new_tokens, mixed_cap) + + # Make sure the input position does not exceed the max model len. +""" + +WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + + # chunked prefill has to be enabled explicitly to allow +""" +WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold + if 0 < threshold < num_new_tokens: + num_new_tokens = threshold + mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] + if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): + if mixed_cap <= 0: + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + num_new_tokens = min(num_new_tokens, mixed_cap) + + # chunked prefill has to be enabled explicitly to allow +""" + +ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" +ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens + long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold + if long_prefill_threshold > 0: + max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) + _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] + if _align_cap is not None and _align_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, _align_cap) + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end +""" + +RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) +""" +RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. + if self.need_mamba_block_aligned_split: + num_new_tokens = self._mamba_block_aligned_split( + request, num_new_tokens + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] +""" + +WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + if num_new_tokens == 0: + break +""" +WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( + request, + num_new_tokens, + num_new_local_computed_tokens, + num_external_computed_tokens, + ) + _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] + if num_new_tokens == 0: + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" +FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] + with record_function_or_nullcontext("schedule: update_after_schedule"): + self._update_after_schedule(scheduler_output) +""" + +RUNNING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled because one of the following +""" +RUNNING_ZERO_NEW = """ if num_new_tokens == 0: + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + # The request cannot be scheduled because one of the following +""" + +WAITING_ZERO_OLD = """ if num_new_tokens == 0: + # The request cannot be scheduled. + break +""" +WAITING_ZERO_NEW = """ if num_new_tokens == 0: + # The request cannot be scheduled. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.mode == "fair": + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break +""" + +PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. + # Preempt the lowest-priority request. +""" +PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. + if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] + break + # Preempt the lowest-priority request. +""" + +RUNNING_ALLOC_OLD = """ if new_blocks is None: + # Cannot schedule this request. + break +""" +RUNNING_ALLOC_NEW = """ if new_blocks is None: + # Cannot schedule this request. + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + req_index += 1 + continue + break +""" + +WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + break + + # KVTransfer:""" +WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: + self.encoder_cache_manager.free(request) + _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] + if _GLM53_MIXED.protect_decode(request): + request_queue.pop_request() + step_skipped_waiting.prepend_request(request) + continue + break + + # KVTransfer:""" + + +def replace_once(text: str, old: str, new: str, label: str) -> str: + n = text.count(old) + if n != 1: + raise SystemExit(f"{P}: expected one {label} target, found {n}") + return text.replace(old, new, 1) + + +V2_PAIRS = ( + (V2_BEGIN_NEW, BEGIN_OLD, 'v2-begin'), + (V2_OBS_NEW, OBS_OLD, 'v2-obs'), + (V2_RUNNING_NEW, RUNNING_OLD, 'v2-running'), + (V2_WAITING_NEW, WAITING_OLD, 'v2-waiting'), + (V2_ALIGN_NEW, ALIGN_OLD, 'v2-align'), + (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'v2-running-mamba'), + (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'v2-waiting-mamba'), +) + +V4_PAIRS = ( + (BEGIN_NEW, BEGIN_OLD, 'begin'), + (OBS_NEW, OBS_OLD, 'obs'), + (RUNNING_NEW, RUNNING_OLD, 'running'), + (WAITING_NEW, WAITING_OLD, 'waiting'), + (ALIGN_NEW, ALIGN_OLD, 'align'), + (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), + (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), + (FIN_NEW, FIN_OLD, 'fin'), + (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), + (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), + (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), + (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), + (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), +) + +V1_PAIRS = ( + (V1_RUNNING_NEW, RUNNING_OLD, 'v1-running'), + (V1_WAITING_NEW, WAITING_OLD, 'v1-waiting'), +) + +# v5 and v6 use the same scheduler anchors as v4 with the marker advanced; the +# v4 insertions above stay frozen so a v4 image can be unpatched exactly. v6 +# only changes the helper (opt-in warm bypass + deadline, both off by default), +# so an unpatch/re-patch of a v6 image is byte-identical. +V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) +V6_PAIRS = tuple((new.replace(MARK_V5, MARK_V6), old, label) for new, old, label in V5_PAIRS) + +LEGACY_MARK = {1: MARK, 2: MARK_V2, 3: MARK_V3, 4: MARK_V4, 5: MARK_V5, 6: MARK_V6} +LEGACY_PAIRS = {1: V1_PAIRS, 2: V2_PAIRS, 5: V5_PAIRS, 6: V6_PAIRS} # v3/v4 are refused, not migrated + +NEEDLE = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" +ADAPTIVE_K_HEAD = "class _Glm53AdaptiveK:" +CLASS_HEAD = "class _Glm53MixedPrefill:" + +# --------------------------------------------------------------------------- +# Canonical legacy helper registry. +# +# Every migratable version has exactly one accepted helper site, and the site is +# validated *before* anything is removed. The versions whose helper text is +# published are compared byte-for-byte (sha256 and length). v3 and v4 were +# introduced by 180725a5ce33 as markers for intermediate builds whose helper +# bodies were never recovered from public history: they are refused as +# unsupported rather than migrated on a structural guess, because an invented +# body would not be a canonical image. A site that validates is removed as one +# exact byte range, so text between the helper and its anchors can never be +# dropped by accident. +# +# provenance, public history of MiaAI-Lab/GLM-5.3-Flash-EXL3-2x-DGX-Sparks: +# v1 f3043c95bbf9 overlay/patch_scheduler_decode_floor.py:HELPER +# cross-attested byte-identical by cad980f7263:HELPER_V1, which documents +# itself as "byte-identical to the shipped v1 overlay" +# v2 9a23b2d8f061 overlay/patch_scheduler_decode_floor.py:HELPER +# v5 a9bbbadc4c72 overlay/patch_scheduler_decode_floor.py:_helper_text() +# v3/v4 no public revision carries a body (180725a5ce33 adds only the +# markers); refused, not migrated (see above) +# +# The published variants that share the unversioned marker but bake a different +# policy default (the "0" default of 83f0067/d9758a6/9a557cf/14b6a9f) are NOT +# accepted: migrating them would silently land a held default change, so they +# are refused like any other drifted site. +# --------------------------------------------------------------------------- +LEGACY_HELPER_SHA256 = { + 1: "c0c10c6385bd7e75bfc0f724378960b9d9cb943cf18c386a5a780cf03cb4c48d", + 2: "e170f2bf6d0c0fe0493a91ff54bd719036f6c4b5ea9fcf2ee1346968a4e3fbed", + 5: "c05769276df6da0a24b3b8c21e251f39ee7aa2a74ee92c6d6ce3d85c38151a53", +} +LEGACY_HELPER_LEN = {1: 784, 2: 13378, 5: 20866} + + +def _refuse(reason: str): + raise SystemExit(f"{P}: refusing to rewrite the scheduler: {reason}") + + +def _site_reasons(span: str, version: int, exact=None) -> str: + """Return '' when `span` is exactly this installer's canonical v{version} helper site. + + v1/v2/v5 are pinned to their published helper text (sha256 and length); v6 is + compared against this installer's own `_helper_text()`. A version with no + published text has no accepted site at all. + """ + if exact is not None: + return "" if span == exact else f"helper site is not this installer's v{version} helper text" + expected = LEGACY_HELPER_SHA256.get(version) + if expected is None: + return f"v{version} helper text was never published; there is no accepted site to compare" + got = hashlib.sha256(span.encode("utf-8")).hexdigest() + if got != expected or len(span) != LEGACY_HELPER_LEN[version]: + return (f"helper site is not the published v{version} helper (sha256 {got[:16]}, " + f"{len(span)} bytes; expected {expected[:16]}, {LEGACY_HELPER_LEN[version]} bytes)") + return "" + + +def _legacy_span(text: str, version: int, exact=None): + """Validate and return (start, end, span) for the unique v{version} helper site. + + Returns None when no helper definition is present at all. Refuses (no write) + on a missing, duplicated, decorated, drifted or otherwise unattested site. + """ + head_token = V1_HELPER_START if version == 1 else CLASS_HEAD + found = text.count(head_token) + if found == 0: + return None + if found != 1: + _refuse(f"{found} {head_token!r} definitions found; v{version} installs exactly one") + head = text.find(head_token) + if head < 2 or text[head - 1] != "\n" or text[head - 2] != "\n": + _refuse(f"v{version} helper is not preceded by the installer's blank-line boundary " + f"(indented or decorated site)") + start = head - 1 + reasons = [] + for anchor in (NEEDLE, ADAPTIVE_K_HEAD): + end = text.find(anchor, head) + if end < 0: + continue + span = text[start:end] + why = _site_reasons(span, version, exact) + if not why: + return start, end, span + reasons.append(why) + _refuse(f"v{version} helper site rejected: " + "; ".join(reasons)) + + +def _unpatch(text: str, version: int, exact=None): + """Invert version `version`: its frozen gate sites, then its validated helper site.""" + if version not in LEGACY_PAIRS: + _refuse(f"v{version} is not a migratable version: its helper text was never published") + for new, old, label in LEGACY_PAIRS[version]: + if new not in text: + _refuse(f"{label} insertion is missing or drifted") + text = replace_once(text, new, old, label) + site = _legacy_span(text, version, exact) + if site is None: + _refuse(f"v{version} helper site is missing") + start, end, span = site + clean = text[:start] + text[end:] + mark = LEGACY_MARK[version] + leftover = [s for s in (mark, V1_HELPER_START if version == 1 else CLASS_HEAD) if s in clean] + if leftover: + _refuse(f"v{version} marker or helper definition still present after unpatch") + return clean, start, span + + +def _round_trip(original: str, clean: str, version: int, start: int, span: str) -> None: + """Refuse unless re-applying version `version` to `clean` reproduces `original`.""" + again = clean[:start] + span + clean[start:] + for new, old, label in LEGACY_PAIRS[version]: + again = replace_once(again, old, new, label) + if again != original: + _refuse(f"v{version} round trip is not byte-identical; the image is not a canonical " + f"v{version} install") + + +def _helper_text() -> str: + body = inspect.getsource(_Glm53MixedPrefill) + return ( + "\n" + + body.replace(MARK_V5, MARK_V6) + + f"\n_GLM53_MIXED = _Glm53MixedPrefill() # {MARK_V6}\n\n" + + f"def _glm53_mixed_prefill_policy(sched, request, computed=None): # {MARK_V6}\n" + + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" + ) + + +def apply_v6(text: str) -> str: + if "import os\n" not in text.split("import time\n", 1)[0]: + text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") + text = replace_once(text, NEEDLE, _helper_text() + NEEDLE, "helper") + for new, old, label in V6_PAIRS: + text = replace_once(text, old, new, label) + compile(text, str(P), "exec") + return text + + +def main() -> int: + if not P.is_file(): + raise SystemExit(f"missing {P}") + text = P.read_text() + original = text + if MARK_V6 in text: + # Validate existing anchors and helper instead of trusting the marker: a + # marker-only or hand-edited patch cannot round-trip and is refused + # without a write. + clean, start, span = _unpatch(text, 6, exact=_helper_text()) + if apply_v6(clean) != text: + _refuse("v6 helper drifted; unpatch/re-patch is not byte-identical") + print(f"{P.name}: {MARK_V6} already present — verified") + return 0 + for unsupported in (3, 4): + if LEGACY_MARK[unsupported] in text: + _refuse(f"v{unsupported} is refused as unsupported: no authenticated producer of that " + f"helper body was recovered from public history, so a canonical v{unsupported} " + f"image cannot be established and the source is left unchanged") + version = next((v for v in (5, 2, 1) + if LEGACY_MARK[v] in text or (v == 1 and V1_HELPER_START in text)), None) + if version is not None: + clean, start, span = _unpatch(text, version) + _round_trip(text, clean, version, start, span) + text = clean + text = apply_v6(text) + if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text or MARK_V5 in text: + raise SystemExit(f"{P}: older marker left after migration") + if text != original: + P.write_text(text) + print(f"patched {P.name} ({MARK_V6})") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_scheduler_decode_floor.py b/tests/test_scheduler_decode_floor.py index 401251d2..da2f737e 100644 --- a/tests/test_scheduler_decode_floor.py +++ b/tests/test_scheduler_decode_floor.py @@ -385,6 +385,37 @@ def test_zero_progress_promotes_next_candidate(self): self.p.note_scheduled(self.b, 0) self.assertEqual(self.p.cap_for(self.s, c), 256) + def test_native_kv_refusal_keeps_waiter_retryable_and_peer_progressing(self): + # #246: a selected waiter that vLLM refuses KV must not stall a runnable + # prefill, must stay a candidate, and must recover once memory frees. + c = Req('C') + self.s.waiting.append(c) + self.s.refresh() + self.assertEqual(self.p.cap_for(self.s, self.b), 256) + self.assertEqual(self.p.cap_for(self.s, c), 0) + self.p.note_alloc_failed(self.b) # native refusal (new_blocks is None) + cap_c = self.p.cap_for(self.s, c) + self.assertGreater(cap_c, 0) # the runnable peer progresses in the same step + self.complete(self.submit({'A': 8, 'C': cap_c}), 0.2) + self.p.begin_step(self.s) # next step: refusal evidence lasts one step + self.assertEqual(self.p._refused_prev, {'B'}) + self.assertNotIn('B', self.p.selected) # not preselected right after its refusal + self.assertIn('B', [r.request_id for r in self.p._candidates]) # but still queued for retry + served = {'B': 0, 'C': cap_c} + for _ in range(20): # memory has freed: no further refusals + counts = {'A': 8} + for r in (self.b, c): + cap = self.p.cap_for(self.s, r) + if cap: + counts[r.request_id] = cap + served[r.request_id] += cap + self.complete(self.submit(counts), 0.2) + self.p.begin_step(self.s) + if served['B'] and served['C'] > cap_c: + break + self.assertGreater(served['B'], 0) # the refused waiter recovers + self.assertGreater(served['C'], cap_c) # and the peer keeps progressing + def test_full_prefix_hit_is_not_blocked_as_cold_prefill(self): self.p.begin_step(self.s) self.p.credit = -10 @@ -461,8 +492,15 @@ def installation_tests(): if src is None: raise SystemExit('Set GLM53_SCHEDULER_PY_SRC to the pinned scheduler source') image = src.read_text() - found = mod._identify(image) - clean = image if found is None else found[1] + current = None + if mod.MARK_V7 in image: + # An already-current v7 install: verify it exactly (as main() does), then unpatch. + clean, _, _ = mod._unpatch(image, 'v7', mod.MARK_V7, mod.CLASS_HEAD, mod.V7_PAIRS, + exact=mod._helper_text()) + found, current = None, image + else: + found = mod._identify(image) + clean = image if found is None else found[1] assert mod.MARK_V7 not in clean and mod.CLASS_HEAD not in clean fixtures = ROOT / 'tests' / 'fixtures' @@ -483,11 +521,19 @@ def legacy_producer(version): return text producers = {'pristine': clean, 'v1': legacy_producer(1), 'v2': legacy_producer(2)} + image_helper = load_fixture('v1_image_helper', fixtures / 'decode_floor_producers' / 'v1_image_d9758a6_helper.py') + v1_image = clean.replace(mod.NEEDLE, image_helper.HELPER + mod.NEEDLE, 1) + for new, old, label in mod.V1_PAIRS: + v1_image = mod.replace_once(v1_image, old, new, label) + producers['v1-image-d9758a6'] = v1_image if found is not None: producers[found[0]] = image - for ident, fname in (('v5-main', 'v5_main_d5713ec6.py'), ('v5-priority', 'v5_priority_763bd30e.py')): - producers[ident] = load_fixture('producer_' + ident.replace('-', '_'), - fixtures / 'decode_floor_producers' / fname).apply_v5(clean) + # Byte-exact public producers of every other non-v7 identity the installer accepts. + for ident, fname in (('v5-historical', 'v5_historical_8e2a32b2.py'), ('v5-main', 'v5_main_d5713ec6.py'), + ('v5-priority', 'v5_priority_763bd30e.py'), ('v6-warm-deadline', 'v6_warm_deadline_1f74558d.py'), + ('v6-carry', 'v6_carry_0870fab9.py')): + producer = load_fixture('producer_' + ident.replace('-', '_'), fixtures / 'decode_floor_producers' / fname) + producers[ident] = (getattr(producer, 'apply_v6', None) or producer.apply_v5)(clean) def run(text, temp): target = Path(temp) / 'scheduler.py' @@ -512,6 +558,11 @@ def run(text, temp): again, repeat = run(after, temp) assert again.returncode == 0 and repeat == after and 'already present' in again.stdout, ident + if current is not None: + assert current == installed, 'the installed v7 differs from a fresh install' + again, repeat = run(current, temp) + assert again.returncode == 0 and repeat == current and 'already present' in again.stdout + v5 = producers['v5-priority'] refused = { 'v3-marker': v5.replace(mod.MARK_V5, mod.MARK_V3), From 73406442937f171ffe690e2a9206cdfd43fea7bf Mon Sep 17 00:00:00 2001 From: Hector Isaza <76170012+hecisaza@users.noreply.github.com> Date: Thu, 1 Oct 2026 01:23:27 +0000 Subject: [PATCH 4/8] scheduler: refuse leftover decode-floor markers after unpatch; canary rejects any URL userinfo Address review on #309: - _unpatch() now refuses when any decode-floor marker (any version) or either helper definition remains after removing the identified install. A canonical v7 with a stray v5 or v1 marker no longer verifies as "already present". Installer-only change: on image ef9f5013 the installed scheduler is still sha256 97c90a18fd0df6f9aeb1452c39801455d90756d315c6846ca0642e934fe4e269. - check_concurrent_agents.py rejects password-only and empty-user userinfo (e.g. http://:secret@host), not only a non-empty username. - Regressions: v7-plus-v5/v1-marker refusals in installation_tests() and tests/test_check_concurrent_agents.py; both fail on the previous code. --- overlay/patch_scheduler_decode_floor.py | 5 +++- tests/check_concurrent_agents.py | 2 +- tests/test_check_concurrent_agents.py | 32 +++++++++++++++++++++++++ tests/test_scheduler_decode_floor.py | 3 +++ 4 files changed, 40 insertions(+), 2 deletions(-) create mode 100644 tests/test_check_concurrent_agents.py diff --git a/overlay/patch_scheduler_decode_floor.py b/overlay/patch_scheduler_decode_floor.py index 4ae9e8e9..3d6ba37a 100644 --- a/overlay/patch_scheduler_decode_floor.py +++ b/overlay/patch_scheduler_decode_floor.py @@ -1312,7 +1312,10 @@ def _unpatch(text: str, name, marker, head_token, pairs, sha=None, length=None, if not _followed_by_anchor(text, start + n): _refuse(f"{name}: unexpected text between the helper and its anchor") clean = text[:start] + text[start + n:] - if marker in clean or head_token in clean: + # Fail closed on mixed state: after removing this identity, no decode-floor + # marker of any version and no helper definition of either form may remain. + if (marker in clean or head_token in clean or MARK.rstrip("]") in clean + or CLASS_HEAD in clean or V1_HELPER_START in clean): _refuse(f"{name}: marker or helper definition left after unpatch") return clean, start, span diff --git a/tests/check_concurrent_agents.py b/tests/check_concurrent_agents.py index 4beee3df..92bcd52b 100644 --- a/tests/check_concurrent_agents.py +++ b/tests/check_concurrent_agents.py @@ -135,7 +135,7 @@ def main(): parser.add_argument("--out", type=Path) args = parser.parse_args() url = urlsplit(args.base_url) - if url.scheme not in ("http", "https") or not url.hostname or url.username or url.query or url.fragment: + if url.scheme not in ("http", "https") or not url.hostname or url.username is not None or url.password is not None or url.query or url.fragment: parser.error("base URL must be HTTP(S) without credentials, query, or fragment") if args.filler_words < 1 or args.long_max_tokens < 128 or args.timeout <= 0: parser.error("positive filler/timeout and at least 128 long-request tokens required") diff --git a/tests/test_check_concurrent_agents.py b/tests/test_check_concurrent_agents.py new file mode 100644 index 00000000..a9d2c00b --- /dev/null +++ b/tests/test_check_concurrent_agents.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +"""Argument validation of tests/check_concurrent_agents.py (no network: rejected before any request).""" +from __future__ import annotations + +from pathlib import Path +import subprocess +import sys +import unittest + +CANARY = Path(__file__).resolve().parent / "check_concurrent_agents.py" + + +class BaseUrlValidation(unittest.TestCase): + def rejected(self, url): + result = subprocess.run([sys.executable, str(CANARY), "--base-url", url, "--timeout", "1"], + capture_output=True, text=True, timeout=60) + return result.returncode == 2 and "without credentials" in result.stderr + + def test_userinfo_in_any_form_is_rejected(self): + for url in ("http://user@127.0.0.1:8888/v1", "http://user:pw@127.0.0.1:8888/v1", + "http://:secret@127.0.0.1:8888/v1", "http://@127.0.0.1:8888/v1"): + with self.subTest(url=url): + self.assertTrue(self.rejected(url)) + + def test_query_fragment_and_scheme_are_rejected(self): + for url in ("http://127.0.0.1:8888/v1?x=1", "http://127.0.0.1:8888/v1#f", "ftp://127.0.0.1/v1"): + with self.subTest(url=url): + self.assertTrue(self.rejected(url)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_scheduler_decode_floor.py b/tests/test_scheduler_decode_floor.py index da2f737e..e2a36a38 100644 --- a/tests/test_scheduler_decode_floor.py +++ b/tests/test_scheduler_decode_floor.py @@ -575,6 +575,9 @@ def run(text, temp): '_GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v7]', 1), 'v7-duplicate-helper': installed.replace(mod.CLASS_HEAD, mod.CLASS_HEAD + ' pass\n\n' + mod.CLASS_HEAD, 1), 'unmarked-helper': clean.replace(mod.NEEDLE, '\n\nclass _Glm53MixedPrefill:\n pass\n' + mod.NEEDLE, 1), + # Mixed state: a canonical v7 plus a leftover older marker must not verify. + 'v7-plus-v5-marker': mod.MARK_V5 + '\n' + installed, + 'v7-plus-v1-marker': installed + '\n' + mod.MARK + '\n', } for case, text in refused.items(): assert text != installed and text != v5, f'{case}: mutation did not apply' From dcd57d9e83379d49941b864dd92e8ae7202f96bb Mon Sep 17 00:00:00 2001 From: Hector Isaza <76170012+hecisaza@users.noreply.github.com> Date: Thu, 1 Oct 2026 02:08:45 +0000 Subject: [PATCH 5/8] scheduler: refuse unknown markers and duplicate wrappers; test deployed v7 at its own position Address independent verification of #309 at 7340644: - _identify() treats a source as pristine only when no decode-floor marker of any version and no helper name (_Glm53MixedPrefill, _GLM53_MIXED, _glm53_mixed_prefill_policy) is present; an unknown marker such as "# [glm53-decode-floor:v99]" is refused without writing instead of being installed beside. - _unpatch() applies the same rule to whatever remains after removing the identified install, so a later duplicate policy wrapper (which would rebind the global the hooks call) no longer verifies as current. - installation_tests() compares a deployed v7 against apply_v7() at the helper's own position, so a scheduler that also carries adaptive-K (the deployed layout, helper before _Glm53AdaptiveK) passes; both compositions are checked explicitly. New refusal cases: duplicate wrapper on v7 and on pristine, unknown marker, unknown text after the helper. Installer-only change: on image ef9f5013 the installed scheduler is still sha256 97c90a18fd0df6f9aeb1452c39801455d90756d315c6846ca0642e934fe4e269 and re-applying on the deployed layout verifies without writing. The two unpublished v7 pre-release identities still migrate to the fresh-install bytes. --- CHANGELOG.md | 5 ++-- overlay/patch_scheduler_decode_floor.py | 19 +++++++++++--- tests/test_scheduler_decode_floor.py | 34 ++++++++++++++++++++++--- 3 files changed, 50 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 06d9622f..08628ff4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -258,8 +258,9 @@ There were no git tags for 1.0.0–1.4.0; 1.5.0 is the first cut named as a rele - Installer migration follows #180's fail-closed contract. Pristine, v1, v2, v5 and v5+priority (the #221 layout on `main` until now) installations migrate to the same bytes as a fresh v7 install, and a re-apply is a - verified no-op. v3/v4, drifted, duplicated, marker-only and unmarked - helpers are refused without writing. + verified no-op. v3/v4 and unknown version markers, drifted, duplicated, + marker-only and unmarked helpers, and a duplicate policy wrapper are + refused without writing. - `overlay/patch_mamba_align_chunking.py` accepts decode-floor v7 as well as v5. - #180's concurrency canary (`tests/check_concurrent_agents.py`) and its CPU diff --git a/overlay/patch_scheduler_decode_floor.py b/overlay/patch_scheduler_decode_floor.py index 3d6ba37a..1664a00a 100644 --- a/overlay/patch_scheduler_decode_floor.py +++ b/overlay/patch_scheduler_decode_floor.py @@ -1290,6 +1290,16 @@ def _followed_by_anchor(text: str, end: int) -> bool: or rest.startswith("\n" + ADAPTIVE_K_HEAD)) +# Any decode-floor marker (every version, including unknown ones) or any helper +# name. A pristine scheduler contains none; neither do the sibling overlays. +_LEFTOVER_TOKENS = (MARK.rstrip("]"), CLASS_HEAD, V1_HELPER_START, + "_Glm53MixedPrefill", "_GLM53_MIXED", "_glm53_mixed_prefill_policy") + + +def _leftover(text: str): + return next((t for t in _LEFTOVER_TOKENS if t in text), None) + + def _unpatch(text: str, name, marker, head_token, pairs, sha=None, length=None, exact=None): """Invert one identity: exact hook insertions, then its authenticated helper site. @@ -1313,9 +1323,8 @@ def _unpatch(text: str, name, marker, head_token, pairs, sha=None, length=None, _refuse(f"{name}: unexpected text between the helper and its anchor") clean = text[:start] + text[start + n:] # Fail closed on mixed state: after removing this identity, no decode-floor - # marker of any version and no helper definition of either form may remain. - if (marker in clean or head_token in clean or MARK.rstrip("]") in clean - or CLASS_HEAD in clean or V1_HELPER_START in clean): + # marker of any version and no helper definition or name may remain. + if marker in clean or head_token in clean or _leftover(clean): _refuse(f"{name}: marker or helper definition left after unpatch") return clean, start, span @@ -1364,6 +1373,10 @@ def _identify(text: str): else: if CLASS_HEAD in text: _refuse("unmarked _Glm53MixedPrefill helper present") + token = _leftover(text) + if token is not None: + _refuse(f"unrecognized decode-floor state ({token!r} present); " + "restore an authenticated/pristine scheduler first") return None reasons = [] for name, mark, head, pairs, sha, length in IDENTITIES: diff --git a/tests/test_scheduler_decode_floor.py b/tests/test_scheduler_decode_floor.py index e2a36a38..f5fde3bf 100644 --- a/tests/test_scheduler_decode_floor.py +++ b/tests/test_scheduler_decode_floor.py @@ -495,8 +495,8 @@ def installation_tests(): current = None if mod.MARK_V7 in image: # An already-current v7 install: verify it exactly (as main() does), then unpatch. - clean, _, _ = mod._unpatch(image, 'v7', mod.MARK_V7, mod.CLASS_HEAD, mod.V7_PAIRS, - exact=mod._helper_text()) + clean, current_at, _ = mod._unpatch(image, 'v7', mod.MARK_V7, mod.CLASS_HEAD, mod.V7_PAIRS, + exact=mod._helper_text()) found, current = None, image else: found = mod._identify(image) @@ -559,10 +559,30 @@ def run(text, temp): assert again.returncode == 0 and repeat == after and 'already present' in again.stdout, ident if current is not None: - assert current == installed, 'the installed v7 differs from a fresh install' + # A deployed v7 keeps its helper where it was installed (e.g. before a later + # overlay such as adaptive-K), so compare at that position, not the default anchor. + assert current == mod.apply_v7(clean, at=current_at), 'the installed v7 is not canonical' again, repeat = run(current, temp) assert again.returncode == 0 and repeat == current and 'already present' in again.stdout + # Composition with patch_adaptive_k, which inserts its class at the same anchor: + # the helper may sit before it (decode-floor applied first, as deployed) or after + # it, and both layouts verify as a no-op. + if mod.ADAPTIVE_K_HEAD in clean: # a deployed source already carries the real class + with_adaptive = clean + adaptive_at = clean.index(mod.ADAPTIVE_K_HEAD) + else: + adaptive = mod.ADAPTIVE_K_HEAD + ' # [glm53-adaptive-k]\n pass\n\n' + with_adaptive = clean.replace(mod.NEEDLE, adaptive + mod.NEEDLE, 1) + adaptive_at = with_adaptive.index(adaptive) + helper_after = mod.apply_v7(with_adaptive) + helper_before = mod.apply_v7(with_adaptive, at=adaptive_at) + assert helper_after.index(mod.CLASS_HEAD) > helper_after.index(mod.ADAPTIVE_K_HEAD) + assert helper_before.index(mod.CLASS_HEAD) < helper_before.index(mod.ADAPTIVE_K_HEAD) + for layout, text in (('helper-after-adaptive-k', helper_after), ('helper-before-adaptive-k', helper_before)): + again, repeat = run(text, temp) + assert again.returncode == 0 and repeat == text and 'already present' in again.stdout, layout + v5 = producers['v5-priority'] refused = { 'v3-marker': v5.replace(mod.MARK_V5, mod.MARK_V3), @@ -578,6 +598,14 @@ def run(text, temp): # Mixed state: a canonical v7 plus a leftover older marker must not verify. 'v7-plus-v5-marker': mod.MARK_V5 + '\n' + installed, 'v7-plus-v1-marker': installed + '\n' + mod.MARK + '\n', + # A later duplicate wrapper would rebind the global the hooks call. + 'v7-plus-duplicate-wrapper': installed + '\n\ndef _glm53_mixed_prefill_policy(sched, request, computed=None):\n return 0\n', + 'pristine-plus-wrapper': clean + '\n\ndef _glm53_mixed_prefill_policy(sched, request, computed=None):\n return 0\n', + # An unknown version is not pristine: refuse instead of installing beside it. + 'unknown-marker-v99': '# [glm53-decode-floor:v99]\n' + clean, + # Only the import anchor or adaptive-K may follow the helper. + 'v7-unknown-text-after-helper': installed.replace( + mod.NEEDLE, 'class _Glm53OtherOverlay:\n pass\n\n' + mod.NEEDLE, 1), } for case, text in refused.items(): assert text != installed and text != v5, f'{case}: mutation did not apply' From b852f57cd0e26f6526a1ed6ab494873b46c75dfb Mon Sep 17 00:00:00 2001 From: Hector Isaza <76170012+hecisaza@users.noreply.github.com> Date: Thu, 1 Oct 2026 02:44:24 +0000 Subject: [PATCH 6/8] tests: consolidate decode-floor producer fixtures into legacy_scheduler_helpers.py Replace the six vendored copies of older patcher scripts under tests/fixtures/decode_floor_producers/ with the byte-exact helper texts of the same identities (v1-image-d9758a6, v5-main, v5-priority, v6-warm-deadline, v6-carry), appended to tests/fixtures/legacy_scheduler_helpers.py in the format #180 established for v1/v2/v5-historical, with provenance in its docstring. installation_tests() rebuilds every accepted public identity the same way: helper text at the installer's anchor plus that version's hook pairs. Each rebuild was checked byte-identical to the output of the removed patcher copy, and the installer still authenticates each by sha256 and length. v5-historical is now exercised directly as well. Same coverage (9 of 11), about 3,600 fewer lines, no new file layout. --- .../v1_image_d9758a6_helper.py | 6 - .../v5_historical_8e2a32b2.py | 1113 --------- .../v5_main_d5713ec6.py | 1123 --------- .../v5_priority_763bd30e.py | 1152 --------- .../v6_carry_0870fab9.py | 1183 ---------- .../v6_warm_deadline_1f74558d.py | 1143 --------- tests/fixtures/legacy_scheduler_helpers.py | 2086 +++++++++++++++++ tests/test_scheduler_decode_floor.py | 29 +- 8 files changed, 2100 insertions(+), 5735 deletions(-) delete mode 100644 tests/fixtures/decode_floor_producers/v1_image_d9758a6_helper.py delete mode 100644 tests/fixtures/decode_floor_producers/v5_historical_8e2a32b2.py delete mode 100644 tests/fixtures/decode_floor_producers/v5_main_d5713ec6.py delete mode 100644 tests/fixtures/decode_floor_producers/v5_priority_763bd30e.py delete mode 100644 tests/fixtures/decode_floor_producers/v6_carry_0870fab9.py delete mode 100644 tests/fixtures/decode_floor_producers/v6_warm_deadline_1f74558d.py diff --git a/tests/fixtures/decode_floor_producers/v1_image_d9758a6_helper.py b/tests/fixtures/decode_floor_producers/v1_image_d9758a6_helper.py deleted file mode 100644 index 60367a08..00000000 --- a/tests/fixtures/decode_floor_producers/v1_image_d9758a6_helper.py +++ /dev/null @@ -1,6 +0,0 @@ -"""The decode-floor v1 helper baked into image ef9f5013 (identity v1-image-d9758a6). - -Inserted at the patcher's NEEDLE with the v1 hook pairs, it rebuilds that image's -scheduler.py byte-for-byte from the pristine source. -""" -HELPER = '\ndef _glm53_mixed_prefill_policy(running, current):\n """Mixed-step prefill policy when a peer in `running` is decoding.\n\n None = no extra policy. 0 = skip this prefill this step. N>0 = cap.\n """\n raw = os.environ.get("GLM53_MIXED_PREFILL_CHUNK", "0").strip().lower()\n if raw in ("0", "off", "no"):\n return None\n if raw in ("skip", "-1"):\n cap = 0\n else:\n try:\n cap = int(raw)\n except ValueError:\n cap = 0\n if cap <= 0:\n return None\n cur_id = getattr(current, "request_id", None)\n for r in running:\n if r is current or getattr(r, "request_id", None) == cur_id:\n continue\n if r.num_computed_tokens >= r.num_prompt_tokens:\n return cap\n return None\n\n\n' diff --git a/tests/fixtures/decode_floor_producers/v5_historical_8e2a32b2.py b/tests/fixtures/decode_floor_producers/v5_historical_8e2a32b2.py deleted file mode 100644 index ddfe522a..00000000 --- a/tests/fixtures/decode_floor_producers/v5_historical_8e2a32b2.py +++ /dev/null @@ -1,1113 +0,0 @@ -#!/usr/bin/env python3 -"""Mixed-prefill policy: skip / cap / off / fair (issue #6). - -A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover -MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill -chunk. Mixed execution leaves the uniform decode FULL-graph path, and a -128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond -budget. - -GLM53_MIXED_PREFILL_CHUNK: - skip / -1 — do not mix prefill with decode. Starves every - waiting/running prefill while any peer is decoding; there is - no age limit. Solo prefill is unchanged. - N>0 — cap mixed prefill chunks to N tokens while a peer decodes. - The cap is fed into hybrid Mamba alignment so N < block_size - still makes sub-block progress. 128 still stalls ~10 tok/s. - 0 / off — disable the extra isolation policy. - fair — service-time mixing (TP=2 default since 2026-09-15, v5; - opt-in on TP=3/TP=4). Decode-only - steps between prefill turns; at most - GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). - Only prefill that contends with a decoder is charged (solo - prefill is cost-sampled, not debt). Credit accrues at SHARE - of accounted engine time. v5 fits a fixed-plus-per-token step - cost from solo and mixed samples, targets the largest ladder - rung (128..2048) whose estimated step fits - GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that - rung instead of spending it on small chunks (v4 priced 1024 - tokens off 128-token samples linearly, ~3x too high, then - stalled at 128-token steps: ~70 tok/s at 20% share). A - never-served newcomer gets one prompt step-bounded probe; - afterwards a 2s age override may borrow one such chunk, only - after all shared debt is repaid. In-flight async prefill - blocks the next mixed turn. Prefills are selected by last - completed positive prefill service plus round-robin. Decode - token/input budget is allocated first by the base scheduler. - Solo prefill retains the base scheduler's limits. Timing is a - host busy-time proxy. - -Fair knobs (read at runtime; identical on every rank): - GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) - GLM53_FAIR_PREFILL_SHARE default 0.20 (credit accrual fraction) - GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 - GLM53_FAIR_PREFILL_MAX_STEP_MS default 1000 (estimated mixed-step limit) - GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 - -Versioned installer: `# [glm53-decode-floor:v5]`. v1 (no version), v2, v3 -and v4 images are unpatched then re-patched. Fail closed if anchors drift. -""" -from __future__ import annotations - -import inspect -import os -import sys -import time -from pathlib import Path - -P = Path( - os.environ.get( - "GLM53_SCHEDULER_PY", - "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", - ) -) -MARK = "# [glm53-decode-floor]" -MARK_V2 = "# [glm53-decode-floor:v2]" -MARK_V3 = "# [glm53-decode-floor:v3]" -MARK_V4 = "# [glm53-decode-floor:v4]" -MARK_V5 = "# [glm53-decode-floor:v5]" - -IMPORT_OLD = """import itertools -import time -""" -IMPORT_NEW = """import itertools -import os -import time -""" - -# v1 helper + insertions (recipe f906ee9 / this overlay before v2). -V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" -V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. -V2_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] - # NOTE(woosuk) on the scheduling algorithm: -""" -V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] - pooler_outputs = model_runner_output.pooler_output -""" -V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" -V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] -""" -V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] - if num_new_tokens == 0: - break -""" - - - -# Frozen v3 anchors for migration from the reviewed implementation. -V3_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v3] - # NOTE(woosuk) on the scheduling algorithm: -""" -V3_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v3] - pooler_outputs = model_runner_output.pooler_output -""" -V3_FIN_OLD = """ # Check if the scheduling constraints are satisfied. - total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) -""" -V3_FIN_NEW = """ # Check if the scheduling constraints are satisfied. - _GLM53_MIXED.note_schedule_output(self, num_scheduled_tokens) # [glm53-decode-floor:v3] - total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) -""" -V3_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( - num_new_tokens, - token_budget, - int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), - _GLM53_MIXED.needs_prefill_compute(request), - ) - - # Make sure the input position does not exceed the max model len. -""" -V3_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( - num_new_tokens, - token_budget, - int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), - _GLM53_MIXED.needs_prefill_compute(request), - ) - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if num_new_tokens <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - - # chunked prefill has to be enabled explicitly to allow -""" -V3_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v3] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V3_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] -""" -V3_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] - if num_new_tokens == 0: - break -""" - -class _Glm53MixedPrefill: # [glm53-decode-floor:v5] - """Bound contention using completion feedback, without synchronizing GPUs.""" - - LADDER = (128, 256, 512, 768, 1024, 1536, 2048) - COLD_TOK_S = 1300.0 - FIT_WINDOW = 24 - - def __init__(self, now=None): - self._now = now or time.monotonic - self.mode = "skip" - self.legacy_cap = 0 - self.logged_boot = False - self.hist_every = 50 - self._parse() - self.last_service = {} - self.arrival = {} - self.rr_seq = {} - self.served_tokens = {} - self.rr_n = 0 - self.steps = 0 - self.credit = 0.0 - self.in_contention = False - self.inflight = {} - self.mixed_samples = [] - self.solo_samples = [] - self.last_account_mono = None - self.last_prefill_turn_mono = 0.0 - self.step_tag = None - self._sched_id = None - self._open_rec = None - self.selected = set() - self._candidates = [] - self._tried = set() - self.step_mode = "solo" - self.defer_reason = "none" - self.missed_prefill = 0 - self._model_cache = None - - def _e(self, name, default): - v = os.environ.get(name) - return default if v is None or not str(v).strip() else str(v).strip() - - def _parse(self) -> None: - raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() - self.legacy_cap = 0 - if raw in ("0", "off", "no"): - self.mode = "off" - elif raw in ("skip", "-1"): - self.mode = "skip" - elif raw == "fair": - self.mode = "fair" - else: - try: - cap = int(raw) - except ValueError: - cap = 0 - if cap <= 0: - self.mode = "off" - else: - self.mode = "cap" - self.legacy_cap = cap - try: - self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) - except ValueError: - self.chunk = 256 - if self.chunk <= 0: - self.chunk = 256 - try: - self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.20")) - except ValueError: - self.share = 0.20 - self.share = min(1.0, max(0.0, self.share)) - try: - self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 - except ValueError: - self.interval_s = 2.0 - if self.interval_s <= 0: - self.interval_s = 2.0 - try: - self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) - except ValueError: - self.max_chunks = 1 - self.max_chunks = max(1, min(self.max_chunks, 16)) - try: - self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "1000")) / 1000.0) - except ValueError: - self.max_step_s = 1.0 - if self.mode == "fair" and not self.logged_boot: - print( - f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " - f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " - f"interval_s={self.interval_s} max_step_s={self.max_step_s} " - f"max_chunks={self.max_chunks}", - flush=True, - ) - self.logged_boot = True - - - @staticmethod - def prefill_remaining(request, computed=None): - prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) - if computed is None: - computed = int(getattr(request, "num_computed_tokens", 0) or 0) - tokens = int(getattr(request, "num_tokens", prompt) or prompt) - return max(0, max(prompt, tokens - 1) - int(computed)) - - def needs_prefill_compute(self, request, computed=None): - return self.prefill_remaining(request, computed) > 0 - - def _iter_waiting(self, sched): - for name in ("waiting", "skipped_waiting"): - q = getattr(sched, name, None) - if not q: - continue - try: - for r in q: - yield r - except TypeError: - continue - - def _live_ids(self, sched): - ids = set() - for r in list(getattr(sched, "running", None) or []): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - for r in self._iter_waiting(sched): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - reqs = getattr(sched, "requests", None) - if isinstance(reqs, dict): - ids.update(reqs.keys()) - return ids - - def _prune(self, live): - for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): - dead = [k for k in store if k not in live] - for k in dead: - store.pop(k, None) - - - @property - def inflight_prefill(self): - return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) - - def _shape(self, decodes, prefills): - history = max((int(r.num_computed_tokens) for r in decodes), default=0) - position = max((int(r.num_computed_tokens) for r in prefills), default=0) - return (len(decodes), (history // 4096).bit_length(), - (position // 4096).bit_length()) - - def _cost_model(self): - """Fit step cost dt = a + b*n over recent prefill-bearing steps. - - Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host - time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to - a solo chunk plus a few decode rows: solo samples are pooled for the - fit. The fit is then scaled so recent mixed samples are not - underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). - Needs two distinct chunk sizes; otherwise returns None and the caller - scales linearly. v4 scaled one sample linearly, priced 1024 tokens off - 128-token samples at ~3x the real cost, and never climbed back. - """ - if self._model_cache is not None: - return self._model_cache - mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] - pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] - if len(pts) < 2 or len({n for n, _ in pts}) < 2: - return None - cnt = float(len(pts)) - sx = float(sum(n for n, _ in pts)) - sy = float(sum(dt for _, dt in pts)) - sxx = float(sum(n * n for n, _ in pts)) - sxy = float(sum(n * dt for n, dt in pts)) - den = cnt * sxx - sx * sx - b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 - a = max(0.0, (sy - b * sx) / cnt) - if mixed: - ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) - r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] - else: - r = 1.1 # decode rows not sampled yet - r = min(1.5, max(1.0, r)) - self._model_cache = (a * r, b * r) - return self._model_cache - - def _est_dt(self, tokens, shape=None): - """Conservative host-time estimate; not a guaranteed execution bound.""" - n = max(1, int(tokens)) - model = self._cost_model() - if model is not None: - a, b = model - return max(0.01, a + b * n) - samples = self.mixed_samples[-8:] - if samples: - # One size only: scale linearly with a fixed-cost floor. - return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) - return max(0.05, n / self.COLD_TOK_S) - - def _target(self, remaining, room): - """Largest rung whose estimated step fits `room` (tokens per accounted - second rise with size under a fixed per-step cost), or None.""" - rungs = {self.chunk} - if self.mixed_samples or self._cost_model() is not None: - rungs.update(self.LADDER) - if remaining is not None: - rungs = {min(n, remaining) for n in rungs} - fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] - fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] - return fitting[-1] if fitting else None - - def _credit_limit(self): - return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) - - def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - self.last_service.get(r.request_id, 0.0), - self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) - - def _promote_next(self): - grants = (self._open_rec or {}).get("grants", {}) - for r in self._candidates: - if len(self.selected | set(grants)) >= self.max_chunks: - break - if r.request_id not in self._tried: - self.selected.add(r.request_id) - - def _release(self, rid, reason="allocation"): - rec = self._open_rec - if rec is None: - return - grant = rec["grants"].pop(rid, None) - if grant: - self.credit += grant[1] - if grant[2]: - rec["borrowed"] = False - self.selected.discard(rid) - self._tried.add(rid) - self.defer_reason = reason - self._promote_next() - - def note_scheduled(self, request, num_new_tokens): - # Called after alignment/encoder caps; final allocation is sealed below. - if num_new_tokens <= 0: - self._release(request.request_id, "zero_progress") - - def protect_decode(self, request): - return (self.mode == "fair" and self._open_rec is not None - and self._open_rec["had_decode"] - and self.needs_prefill_compute(request)) - - def begin_step(self, sched): - now = self._now() - sid = id(sched) - tag = (sid, int(sched.current_step)) - if self.step_tag == tag: - return - if self._sched_id is not None and self._sched_id != sid: - self.inflight.clear() - self._open_rec = None - self.credit = 0.0 - self.in_contention = False - self.last_account_mono = None - self._sched_id = sid - # An unfinished/failed schedule has dispatched nothing: refund its grants. - if self._open_rec: - self.credit += sum(g[1] for g in self._open_rec["grants"].values()) - self.step_tag = tag - self.steps += 1 - self._prune(self._live_ids(sched)) - running = list(sched.running) - waiting = list(self._iter_waiting(sched)) - prefills = list({r.request_id: r for r in running + waiting - if self.needs_prefill_compute(r)}.values()) - decodes = [r for r in running if not self.needs_prefill_compute(r)] - # The base loop checks eligibility and reserves BOTH token and input/draft - # capacity by executing these requests first. Preserve order within groups. - if self.mode == "fair": - sched.running[:] = decodes + [r for r in running - if self.needs_prefill_compute(r)] - for r in running + waiting: - self.arrival.setdefault(r.request_id, now) - self._candidates = self._rank_prefills(prefills) - self._tried = set() - self.selected = set() - sched._glm53_align_prefill_limit = None - self._open_rec = { - "step_id": int(sched.current_step), "t_submit": now, - "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), - "shape": self._shape(decodes, prefills), "grants": {}, - "borrowed": False, - } - # Do not forgive debt while a decoder or its outstanding work remains. - if not decodes and not any(r["had_decode"] for r in self.inflight.values()): - self.in_contention = False - self.credit = 0.0 - self.step_mode = "legacy" if self.mode != "fair" else "solo" - self.defer_reason = "none" - if self.mode != "fair" or not prefills or not decodes: - return - if not self.in_contention: - self.credit = min(self.max_step_s, self._est_dt(self.chunk)) - self.in_contention = True - if self.inflight_prefill: - self.step_mode = "decode_only" - self.defer_reason = "async_inflight" - else: - self.step_mode = "prefill_turn" - self._promote_next() - self._maybe_log() - - def cap_for(self, sched, request, computed=None): - self.begin_step(sched) - sched._glm53_align_prefill_limit = None - remaining = self.prefill_remaining(request, computed) - if remaining <= 0 or self.mode == "off": - return None - peer_decode = any(r is not request and not self.needs_prefill_compute(r) - for r in sched.running) - if self.mode == "skip": - return 0 if peer_decode else None - if self.mode == "cap": - cap = self.legacy_cap if peer_decode else None - sched._glm53_align_prefill_limit = cap - return cap - if self.step_mode == "solo": - return None - rid = request.request_id - rec = self._open_rec - if rid in rec["grants"]: - return rec["grants"][rid][0] - if self.step_mode != "prefill_turn" or rid not in self.selected: - return 0 - reserved = sum(g[1] for g in rec["grants"].values()) - gap_room = max(0.0, self.max_step_s - reserved) - age = self._now() - self.last_service.get(rid, self.arrival[rid]) - pick = self._target(remaining, gap_room) - if pick is None: - self._release(rid, "gap_budget") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - # Save credit for the most efficient rung that fits the step budget - # instead of spending it on a smaller chunk now (v4 did, and stalled - # at 128-token steps once one expensive sample priced 256 too high). - cap, cost = pick - borrowed = False - if cost > self.credit + 1e-9: - # A never-served newcomer gets a prompt probe; afterwards service - # ages. Either may borrow ONE step-bounded chunk globally, only - # after all shared debt is repaid: queue churn or several aged - # newcomers cannot repeatedly overdraw the decoder. - due = rid not in self.last_service or age >= self.interval_s - if (due and self.credit >= -1e-9 and not rec["grants"] - and not rec["borrowed"]): - borrowed = True - else: - self._release(rid, "credit") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - self.credit -= cost - rec["grants"][rid] = (cap, cost, borrowed) - rec["borrowed"] |= borrowed - self._tried.add(rid) - sched._glm53_align_prefill_limit = cap - return cap - - def finish_step(self, sched, scheduler_output): - """Seal final work before _update_after_schedule advances request state.""" - rec = self._open_rec - self._open_rec = None - if rec is None: - return - scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() - if n > 0} - requests = getattr(sched, "requests", {}) - prefill = {} - for rid, n in scheduled.items(): - request = requests.get(rid) - if request is not None: - amount = min(n, self.prefill_remaining(request)) - if amount > 0: - prefill[rid] = amount - reserved = 0.0 - for rid, (_, estimate, _) in rec["grants"].items(): - # Never publish/charge a tentative grant that alignment, allocation, - # encoder caps or preemption removed from the final scheduler output. - actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 - self.credit += estimate - actual_est - reserved += actual_est - if not scheduled: - # Empty schedules do not necessarily have a completion callback. - return - rec.update(output=scheduler_output, scheduled=scheduled, - prefill_tokens=prefill, reserved=reserved) - del rec["grants"] - if self.mode == "fair": - self.inflight[id(scheduler_output)] = rec - - def observe_output(self, sched, scheduler_output): - rec = self.inflight.pop(id(scheduler_output), None) - if rec is None or rec["output"] is not scheduler_output: - return # Empty, duplicate or unrelated callback; never pop another step. - now = self._now() - # Partition observed busy wall time instead of adding overlapping - # submit-to-completion latencies of queued async steps. Queue residence - # and engine idle gaps therefore cannot mint decode credit twice. - start = rec["t_submit"] - if self.last_account_mono is not None: - start = max(start, self.last_account_mono) - dt = max(0.0, now - start) - self.last_account_mono = max(now, self.last_account_mono or now) - actual = scheduler_output.num_scheduled_tokens - served = {rid: min(n, max(0, int(actual.get(rid, 0)))) - for rid, n in rec["prefill_tokens"].items() - if int(actual.get(rid, 0)) > 0} - for rid, n in served.items(): - self.last_service[rid] = now - self.rr_n += 1 - self.rr_seq[rid] = self.rr_n - self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n - n = sum(served.values()) - if rec["had_decode"]: - # Reservation was already debited; settle once, including overruns. - self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) - self.credit = min(self.credit, self._credit_limit()) - if n and dt > 0: - self.mixed_samples.append((rec["shape"], n, dt)) - self.mixed_samples = self.mixed_samples[-64:] - self.last_prefill_turn_mono = now - self._model_cache = None - elif n and dt > 0: - self.solo_samples.append((n, dt)) - self.solo_samples = self.solo_samples[-32:] - self._model_cache = None - if n and self.hist_every > 0: - print(f"[glm53-decode-floor] completed_step={rec['step_id']} " - f"contention={int(rec['had_decode'])} prefill_tokens={served} " - f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " - f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) - self._prune(self._live_ids(sched)) - - def _maybe_log(self): - if self.hist_every <= 0 or self.steps % self.hist_every != 1: - return - rec = self._open_rec or {} - remaining = sum(self.prefill_remaining(r) for r in self._candidates) - target, estimate = (self._target(None, self.max_step_s) - or (self.chunk, self._est_dt(self.chunk))) - rate = target / estimate if estimate > 0 else 0.0 - eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") - model = self._cost_model() - fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" - if model else "fit=none") - print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " - f"mode={self.step_mode} defer={self.defer_reason} " - f"inflight={self.inflight_prefill} credit={self.credit:.3f} " - f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " - f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " - f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) - - @staticmethod - def aligned_new_tokens( - start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None - ) -> int: - """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" - if policy_cap is not None and policy_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, policy_cap) - end = start + num_new - if end < prefill_end: - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end - return max(0, end - start) - - - -def _helper_text() -> str: - body = inspect.getsource(_Glm53MixedPrefill) - return ( - "\n" - + body - + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5]\n\n" - + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5]\n" - + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" - ) - - -HELPER = None # filled at apply time so tests can call _helper_text() - - -BEGIN_OLD = """ self.current_step += 1 - # NOTE(woosuk) on the scheduling algorithm: -""" -BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] - # NOTE(woosuk) on the scheduling algorithm: -""" - -OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - pooler_outputs = model_runner_output.pooler_output -""" -OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] - pooler_outputs = model_runner_output.pooler_output -""" - -RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - - # Make sure the input position does not exceed the max model len. -""" -RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" - -WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - - # chunked prefill has to be enabled explicitly to allow -""" -WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" - -RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) -""" -RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] -""" - -WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - if num_new_tokens == 0: - break -""" -WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] - if num_new_tokens == 0: - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" -FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] - with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" - -RUNNING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled because one of the following -""" -RUNNING_ZERO_NEW = """ if num_new_tokens == 0: - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - # The request cannot be scheduled because one of the following -""" - -WAITING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled. - break -""" -WAITING_ZERO_NEW = """ if num_new_tokens == 0: - # The request cannot be scheduled. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. - # Preempt the lowest-priority request. -""" -PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. - if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] - break - # Preempt the lowest-priority request. -""" - -RUNNING_ALLOC_OLD = """ if new_blocks is None: - # Cannot schedule this request. - break -""" -RUNNING_ALLOC_NEW = """ if new_blocks is None: - # Cannot schedule this request. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - req_index += 1 - continue - break -""" - -WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - break - - # KVTransfer:""" -WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break - - # KVTransfer:""" - -V4_PAIRS = ( - (BEGIN_NEW, BEGIN_OLD, 'begin'), - (OBS_NEW, OBS_OLD, 'obs'), - (RUNNING_NEW, RUNNING_OLD, 'running'), - (WAITING_NEW, WAITING_OLD, 'waiting'), - (ALIGN_NEW, ALIGN_OLD, 'align'), - (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), - (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), - (FIN_NEW, FIN_OLD, 'fin'), - (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), - (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), - (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), - (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), - (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), -) - - -def replace_once(text: str, old: str, new: str, label: str) -> str: - n = text.count(old) - if n != 1: - raise SystemExit(f"{P}: expected one {label} target, found {n}") - return text.replace(old, new, 1) - - -def _strip_helper(text: str, label: str) -> str: - start = text.find("class _Glm53MixedPrefill:") - if start < 0: - raise SystemExit(f"{P}: {label} helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: {label} helper end not found") - end = min(candidates) - return text[:start] + text[end:] - - -def _strip_v1_helper(text: str) -> str: - start = text.find(V1_HELPER_START) - if start < 0: - raise SystemExit(f"{P}: v1 helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: v1 helper end not found") - end = min(candidates) - return text[:start] + "\n" + text[end:] - - -def unpatch_v1(text: str) -> str: - if V1_RUNNING_NEW in text: - text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 running insertion drifted") - if V1_WAITING_NEW in text: - text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 waiting insertion drifted") - if V1_HELPER_START in text: - text = _strip_v1_helper(text) - leftover = [ - "_glm53_mixed_prefill_policy(self.running, request)", - V1_HELPER_START, - ] - for s in leftover: - if s in text: - raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") - return text - - -def unpatch_v2(text: str) -> str: - pairs = ( - (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), - (V2_OBS_NEW, OBS_OLD, "v2-obs"), - (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), - (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), - (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), - (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), - (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), - ) - for new, old, label in pairs: - if new in text: - text = replace_once(text, new, old, label) - elif MARK_V2 in new: - # Some insertions may already have been removed; fail if marker remains. - pass - if "class _Glm53MixedPrefill:" in text: - text = _strip_helper(text, "v2") - if MARK_V2 in text: - raise SystemExit(f"{P}: v2 leftover after unpatch") - return text - - -def unpatch_v3(text: str) -> str: - pairs = ( - (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), - (V3_OBS_NEW, OBS_OLD, "v3-obs"), - (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), - (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), - (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), - (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), - (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), - (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), - ) - for new, old, label in pairs: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v3") - if MARK_V3 in text: - raise SystemExit(f"{P}: v3 leftover after unpatch") - return text - - -def unpatch_v4(text: str) -> str: - for new, old, label in V4_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v4") - if MARK_V4 in text: - raise SystemExit(f"{P}: v4 leftover after unpatch") - return text - - -# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 -# insertions above stay frozen so a v4 image can be unpatched exactly. -V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) - - -def unpatch_v5(text: str) -> str: - for new, old, label in V5_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v5") - if MARK_V5 in text: - raise SystemExit(f"{P}: v5 leftover after unpatch") - return text - - -def apply_v5(text: str) -> str: - if "import os\n" not in text.split("import time\n", 1)[0]: - text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") - needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" - text = replace_once(text, needle, _helper_text() + needle, "helper") - for new, old, label in V5_PAIRS: - text = replace_once(text, old, new, label) - compile(text, str(P), "exec") - return text - - -def main() -> int: - if not P.is_file(): - raise SystemExit(f"missing {P}") - text = P.read_text() - original = text - if MARK_V5 in text: - # Validate existing anchors/helper instead of trusting the marker alone. - clean = unpatch_v5(text) - if apply_v5(clean) != text: - raise SystemExit(f"{P}: v5 helper drifted") - print(f"{P.name}: {MARK_V5} already present — verified") - return 0 - if MARK_V4 in text: - text = unpatch_v4(text) - elif MARK_V3 in text: - text = unpatch_v3(text) - elif MARK_V2 in text: - text = unpatch_v2(text) - elif MARK in text or V1_HELPER_START in text: - text = unpatch_v1(text) - text = apply_v5(text) - if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text: - raise SystemExit(f"{P}: older marker left after migration") - if text != original: - P.write_text(text) - print(f"patched {P.name} ({MARK_V5})") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/fixtures/decode_floor_producers/v5_main_d5713ec6.py b/tests/fixtures/decode_floor_producers/v5_main_d5713ec6.py deleted file mode 100644 index 78c53778..00000000 --- a/tests/fixtures/decode_floor_producers/v5_main_d5713ec6.py +++ /dev/null @@ -1,1123 +0,0 @@ -#!/usr/bin/env python3 -"""Mixed-prefill policy: skip / cap / off / fair (issue #6). - -A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover -MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill -chunk. Mixed execution leaves the uniform decode FULL-graph path, and a -128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond -budget. - -GLM53_MIXED_PREFILL_CHUNK: - skip / -1 — do not mix prefill with decode. Starves every - waiting/running prefill while any peer is decoding; there is - no age limit. Solo prefill is unchanged. - N>0 — cap mixed prefill chunks to N tokens while a peer decodes. - The cap is fed into hybrid Mamba alignment so N < block_size - still makes sub-block progress. 128 still stalls ~10 tok/s. - 0 / off — disable the extra isolation policy. - fair — service-time mixing (default on TP=2/3/4 since 2026-09-15, - v5). Decode-only - steps between prefill turns; at most - GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). - Only prefill that contends with a decoder is charged (solo - prefill is cost-sampled, not debt). Credit accrues at SHARE - of accounted engine time. v5 fits a fixed-plus-per-token step - cost from solo and mixed samples, targets the largest ladder - rung (128..2048) whose estimated step fits - GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that - rung instead of spending it on small chunks (v4 priced 1024 - tokens off 128-token samples linearly, ~3x too high, then - stalled at 128-token steps: ~70 tok/s at 20% share). A - never-served newcomer gets one prompt step-bounded probe; - afterwards a 2s age override may borrow one such chunk, only - after all shared debt is repaid. In-flight async prefill - blocks the next mixed turn. Prefills are selected by last - completed positive prefill service plus round-robin. Decode - token/input budget is allocated first by the base scheduler. - Solo prefill retains the base scheduler's limits. Timing is a - host busy-time proxy. - -Fair knobs (read at runtime; identical on every rank): - GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) - GLM53_FAIR_PREFILL_SHARE default 0.30 (credit accrual fraction) - GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 - GLM53_FAIR_PREFILL_MAX_STEP_MS default 2000 (estimated mixed-step limit) - GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 - -Versioned installer: `# [glm53-decode-floor:v5]`. v1 (no version), v2, v3 -and v4 images are unpatched then re-patched. Fail closed if anchors drift. -""" -from __future__ import annotations - -import inspect -import os -import sys -import time -from pathlib import Path - -P = Path( - os.environ.get( - "GLM53_SCHEDULER_PY", - "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", - ) -) -MARK = "# [glm53-decode-floor]" -MARK_V2 = "# [glm53-decode-floor:v2]" -MARK_V3 = "# [glm53-decode-floor:v3]" -MARK_V4 = "# [glm53-decode-floor:v4]" -MARK_V5 = "# [glm53-decode-floor:v5]" - -IMPORT_OLD = """import itertools -import time -""" -IMPORT_NEW = """import itertools -import os -import time -""" - -# v1 helper + insertions (recipe f906ee9 / this overlay before v2). -V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" -V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. -V2_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] - # NOTE(woosuk) on the scheduling algorithm: -""" -V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] - pooler_outputs = model_runner_output.pooler_output -""" -V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" -V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] -""" -V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] - if num_new_tokens == 0: - break -""" - - - -# Frozen v3 anchors for migration from the reviewed implementation. -V3_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v3] - # NOTE(woosuk) on the scheduling algorithm: -""" -V3_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v3] - pooler_outputs = model_runner_output.pooler_output -""" -V3_FIN_OLD = """ # Check if the scheduling constraints are satisfied. - total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) -""" -V3_FIN_NEW = """ # Check if the scheduling constraints are satisfied. - _GLM53_MIXED.note_schedule_output(self, num_scheduled_tokens) # [glm53-decode-floor:v3] - total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) -""" -V3_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( - num_new_tokens, - token_budget, - int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), - _GLM53_MIXED.needs_prefill_compute(request), - ) - - # Make sure the input position does not exceed the max model len. -""" -V3_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( - num_new_tokens, - token_budget, - int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), - _GLM53_MIXED.needs_prefill_compute(request), - ) - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if num_new_tokens <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - - # chunked prefill has to be enabled explicitly to allow -""" -V3_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v3] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V3_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] -""" -V3_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] - if num_new_tokens == 0: - break -""" - -class _Glm53MixedPrefill: # [glm53-decode-floor:v5] - """Bound contention using completion feedback, without synchronizing GPUs.""" - - LADDER = (128, 256, 512, 768, 1024, 1536, 2048) - COLD_TOK_S = 1300.0 - FIT_WINDOW = 24 - - def __init__(self, now=None): - self._now = now or time.monotonic - self.mode = "skip" - self.legacy_cap = 0 - self.logged_boot = False - self.hist_every = 50 - self._parse() - self.last_service = {} - self.arrival = {} - self.rr_seq = {} - self.served_tokens = {} - self.rr_n = 0 - self.steps = 0 - self.credit = 0.0 - self.in_contention = False - self.inflight = {} - self.mixed_samples = [] - self.solo_samples = [] - self.last_account_mono = None - self.last_prefill_turn_mono = 0.0 - self.step_tag = None - self._sched_id = None - self._open_rec = None - self.selected = set() - self._candidates = [] - self._tried = set() - self.step_mode = "solo" - self.defer_reason = "none" - self.missed_prefill = 0 - self._model_cache = None - - def _e(self, name, default): - v = os.environ.get(name) - return default if v is None or not str(v).strip() else str(v).strip() - - def _parse(self) -> None: - raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() - self.legacy_cap = 0 - if raw in ("0", "off", "no"): - self.mode = "off" - elif raw in ("skip", "-1"): - self.mode = "skip" - elif raw == "fair": - self.mode = "fair" - else: - try: - cap = int(raw) - except ValueError: - cap = 0 - if cap <= 0: - self.mode = "off" - else: - self.mode = "cap" - self.legacy_cap = cap - try: - self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) - except ValueError: - self.chunk = 256 - if self.chunk <= 0: - self.chunk = 256 - try: - self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) - except ValueError: - self.share = 0.30 - self.share = min(1.0, max(0.0, self.share)) - try: - self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 - except ValueError: - self.interval_s = 2.0 - if self.interval_s <= 0: - self.interval_s = 2.0 - try: - self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) - except ValueError: - self.max_chunks = 1 - self.max_chunks = max(1, min(self.max_chunks, 16)) - try: - self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) - except ValueError: - self.max_step_s = 2.0 - if self.mode == "fair" and not self.logged_boot: - print( - f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " - f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " - f"interval_s={self.interval_s} max_step_s={self.max_step_s} " - f"max_chunks={self.max_chunks}", - flush=True, - ) - self.logged_boot = True - - - @staticmethod - def prefill_remaining(request, computed=None): - prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) - if computed is None: - computed = int(getattr(request, "num_computed_tokens", 0) or 0) - tokens = int(getattr(request, "num_tokens", prompt) or prompt) - return max(0, max(prompt, tokens - 1) - int(computed)) - - def needs_prefill_compute(self, request, computed=None): - return self.prefill_remaining(request, computed) > 0 - - def _iter_waiting(self, sched): - for name in ("waiting", "skipped_waiting"): - q = getattr(sched, name, None) - if not q: - continue - try: - for r in q: - yield r - except TypeError: - continue - - def _live_ids(self, sched): - ids = set() - for r in list(getattr(sched, "running", None) or []): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - for r in self._iter_waiting(sched): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - reqs = getattr(sched, "requests", None) - if isinstance(reqs, dict): - ids.update(reqs.keys()) - return ids - - def _prune(self, live): - for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): - dead = [k for k in store if k not in live] - for k in dead: - store.pop(k, None) - - - @property - def inflight_prefill(self): - return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) - - def _shape(self, decodes, prefills): - history = max((int(r.num_computed_tokens) for r in decodes), default=0) - position = max((int(r.num_computed_tokens) for r in prefills), default=0) - return (len(decodes), (history // 4096).bit_length(), - (position // 4096).bit_length()) - - def _cost_model(self): - """Fit step cost dt = a + b*n over recent prefill-bearing steps. - - Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host - time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to - a solo chunk plus a few decode rows: solo samples are pooled for the - fit. The fit is then scaled so recent mixed samples are not - underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). - Needs two distinct chunk sizes; otherwise returns None and the caller - scales linearly. v4 scaled one sample linearly, priced 1024 tokens off - 128-token samples at ~3x the real cost, and never climbed back. - """ - if self._model_cache is not None: - return self._model_cache - mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] - pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] - if len(pts) < 2 or len({n for n, _ in pts}) < 2: - return None - cnt = float(len(pts)) - sx = float(sum(n for n, _ in pts)) - sy = float(sum(dt for _, dt in pts)) - sxx = float(sum(n * n for n, _ in pts)) - sxy = float(sum(n * dt for n, dt in pts)) - den = cnt * sxx - sx * sx - b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 - a = max(0.0, (sy - b * sx) / cnt) - if mixed: - ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) - r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] - else: - r = 1.1 # decode rows not sampled yet - r = min(1.5, max(1.0, r)) - self._model_cache = (a * r, b * r) - return self._model_cache - - def _est_dt(self, tokens, shape=None): - """Conservative host-time estimate; not a guaranteed execution bound.""" - n = max(1, int(tokens)) - model = self._cost_model() - if model is not None: - a, b = model - return max(0.01, a + b * n) - samples = self.mixed_samples[-8:] - if samples: - # One size only: scale linearly with a fixed-cost floor. - return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) - return max(0.05, n / self.COLD_TOK_S) - - def _target(self, remaining, room): - """Largest rung whose estimated step fits `room` (tokens per accounted - second rise with size under a fixed per-step cost), or None.""" - rungs = {self.chunk} - if self.mixed_samples or self._cost_model() is not None: - rungs.update(self.LADDER) - if remaining is not None: - rungs = {min(n, remaining) for n in rungs} - fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] - fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] - return fitting[-1] if fitting else None - - def _credit_limit(self): - return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) - - def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - self.last_service.get(r.request_id, 0.0), - self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) - - def _promote_next(self): - grants = (self._open_rec or {}).get("grants", {}) - for r in self._candidates: - if len(self.selected | set(grants)) >= self.max_chunks: - break - if r.request_id not in self._tried: - self.selected.add(r.request_id) - - def _release(self, rid, reason="allocation"): - rec = self._open_rec - if rec is None: - return - grant = rec["grants"].pop(rid, None) - if grant: - self.credit += grant[1] - if grant[2]: - rec["borrowed"] = False - self.selected.discard(rid) - self._tried.add(rid) - self.defer_reason = reason - self._promote_next() - - def note_scheduled(self, request, num_new_tokens): - # Called after alignment/encoder caps; final allocation is sealed below. - if num_new_tokens <= 0: - self._release(request.request_id, "zero_progress") - - def protect_decode(self, request): - return (self.mode == "fair" and self._open_rec is not None - and self._open_rec["had_decode"] - and self.needs_prefill_compute(request)) - - def begin_step(self, sched): - now = self._now() - sid = id(sched) - tag = (sid, int(sched.current_step)) - if self.step_tag == tag: - return - if self._sched_id is not None and self._sched_id != sid: - self.inflight.clear() - self._open_rec = None - self.credit = 0.0 - self.in_contention = False - self.last_account_mono = None - self._sched_id = sid - # An unfinished/failed schedule has dispatched nothing: refund its grants. - if self._open_rec: - self.credit += sum(g[1] for g in self._open_rec["grants"].values()) - self.step_tag = tag - self.steps += 1 - self._prune(self._live_ids(sched)) - running = list(sched.running) - waiting = list(self._iter_waiting(sched)) - prefills = list({r.request_id: r for r in running + waiting - if self.needs_prefill_compute(r)}.values()) - decodes = [r for r in running if not self.needs_prefill_compute(r)] - # The base loop checks eligibility and reserves BOTH token and input/draft - # capacity by executing these requests first. Preserve order within groups. - if self.mode == "fair": - sched.running[:] = decodes + [r for r in running - if self.needs_prefill_compute(r)] - for r in running + waiting: - self.arrival.setdefault(r.request_id, now) - self._candidates = self._rank_prefills(prefills) - self._tried = set() - self.selected = set() - sched._glm53_align_prefill_limit = None - self._open_rec = { - "step_id": int(sched.current_step), "t_submit": now, - "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), - "shape": self._shape(decodes, prefills), "grants": {}, - "borrowed": False, - } - # Do not forgive debt while a decoder or its outstanding work remains. - if not decodes and not any(r["had_decode"] for r in self.inflight.values()): - self.in_contention = False - self.credit = 0.0 - self.step_mode = "legacy" if self.mode != "fair" else "solo" - self.defer_reason = "none" - if self.mode != "fair" or not prefills or not decodes: - return - if not self.in_contention: - self.credit = min(self.max_step_s, self._est_dt(self.chunk)) - self.in_contention = True - if self.inflight_prefill: - self.step_mode = "decode_only" - self.defer_reason = "async_inflight" - else: - self.step_mode = "prefill_turn" - self._promote_next() - self._maybe_log() - - def cap_for(self, sched, request, computed=None): - self.begin_step(sched) - sched._glm53_align_prefill_limit = None - remaining = self.prefill_remaining(request, computed) - if remaining <= 0 or self.mode == "off": - return None - peer_decode = any(r is not request and not self.needs_prefill_compute(r) - for r in sched.running) - if self.mode == "skip": - return 0 if peer_decode else None - if self.mode == "cap": - cap = self.legacy_cap if peer_decode else None - sched._glm53_align_prefill_limit = cap - return cap - if self.step_mode == "solo": - return None - rid = request.request_id - rec = self._open_rec - if rid in rec["grants"]: - return rec["grants"][rid][0] - if self.step_mode != "prefill_turn" or rid not in self.selected: - return 0 - reserved = sum(g[1] for g in rec["grants"].values()) - gap_room = max(0.0, self.max_step_s - reserved) - age = self._now() - self.last_service.get(rid, self.arrival[rid]) - pick = self._target(remaining, gap_room) - if pick is None: - self._release(rid, "gap_budget") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - # Save credit for the most efficient rung that fits the step budget - # instead of spending it on a smaller chunk now (v4 did, and stalled - # at 128-token steps once one expensive sample priced 256 too high). - cap, cost = pick - borrowed = False - if cost > self.credit + 1e-9: - # A never-served newcomer gets a prompt probe; afterwards service - # ages. Either may borrow ONE step-bounded chunk globally, only - # after all shared debt is repaid: queue churn or several aged - # newcomers cannot repeatedly overdraw the decoder. - due = rid not in self.last_service or age >= self.interval_s - if (due and self.credit >= -1e-9 and not rec["grants"] - and not rec["borrowed"]): - borrowed = True - else: - self._release(rid, "credit") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - self.credit -= cost - rec["grants"][rid] = (cap, cost, borrowed) - rec["borrowed"] |= borrowed - self._tried.add(rid) - sched._glm53_align_prefill_limit = cap - return cap - - def finish_step(self, sched, scheduler_output): - """Seal final work before _update_after_schedule advances request state.""" - rec = self._open_rec - self._open_rec = None - if rec is None: - return - scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() - if n > 0} - requests = getattr(sched, "requests", {}) - prefill = {} - for rid, n in scheduled.items(): - request = requests.get(rid) - if request is not None: - amount = min(n, self.prefill_remaining(request)) - if amount > 0: - prefill[rid] = amount - reserved = 0.0 - for rid, (_, estimate, _) in rec["grants"].items(): - # Never publish/charge a tentative grant that alignment, allocation, - # encoder caps or preemption removed from the final scheduler output. - actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 - self.credit += estimate - actual_est - reserved += actual_est - if not scheduled: - # Empty schedules do not necessarily have a completion callback. - return - rec.update(output=scheduler_output, scheduled=scheduled, - prefill_tokens=prefill, reserved=reserved) - del rec["grants"] - if self.mode == "fair": - self.inflight[id(scheduler_output)] = rec - - def observe_output(self, sched, scheduler_output): - rec = self.inflight.pop(id(scheduler_output), None) - if rec is None or rec["output"] is not scheduler_output: - return # Empty, duplicate or unrelated callback; never pop another step. - now = self._now() - # Partition observed busy wall time instead of adding overlapping - # submit-to-completion latencies of queued async steps. Queue residence - # and engine idle gaps therefore cannot mint decode credit twice. - start = rec["t_submit"] - if self.last_account_mono is not None: - start = max(start, self.last_account_mono) - dt = max(0.0, now - start) - self.last_account_mono = max(now, self.last_account_mono or now) - actual = scheduler_output.num_scheduled_tokens - served = {rid: min(n, max(0, int(actual.get(rid, 0)))) - for rid, n in rec["prefill_tokens"].items() - if int(actual.get(rid, 0)) > 0} - for rid, n in served.items(): - self.last_service[rid] = now - self.rr_n += 1 - self.rr_seq[rid] = self.rr_n - self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n - n = sum(served.values()) - if rec["had_decode"]: - # Reservation was already debited; settle once, including overruns. - self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) - self.credit = min(self.credit, self._credit_limit()) - if n and dt > 0: - self.mixed_samples.append((rec["shape"], n, dt)) - self.mixed_samples = self.mixed_samples[-64:] - self.last_prefill_turn_mono = now - self._model_cache = None - elif n and dt > 0: - self.solo_samples.append((n, dt)) - self.solo_samples = self.solo_samples[-32:] - self._model_cache = None - if n and self.hist_every > 0: - print(f"[glm53-decode-floor] completed_step={rec['step_id']} " - f"contention={int(rec['had_decode'])} prefill_tokens={served} " - f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " - f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) - self._prune(self._live_ids(sched)) - - def _maybe_log(self): - if self.hist_every <= 0 or self.steps % self.hist_every != 1: - return - rec = self._open_rec or {} - remaining = sum(self.prefill_remaining(r) for r in self._candidates) - target, estimate = (self._target(None, self.max_step_s) - or (self.chunk, self._est_dt(self.chunk))) - rate = target / estimate if estimate > 0 else 0.0 - eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") - model = self._cost_model() - fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" - if model else "fit=none") - print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " - f"mode={self.step_mode} defer={self.defer_reason} " - f"inflight={self.inflight_prefill} credit={self.credit:.3f} " - f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " - f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " - f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) - - @staticmethod - def aligned_new_tokens( - start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None - ) -> int: - """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" - if policy_cap is not None and policy_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, policy_cap) - end = start + num_new - if end < prefill_end: - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end - return max(0, end - start) - - - -def _helper_text() -> str: - body = inspect.getsource(_Glm53MixedPrefill) - return ( - "\n" - + body - + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5]\n\n" - + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5]\n" - + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" - ) - - -HELPER = None # filled at apply time so tests can call _helper_text() - - -BEGIN_OLD = """ self.current_step += 1 - # NOTE(woosuk) on the scheduling algorithm: -""" -BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] - # NOTE(woosuk) on the scheduling algorithm: -""" - -OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - pooler_outputs = model_runner_output.pooler_output -""" -OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] - pooler_outputs = model_runner_output.pooler_output -""" - -RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - - # Make sure the input position does not exceed the max model len. -""" -RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" - -WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - - # chunked prefill has to be enabled explicitly to allow -""" -WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" - -RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) -""" -RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] -""" - -WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - if num_new_tokens == 0: - break -""" -WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] - if num_new_tokens == 0: - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" -FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] - with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" - -RUNNING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled because one of the following -""" -RUNNING_ZERO_NEW = """ if num_new_tokens == 0: - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - # The request cannot be scheduled because one of the following -""" - -WAITING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled. - break -""" -WAITING_ZERO_NEW = """ if num_new_tokens == 0: - # The request cannot be scheduled. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. - # Preempt the lowest-priority request. -""" -PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. - if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] - break - # Preempt the lowest-priority request. -""" - -RUNNING_ALLOC_OLD = """ if new_blocks is None: - # Cannot schedule this request. - break -""" -RUNNING_ALLOC_NEW = """ if new_blocks is None: - # Cannot schedule this request. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - req_index += 1 - continue - break -""" - -WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - break - - # KVTransfer:""" -WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break - - # KVTransfer:""" - -V4_PAIRS = ( - (BEGIN_NEW, BEGIN_OLD, 'begin'), - (OBS_NEW, OBS_OLD, 'obs'), - (RUNNING_NEW, RUNNING_OLD, 'running'), - (WAITING_NEW, WAITING_OLD, 'waiting'), - (ALIGN_NEW, ALIGN_OLD, 'align'), - (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), - (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), - (FIN_NEW, FIN_OLD, 'fin'), - (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), - (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), - (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), - (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), - (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), -) - - -def replace_once(text: str, old: str, new: str, label: str) -> str: - n = text.count(old) - if n != 1: - raise SystemExit(f"{P}: expected one {label} target, found {n}") - return text.replace(old, new, 1) - - -def _strip_helper(text: str, label: str, *, expected: str | None = None) -> str: - start = text.find("class _Glm53MixedPrefill:") - if start < 0: - raise SystemExit(f"{P}: {label} helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: {label} helper end not found") - end = min(candidates) - if expected is not None and text[start:end].strip() != expected.strip(): - raise SystemExit(f"{P}: {label} helper drifted") - return text[:start] + text[end:] - - -def _strip_v1_helper(text: str) -> str: - start = text.find(V1_HELPER_START) - if start < 0: - raise SystemExit(f"{P}: v1 helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: v1 helper end not found") - end = min(candidates) - return text[:start] + "\n" + text[end:] - - -def unpatch_v1(text: str) -> str: - if V1_RUNNING_NEW in text: - text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 running insertion drifted") - if V1_WAITING_NEW in text: - text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 waiting insertion drifted") - if V1_HELPER_START in text: - text = _strip_v1_helper(text) - leftover = [ - "_glm53_mixed_prefill_policy(self.running, request)", - V1_HELPER_START, - ] - for s in leftover: - if s in text: - raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") - return text - - -def unpatch_v2(text: str) -> str: - pairs = ( - (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), - (V2_OBS_NEW, OBS_OLD, "v2-obs"), - (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), - (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), - (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), - (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), - (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), - ) - for new, old, label in pairs: - if new in text: - text = replace_once(text, new, old, label) - elif MARK_V2 in new: - # Some insertions may already have been removed; fail if marker remains. - pass - if "class _Glm53MixedPrefill:" in text: - text = _strip_helper(text, "v2") - if MARK_V2 in text: - raise SystemExit(f"{P}: v2 leftover after unpatch") - return text - - -def unpatch_v3(text: str) -> str: - pairs = ( - (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), - (V3_OBS_NEW, OBS_OLD, "v3-obs"), - (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), - (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), - (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), - (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), - (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), - (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), - ) - for new, old, label in pairs: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v3") - if MARK_V3 in text: - raise SystemExit(f"{P}: v3 leftover after unpatch") - return text - - -def unpatch_v4(text: str) -> str: - for new, old, label in V4_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v4") - if MARK_V4 in text: - raise SystemExit(f"{P}: v4 leftover after unpatch") - return text - - -# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 -# insertions above stay frozen so a v4 image can be unpatched exactly. -V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) - - -def unpatch_v5(text: str) -> str: - for new, old, label in V5_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v5", expected=_helper_text()) - if MARK_V5 in text: - raise SystemExit(f"{P}: v5 leftover after unpatch") - return text - - -def apply_v5(text: str) -> str: - if "import os\n" not in text.split("import time\n", 1)[0]: - text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") - needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" - text = replace_once(text, needle, _helper_text() + needle, "helper") - for new, old, label in V5_PAIRS: - text = replace_once(text, old, new, label) - compile(text, str(P), "exec") - return text - - -def main() -> int: - if not P.is_file(): - raise SystemExit(f"missing {P}") - text = P.read_text() - original = text - if MARK_V5 in text: - # Validate existing anchors/helper instead of trusting the marker alone. - # unpatch_v5 checks every v5 insertion occurs exactly once, strips the - # helper, and rejects leftover markers. The result is discarded: a - # later overlay may legitimately sit between the helper and the - # cuda_graph import anchor, so re-applying at that fixed anchor would - # relocate the helper and fail a byte-compare on a healthy file. - unpatch_v5(text) - # The import edit is part of the applied state; the byte-compare used - # to cover it implicitly. - if "import os\n" not in text.split("import time\n", 1)[0]: - raise SystemExit(f"{P}: v5 import drifted") - compile(text, str(P), "exec") - print(f"{P.name}: {MARK_V5} already present — verified") - return 0 - if MARK_V4 in text: - text = unpatch_v4(text) - elif MARK_V3 in text: - text = unpatch_v3(text) - elif MARK_V2 in text: - text = unpatch_v2(text) - elif MARK in text or V1_HELPER_START in text: - text = unpatch_v1(text) - text = apply_v5(text) - if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text: - raise SystemExit(f"{P}: older marker left after migration") - if text != original: - P.write_text(text) - print(f"patched {P.name} ({MARK_V5})") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/fixtures/decode_floor_producers/v5_priority_763bd30e.py b/tests/fixtures/decode_floor_producers/v5_priority_763bd30e.py deleted file mode 100644 index 1019e530..00000000 --- a/tests/fixtures/decode_floor_producers/v5_priority_763bd30e.py +++ /dev/null @@ -1,1152 +0,0 @@ -#!/usr/bin/env python3 -"""Mixed-prefill policy: skip / cap / off / fair (issue #6). - -A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover -MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill -chunk. Mixed execution leaves the uniform decode FULL-graph path, and a -128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond -budget. - -GLM53_MIXED_PREFILL_CHUNK: - skip / -1 — do not mix prefill with decode. Starves every - waiting/running prefill while any peer is decoding; there is - no age limit. Solo prefill is unchanged. - N>0 — cap mixed prefill chunks to N tokens while a peer decodes. - The cap is fed into hybrid Mamba alignment so N < block_size - still makes sub-block progress. 128 still stalls ~10 tok/s. - 0 / off — disable the extra isolation policy. - fair — service-time mixing (default on TP=2/3/4 since 2026-09-15, - v5). Decode-only - steps between prefill turns; at most - GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). - Only prefill that contends with a decoder is charged (solo - prefill is cost-sampled, not debt). Credit accrues at SHARE - of accounted engine time. v5 fits a fixed-plus-per-token step - cost from solo and mixed samples, targets the largest ladder - rung (128..2048) whose estimated step fits - GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that - rung instead of spending it on small chunks (v4 priced 1024 - tokens off 128-token samples linearly, ~3x too high, then - stalled at 128-token steps: ~70 tok/s at 20% share). A - never-served newcomer gets one prompt step-bounded probe; - afterwards a 2s age override may borrow one such chunk, only - after all shared debt is repaid. In-flight async prefill - blocks the next mixed turn. Prefills are selected by last - completed positive prefill service plus round-robin within each - request-priority tier when scheduling-policy=priority. Decode - token/input budget is allocated first by the base scheduler. - Solo prefill retains the base scheduler's limits. Timing is a - host busy-time proxy. - -Fair knobs (read at runtime; identical on every rank): - GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) - GLM53_FAIR_PREFILL_SHARE default 0.30 (credit accrual fraction) - GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 - GLM53_FAIR_PREFILL_MAX_STEP_MS default 2000 (estimated mixed-step limit) - GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 - -Versioned installer: `# [glm53-decode-floor:v5]`. v1 (no version), v2, v3 -and v4 images are unpatched then re-patched. Fail closed if anchors drift. -""" -from __future__ import annotations - -import inspect -import os -import sys -import time -from pathlib import Path - -P = Path( - os.environ.get( - "GLM53_SCHEDULER_PY", - "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", - ) -) -MARK = "# [glm53-decode-floor]" -MARK_V2 = "# [glm53-decode-floor:v2]" -MARK_V3 = "# [glm53-decode-floor:v3]" -MARK_V4 = "# [glm53-decode-floor:v4]" -MARK_V5 = "# [glm53-decode-floor:v5]" -PRIORITY_MARK = "# [glm53-prefill-priority:v1]" -LEGACY_RANK_PREFIX = """ def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - self.last_service.get(r.request_id, 0.0), -""" -PRIORITY_RANK_PREFIX = """ def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - # [glm53-prefill-priority:v1] Preserve fairness within priority tiers. - int(getattr(r, "priority", 0) or 0) if getattr(self, "priority_mode", False) else 0, - self.last_service.get(r.request_id, 0.0), -""" -LEGACY_CANDIDATES = """ self._candidates = self._rank_prefills(prefills) -""" -PRIORITY_CANDIDATES = """ policy = getattr(sched, "policy", None) - self.priority_mode = getattr(policy, "value", policy) == "priority" - self._candidates = self._rank_prefills(prefills) -""" - -IMPORT_OLD = """import itertools -import time -""" -IMPORT_NEW = """import itertools -import os -import time -""" - -# v1 helper + insertions (recipe f906ee9 / this overlay before v2). -V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" -V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. -V2_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] - # NOTE(woosuk) on the scheduling algorithm: -""" -V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] - pooler_outputs = model_runner_output.pooler_output -""" -V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" -V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] -""" -V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] - if num_new_tokens == 0: - break -""" - - - -# Frozen v3 anchors for migration from the reviewed implementation. -V3_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v3] - # NOTE(woosuk) on the scheduling algorithm: -""" -V3_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v3] - pooler_outputs = model_runner_output.pooler_output -""" -V3_FIN_OLD = """ # Check if the scheduling constraints are satisfied. - total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) -""" -V3_FIN_NEW = """ # Check if the scheduling constraints are satisfied. - _GLM53_MIXED.note_schedule_output(self, num_scheduled_tokens) # [glm53-decode-floor:v3] - total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) -""" -V3_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( - num_new_tokens, - token_budget, - int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), - _GLM53_MIXED.needs_prefill_compute(request), - ) - - # Make sure the input position does not exceed the max model len. -""" -V3_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( - num_new_tokens, - token_budget, - int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), - _GLM53_MIXED.needs_prefill_compute(request), - ) - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if num_new_tokens <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - - # chunked prefill has to be enabled explicitly to allow -""" -V3_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v3] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V3_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] -""" -V3_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] - if num_new_tokens == 0: - break -""" - -class _Glm53MixedPrefill: # [glm53-decode-floor:v5] - """Bound contention using completion feedback, without synchronizing GPUs.""" - - LADDER = (128, 256, 512, 768, 1024, 1536, 2048) - COLD_TOK_S = 1300.0 - FIT_WINDOW = 24 - - def __init__(self, now=None): - self._now = now or time.monotonic - self.mode = "skip" - self.legacy_cap = 0 - self.logged_boot = False - self.hist_every = 50 - self._parse() - self.last_service = {} - self.arrival = {} - self.rr_seq = {} - self.served_tokens = {} - self.rr_n = 0 - self.steps = 0 - self.credit = 0.0 - self.in_contention = False - self.inflight = {} - self.mixed_samples = [] - self.solo_samples = [] - self.last_account_mono = None - self.last_prefill_turn_mono = 0.0 - self.step_tag = None - self._sched_id = None - self._open_rec = None - self.selected = set() - self._candidates = [] - self._tried = set() - self.step_mode = "solo" - self.defer_reason = "none" - self.missed_prefill = 0 - self._model_cache = None - - def _e(self, name, default): - v = os.environ.get(name) - return default if v is None or not str(v).strip() else str(v).strip() - - def _parse(self) -> None: - raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() - self.legacy_cap = 0 - if raw in ("0", "off", "no"): - self.mode = "off" - elif raw in ("skip", "-1"): - self.mode = "skip" - elif raw == "fair": - self.mode = "fair" - else: - try: - cap = int(raw) - except ValueError: - cap = 0 - if cap <= 0: - self.mode = "off" - else: - self.mode = "cap" - self.legacy_cap = cap - try: - self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) - except ValueError: - self.chunk = 256 - if self.chunk <= 0: - self.chunk = 256 - try: - self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) - except ValueError: - self.share = 0.30 - self.share = min(1.0, max(0.0, self.share)) - try: - self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 - except ValueError: - self.interval_s = 2.0 - if self.interval_s <= 0: - self.interval_s = 2.0 - try: - self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) - except ValueError: - self.max_chunks = 1 - self.max_chunks = max(1, min(self.max_chunks, 16)) - try: - self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) - except ValueError: - self.max_step_s = 2.0 - if self.mode == "fair" and not self.logged_boot: - print( - f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " - f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " - f"interval_s={self.interval_s} max_step_s={self.max_step_s} " - f"max_chunks={self.max_chunks}", - flush=True, - ) - self.logged_boot = True - - - @staticmethod - def prefill_remaining(request, computed=None): - prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) - if computed is None: - computed = int(getattr(request, "num_computed_tokens", 0) or 0) - tokens = int(getattr(request, "num_tokens", prompt) or prompt) - return max(0, max(prompt, tokens - 1) - int(computed)) - - def needs_prefill_compute(self, request, computed=None): - return self.prefill_remaining(request, computed) > 0 - - def _iter_waiting(self, sched): - for name in ("waiting", "skipped_waiting"): - q = getattr(sched, name, None) - if not q: - continue - try: - for r in q: - yield r - except TypeError: - continue - - def _live_ids(self, sched): - ids = set() - for r in list(getattr(sched, "running", None) or []): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - for r in self._iter_waiting(sched): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - reqs = getattr(sched, "requests", None) - if isinstance(reqs, dict): - ids.update(reqs.keys()) - return ids - - def _prune(self, live): - for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): - dead = [k for k in store if k not in live] - for k in dead: - store.pop(k, None) - - - @property - def inflight_prefill(self): - return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) - - def _shape(self, decodes, prefills): - history = max((int(r.num_computed_tokens) for r in decodes), default=0) - position = max((int(r.num_computed_tokens) for r in prefills), default=0) - return (len(decodes), (history // 4096).bit_length(), - (position // 4096).bit_length()) - - def _cost_model(self): - """Fit step cost dt = a + b*n over recent prefill-bearing steps. - - Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host - time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to - a solo chunk plus a few decode rows: solo samples are pooled for the - fit. The fit is then scaled so recent mixed samples are not - underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). - Needs two distinct chunk sizes; otherwise returns None and the caller - scales linearly. v4 scaled one sample linearly, priced 1024 tokens off - 128-token samples at ~3x the real cost, and never climbed back. - """ - if self._model_cache is not None: - return self._model_cache - mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] - pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] - if len(pts) < 2 or len({n for n, _ in pts}) < 2: - return None - cnt = float(len(pts)) - sx = float(sum(n for n, _ in pts)) - sy = float(sum(dt for _, dt in pts)) - sxx = float(sum(n * n for n, _ in pts)) - sxy = float(sum(n * dt for n, dt in pts)) - den = cnt * sxx - sx * sx - b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 - a = max(0.0, (sy - b * sx) / cnt) - if mixed: - ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) - r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] - else: - r = 1.1 # decode rows not sampled yet - r = min(1.5, max(1.0, r)) - self._model_cache = (a * r, b * r) - return self._model_cache - - def _est_dt(self, tokens, shape=None): - """Conservative host-time estimate; not a guaranteed execution bound.""" - n = max(1, int(tokens)) - model = self._cost_model() - if model is not None: - a, b = model - return max(0.01, a + b * n) - samples = self.mixed_samples[-8:] - if samples: - # One size only: scale linearly with a fixed-cost floor. - return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) - return max(0.05, n / self.COLD_TOK_S) - - def _target(self, remaining, room): - """Largest rung whose estimated step fits `room` (tokens per accounted - second rise with size under a fixed per-step cost), or None.""" - rungs = {self.chunk} - if self.mixed_samples or self._cost_model() is not None: - rungs.update(self.LADDER) - if remaining is not None: - rungs = {min(n, remaining) for n in rungs} - fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] - fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] - return fitting[-1] if fitting else None - - def _credit_limit(self): - return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) - - def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - # [glm53-prefill-priority:v1] Preserve fairness within priority tiers. - int(getattr(r, "priority", 0) or 0) if getattr(self, "priority_mode", False) else 0, - self.last_service.get(r.request_id, 0.0), - self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) - - def _promote_next(self): - grants = (self._open_rec or {}).get("grants", {}) - for r in self._candidates: - if len(self.selected | set(grants)) >= self.max_chunks: - break - if r.request_id not in self._tried: - self.selected.add(r.request_id) - - def _release(self, rid, reason="allocation"): - rec = self._open_rec - if rec is None: - return - grant = rec["grants"].pop(rid, None) - if grant: - self.credit += grant[1] - if grant[2]: - rec["borrowed"] = False - self.selected.discard(rid) - self._tried.add(rid) - self.defer_reason = reason - self._promote_next() - - def note_scheduled(self, request, num_new_tokens): - # Called after alignment/encoder caps; final allocation is sealed below. - if num_new_tokens <= 0: - self._release(request.request_id, "zero_progress") - - def protect_decode(self, request): - return (self.mode == "fair" and self._open_rec is not None - and self._open_rec["had_decode"] - and self.needs_prefill_compute(request)) - - def begin_step(self, sched): - now = self._now() - sid = id(sched) - tag = (sid, int(sched.current_step)) - if self.step_tag == tag: - return - if self._sched_id is not None and self._sched_id != sid: - self.inflight.clear() - self._open_rec = None - self.credit = 0.0 - self.in_contention = False - self.last_account_mono = None - self._sched_id = sid - # An unfinished/failed schedule has dispatched nothing: refund its grants. - if self._open_rec: - self.credit += sum(g[1] for g in self._open_rec["grants"].values()) - self.step_tag = tag - self.steps += 1 - self._prune(self._live_ids(sched)) - running = list(sched.running) - waiting = list(self._iter_waiting(sched)) - prefills = list({r.request_id: r for r in running + waiting - if self.needs_prefill_compute(r)}.values()) - decodes = [r for r in running if not self.needs_prefill_compute(r)] - # The base loop checks eligibility and reserves BOTH token and input/draft - # capacity by executing these requests first. Preserve order within groups. - if self.mode == "fair": - sched.running[:] = decodes + [r for r in running - if self.needs_prefill_compute(r)] - for r in running + waiting: - self.arrival.setdefault(r.request_id, now) - policy = getattr(sched, "policy", None) - self.priority_mode = getattr(policy, "value", policy) == "priority" - self._candidates = self._rank_prefills(prefills) - self._tried = set() - self.selected = set() - sched._glm53_align_prefill_limit = None - self._open_rec = { - "step_id": int(sched.current_step), "t_submit": now, - "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), - "shape": self._shape(decodes, prefills), "grants": {}, - "borrowed": False, - } - # Do not forgive debt while a decoder or its outstanding work remains. - if not decodes and not any(r["had_decode"] for r in self.inflight.values()): - self.in_contention = False - self.credit = 0.0 - self.step_mode = "legacy" if self.mode != "fair" else "solo" - self.defer_reason = "none" - if self.mode != "fair" or not prefills or not decodes: - return - if not self.in_contention: - self.credit = min(self.max_step_s, self._est_dt(self.chunk)) - self.in_contention = True - if self.inflight_prefill: - self.step_mode = "decode_only" - self.defer_reason = "async_inflight" - else: - self.step_mode = "prefill_turn" - self._promote_next() - self._maybe_log() - - def cap_for(self, sched, request, computed=None): - self.begin_step(sched) - sched._glm53_align_prefill_limit = None - remaining = self.prefill_remaining(request, computed) - if remaining <= 0 or self.mode == "off": - return None - peer_decode = any(r is not request and not self.needs_prefill_compute(r) - for r in sched.running) - if self.mode == "skip": - return 0 if peer_decode else None - if self.mode == "cap": - cap = self.legacy_cap if peer_decode else None - sched._glm53_align_prefill_limit = cap - return cap - if self.step_mode == "solo": - return None - rid = request.request_id - rec = self._open_rec - if rid in rec["grants"]: - return rec["grants"][rid][0] - if self.step_mode != "prefill_turn" or rid not in self.selected: - return 0 - reserved = sum(g[1] for g in rec["grants"].values()) - gap_room = max(0.0, self.max_step_s - reserved) - age = self._now() - self.last_service.get(rid, self.arrival[rid]) - pick = self._target(remaining, gap_room) - if pick is None: - self._release(rid, "gap_budget") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - # Save credit for the most efficient rung that fits the step budget - # instead of spending it on a smaller chunk now (v4 did, and stalled - # at 128-token steps once one expensive sample priced 256 too high). - cap, cost = pick - borrowed = False - if cost > self.credit + 1e-9: - # A never-served newcomer gets a prompt probe; afterwards service - # ages. Either may borrow ONE step-bounded chunk globally, only - # after all shared debt is repaid: queue churn or several aged - # newcomers cannot repeatedly overdraw the decoder. - due = rid not in self.last_service or age >= self.interval_s - if (due and self.credit >= -1e-9 and not rec["grants"] - and not rec["borrowed"]): - borrowed = True - else: - self._release(rid, "credit") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - self.credit -= cost - rec["grants"][rid] = (cap, cost, borrowed) - rec["borrowed"] |= borrowed - self._tried.add(rid) - sched._glm53_align_prefill_limit = cap - return cap - - def finish_step(self, sched, scheduler_output): - """Seal final work before _update_after_schedule advances request state.""" - rec = self._open_rec - self._open_rec = None - if rec is None: - return - scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() - if n > 0} - requests = getattr(sched, "requests", {}) - prefill = {} - for rid, n in scheduled.items(): - request = requests.get(rid) - if request is not None: - amount = min(n, self.prefill_remaining(request)) - if amount > 0: - prefill[rid] = amount - reserved = 0.0 - for rid, (_, estimate, _) in rec["grants"].items(): - # Never publish/charge a tentative grant that alignment, allocation, - # encoder caps or preemption removed from the final scheduler output. - actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 - self.credit += estimate - actual_est - reserved += actual_est - if not scheduled: - # Empty schedules do not necessarily have a completion callback. - return - rec.update(output=scheduler_output, scheduled=scheduled, - prefill_tokens=prefill, reserved=reserved) - del rec["grants"] - if self.mode == "fair": - self.inflight[id(scheduler_output)] = rec - - def observe_output(self, sched, scheduler_output): - rec = self.inflight.pop(id(scheduler_output), None) - if rec is None or rec["output"] is not scheduler_output: - return # Empty, duplicate or unrelated callback; never pop another step. - now = self._now() - # Partition observed busy wall time instead of adding overlapping - # submit-to-completion latencies of queued async steps. Queue residence - # and engine idle gaps therefore cannot mint decode credit twice. - start = rec["t_submit"] - if self.last_account_mono is not None: - start = max(start, self.last_account_mono) - dt = max(0.0, now - start) - self.last_account_mono = max(now, self.last_account_mono or now) - actual = scheduler_output.num_scheduled_tokens - served = {rid: min(n, max(0, int(actual.get(rid, 0)))) - for rid, n in rec["prefill_tokens"].items() - if int(actual.get(rid, 0)) > 0} - for rid, n in served.items(): - self.last_service[rid] = now - self.rr_n += 1 - self.rr_seq[rid] = self.rr_n - self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n - n = sum(served.values()) - if rec["had_decode"]: - # Reservation was already debited; settle once, including overruns. - self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) - self.credit = min(self.credit, self._credit_limit()) - if n and dt > 0: - self.mixed_samples.append((rec["shape"], n, dt)) - self.mixed_samples = self.mixed_samples[-64:] - self.last_prefill_turn_mono = now - self._model_cache = None - elif n and dt > 0: - self.solo_samples.append((n, dt)) - self.solo_samples = self.solo_samples[-32:] - self._model_cache = None - if n and self.hist_every > 0: - print(f"[glm53-decode-floor] completed_step={rec['step_id']} " - f"contention={int(rec['had_decode'])} prefill_tokens={served} " - f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " - f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) - self._prune(self._live_ids(sched)) - - def _maybe_log(self): - if self.hist_every <= 0 or self.steps % self.hist_every != 1: - return - rec = self._open_rec or {} - remaining = sum(self.prefill_remaining(r) for r in self._candidates) - target, estimate = (self._target(None, self.max_step_s) - or (self.chunk, self._est_dt(self.chunk))) - rate = target / estimate if estimate > 0 else 0.0 - eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") - model = self._cost_model() - fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" - if model else "fit=none") - print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " - f"mode={self.step_mode} defer={self.defer_reason} " - f"inflight={self.inflight_prefill} credit={self.credit:.3f} " - f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " - f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " - f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) - - @staticmethod - def aligned_new_tokens( - start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None - ) -> int: - """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" - if policy_cap is not None and policy_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, policy_cap) - end = start + num_new - if end < prefill_end: - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end - return max(0, end - start) - - - -def _helper_text() -> str: - body = inspect.getsource(_Glm53MixedPrefill) - return ( - "\n" - + body - + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5]\n\n" - + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5]\n" - + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" - ) - - -HELPER = None # filled at apply time so tests can call _helper_text() - - -BEGIN_OLD = """ self.current_step += 1 - # NOTE(woosuk) on the scheduling algorithm: -""" -BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] - # NOTE(woosuk) on the scheduling algorithm: -""" - -OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - pooler_outputs = model_runner_output.pooler_output -""" -OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] - pooler_outputs = model_runner_output.pooler_output -""" - -RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - - # Make sure the input position does not exceed the max model len. -""" -RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" - -WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - - # chunked prefill has to be enabled explicitly to allow -""" -WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" - -RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) -""" -RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] -""" - -WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - if num_new_tokens == 0: - break -""" -WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] - if num_new_tokens == 0: - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" -FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] - with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" - -RUNNING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled because one of the following -""" -RUNNING_ZERO_NEW = """ if num_new_tokens == 0: - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - # The request cannot be scheduled because one of the following -""" - -WAITING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled. - break -""" -WAITING_ZERO_NEW = """ if num_new_tokens == 0: - # The request cannot be scheduled. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. - # Preempt the lowest-priority request. -""" -PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. - if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] - break - # Preempt the lowest-priority request. -""" - -RUNNING_ALLOC_OLD = """ if new_blocks is None: - # Cannot schedule this request. - break -""" -RUNNING_ALLOC_NEW = """ if new_blocks is None: - # Cannot schedule this request. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - req_index += 1 - continue - break -""" - -WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - break - - # KVTransfer:""" -WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break - - # KVTransfer:""" - -V4_PAIRS = ( - (BEGIN_NEW, BEGIN_OLD, 'begin'), - (OBS_NEW, OBS_OLD, 'obs'), - (RUNNING_NEW, RUNNING_OLD, 'running'), - (WAITING_NEW, WAITING_OLD, 'waiting'), - (ALIGN_NEW, ALIGN_OLD, 'align'), - (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), - (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), - (FIN_NEW, FIN_OLD, 'fin'), - (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), - (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), - (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), - (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), - (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), -) - - -def replace_once(text: str, old: str, new: str, label: str) -> str: - n = text.count(old) - if n != 1: - raise SystemExit(f"{P}: expected one {label} target, found {n}") - return text.replace(old, new, 1) - - -def _strip_helper(text: str, label: str, *, expected: str | None = None) -> str: - start = text.find("class _Glm53MixedPrefill:") - if start < 0: - raise SystemExit(f"{P}: {label} helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: {label} helper end not found") - end = min(candidates) - if expected is not None and text[start:end].strip() != expected.strip(): - raise SystemExit(f"{P}: {label} helper drifted") - return text[:start] + text[end:] - - -def _strip_v1_helper(text: str) -> str: - start = text.find(V1_HELPER_START) - if start < 0: - raise SystemExit(f"{P}: v1 helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: v1 helper end not found") - end = min(candidates) - return text[:start] + "\n" + text[end:] - - -def unpatch_v1(text: str) -> str: - if V1_RUNNING_NEW in text: - text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 running insertion drifted") - if V1_WAITING_NEW in text: - text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 waiting insertion drifted") - if V1_HELPER_START in text: - text = _strip_v1_helper(text) - leftover = [ - "_glm53_mixed_prefill_policy(self.running, request)", - V1_HELPER_START, - ] - for s in leftover: - if s in text: - raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") - return text - - -def unpatch_v2(text: str) -> str: - pairs = ( - (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), - (V2_OBS_NEW, OBS_OLD, "v2-obs"), - (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), - (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), - (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), - (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), - (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), - ) - for new, old, label in pairs: - if new in text: - text = replace_once(text, new, old, label) - elif MARK_V2 in new: - # Some insertions may already have been removed; fail if marker remains. - pass - if "class _Glm53MixedPrefill:" in text: - text = _strip_helper(text, "v2") - if MARK_V2 in text: - raise SystemExit(f"{P}: v2 leftover after unpatch") - return text - - -def unpatch_v3(text: str) -> str: - pairs = ( - (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), - (V3_OBS_NEW, OBS_OLD, "v3-obs"), - (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), - (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), - (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), - (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), - (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), - (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), - ) - for new, old, label in pairs: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v3") - if MARK_V3 in text: - raise SystemExit(f"{P}: v3 leftover after unpatch") - return text - - -def unpatch_v4(text: str) -> str: - for new, old, label in V4_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v4") - if MARK_V4 in text: - raise SystemExit(f"{P}: v4 leftover after unpatch") - return text - - -# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 -# insertions above stay frozen so a v4 image can be unpatched exactly. -V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) - - -def unpatch_v5(text: str) -> str: - for new, old, label in V5_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v5", expected=_helper_text()) - if MARK_V5 in text: - raise SystemExit(f"{P}: v5 leftover after unpatch") - return text - - -def apply_v5(text: str) -> str: - if "import os\n" not in text.split("import time\n", 1)[0]: - text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") - needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" - text = replace_once(text, needle, _helper_text() + needle, "helper") - for new, old, label in V5_PAIRS: - text = replace_once(text, old, new, label) - compile(text, str(P), "exec") - return text - - -def main() -> int: - if not P.is_file(): - raise SystemExit(f"missing {P}") - text = P.read_text() - original = text - if MARK_V5 in text: - if PRIORITY_MARK not in text: - text = replace_once(text, LEGACY_RANK_PREFIX, PRIORITY_RANK_PREFIX, "prefill priority ranking") - text = replace_once(text, LEGACY_CANDIDATES, PRIORITY_CANDIDATES, "prefill scheduler policy") - # Validate existing anchors/helper instead of trusting the marker alone. - # unpatch_v5 checks every v5 insertion occurs exactly once, strips the - # helper, and rejects leftover markers. The result is discarded: a - # later overlay may legitimately sit between the helper and the - # cuda_graph import anchor, so re-applying at that fixed anchor would - # relocate the helper and fail a byte-compare on a healthy file. - unpatch_v5(text) - # The import edit is part of the applied state; the byte-compare used - # to cover it implicitly. - if "import os\n" not in text.split("import time\n", 1)[0]: - raise SystemExit(f"{P}: v5 import drifted") - compile(text, str(P), "exec") - if text != original: - P.write_text(text) - print(f"{P.name}: migrated prefill-priority ranking ({PRIORITY_MARK})") - else: - print(f"{P.name}: {MARK_V5} already present — verified") - return 0 - if MARK_V4 in text: - text = unpatch_v4(text) - elif MARK_V3 in text: - text = unpatch_v3(text) - elif MARK_V2 in text: - text = unpatch_v2(text) - elif MARK in text or V1_HELPER_START in text: - text = unpatch_v1(text) - text = apply_v5(text) - if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text: - raise SystemExit(f"{P}: older marker left after migration") - if text != original: - P.write_text(text) - print(f"patched {P.name} ({MARK_V5})") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/fixtures/decode_floor_producers/v6_carry_0870fab9.py b/tests/fixtures/decode_floor_producers/v6_carry_0870fab9.py deleted file mode 100644 index 0c16d294..00000000 --- a/tests/fixtures/decode_floor_producers/v6_carry_0870fab9.py +++ /dev/null @@ -1,1183 +0,0 @@ -#!/usr/bin/env python3 -"""Mixed-prefill policy: skip / cap / off / fair (issue #6). - -A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover -MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill -chunk. Mixed execution leaves the uniform decode FULL-graph path, and a -128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond -budget. - -GLM53_MIXED_PREFILL_CHUNK: - skip / -1 — do not mix prefill with decode. Starves every - waiting/running prefill while any peer is decoding; there is - no age limit. Solo prefill is unchanged. - N>0 — cap mixed prefill chunks to N tokens while a peer decodes. - The cap is fed into hybrid Mamba alignment so N < block_size - still makes sub-block progress. 128 still stalls ~10 tok/s. - 0 / off — disable the extra isolation policy. - fair — service-time mixing (default on TP=2/3/4 since 2026-09-15, - v6). Decode-only - steps between prefill turns; at most - GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). - Only prefill that contends with a decoder is charged (solo - prefill is cost-sampled, not debt). Credit accrues at SHARE - of accounted engine time. v5 fits a fixed-plus-per-token step - cost from solo and mixed samples, targets the largest ladder - rung (128..2048) whose estimated step fits - GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that - rung instead of spending it on small chunks (v4 priced 1024 - tokens off 128-token samples linearly, ~3x too high, then - stalled at 128-token steps: ~70 tok/s at 20% share). A - never-served newcomer gets one prompt step-bounded probe; - afterwards a 2s age override may borrow one such chunk, only - after all shared debt is repaid. In-flight async prefill - blocks the next mixed turn. Prefills are selected by last - completed positive prefill service plus round-robin. Fair - ordering is allocation-aware: a runnable candidate passed by - the scheduler is carried one step when a later candidate - fails KV allocation, preventing blocked-head reset starvation. - Decode token/input budget is allocated first by the base scheduler. - Solo prefill retains the base scheduler's limits. Timing is a - host busy-time proxy. - -Fair knobs (read at runtime; identical on every rank): - GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) - GLM53_FAIR_PREFILL_SHARE default 0.30 (credit accrual fraction) - GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 - GLM53_FAIR_PREFILL_MAX_STEP_MS default 2000 (estimated mixed-step limit) - GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 - -Versioned installer: `# [glm53-decode-floor:v6]`. v1 (no version), v2, v3, -v4 and v5 images are unpatched then re-patched. Fail closed if anchors drift. -""" -from __future__ import annotations - -import inspect -import hashlib -import os -import sys -import time -from pathlib import Path - -P = Path( - os.environ.get( - "GLM53_SCHEDULER_PY", - "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", - ) -) -MARK = "# [glm53-decode-floor]" -MARK_V2 = "# [glm53-decode-floor:v2]" -MARK_V3 = "# [glm53-decode-floor:v3]" -MARK_V4 = "# [glm53-decode-floor:v4]" -MARK_V5 = "# [glm53-decode-floor:v5]" -MARK_V6 = "# [glm53-decode-floor:v6]" - -IMPORT_OLD = """import itertools -import time -""" -IMPORT_NEW = """import itertools -import os -import time -""" - -# v1 helper + insertions (recipe f906ee9 / this overlay before v2). -V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" -V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. -V2_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] - # NOTE(woosuk) on the scheduling algorithm: -""" -V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] - pooler_outputs = model_runner_output.pooler_output -""" -V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" -V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] -""" -V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] - if num_new_tokens == 0: - break -""" - - - -# Frozen v3 anchors for migration from the reviewed implementation. -V3_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v3] - # NOTE(woosuk) on the scheduling algorithm: -""" -V3_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v3] - pooler_outputs = model_runner_output.pooler_output -""" -V3_FIN_OLD = """ # Check if the scheduling constraints are satisfied. - total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) -""" -V3_FIN_NEW = """ # Check if the scheduling constraints are satisfied. - _GLM53_MIXED.note_schedule_output(self, num_scheduled_tokens) # [glm53-decode-floor:v3] - total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) -""" -V3_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( - num_new_tokens, - token_budget, - int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), - _GLM53_MIXED.needs_prefill_compute(request), - ) - - # Make sure the input position does not exceed the max model len. -""" -V3_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v3] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - num_new_tokens = _GLM53_MIXED.clip_for_decode_reserve( - num_new_tokens, - token_budget, - int(getattr(self, "_glm53_decode_reserve_tokens", 0) or 0), - _GLM53_MIXED.needs_prefill_compute(request), - ) - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if num_new_tokens <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - - # chunked prefill has to be enabled explicitly to allow -""" -V3_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v3] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V3_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] -""" -V3_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v3] - if num_new_tokens == 0: - break -""" - -class _Glm53MixedPrefill: # [glm53-decode-floor:v6] - """Bound contention using completion feedback, without synchronizing GPUs.""" - - LADDER = (128, 256, 512, 768, 1024, 1536, 2048) - COLD_TOK_S = 1300.0 - FIT_WINDOW = 24 - - def __init__(self, now=None): - self._now = now or time.monotonic - self.mode = "skip" - self.legacy_cap = 0 - self.logged_boot = False - self.hist_every = 50 - self._parse() - self.last_service = {} - self.arrival = {} - self.rr_seq = {} - self.served_tokens = {} - self.rr_n = 0 - self.steps = 0 - self.credit = 0.0 - self.in_contention = False - self.inflight = {} - self.mixed_samples = [] - self.solo_samples = [] - self.last_account_mono = None - self.last_prefill_turn_mono = 0.0 - self.step_tag = None - self._sched_id = None - self._open_rec = None - self.selected = set() - self._candidates = [] - self._tried = set() - self._passed = set() - self._carry = set() - self.step_mode = "solo" - self.defer_reason = "none" - self.missed_prefill = 0 - self._model_cache = None - - def _e(self, name, default): - v = os.environ.get(name) - return default if v is None or not str(v).strip() else str(v).strip() - - def _parse(self) -> None: - raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() - self.legacy_cap = 0 - if raw in ("0", "off", "no"): - self.mode = "off" - elif raw in ("skip", "-1"): - self.mode = "skip" - elif raw == "fair": - self.mode = "fair" - else: - try: - cap = int(raw) - except ValueError: - cap = 0 - if cap <= 0: - self.mode = "off" - else: - self.mode = "cap" - self.legacy_cap = cap - try: - self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) - except ValueError: - self.chunk = 256 - if self.chunk <= 0: - self.chunk = 256 - try: - self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) - except ValueError: - self.share = 0.30 - self.share = min(1.0, max(0.0, self.share)) - try: - self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 - except ValueError: - self.interval_s = 2.0 - if self.interval_s <= 0: - self.interval_s = 2.0 - try: - self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) - except ValueError: - self.max_chunks = 1 - self.max_chunks = max(1, min(self.max_chunks, 16)) - try: - self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) - except ValueError: - self.max_step_s = 2.0 - if self.mode == "fair" and not self.logged_boot: - print( - f"[glm53-decode-floor] fair v6 probe_chunk={self.chunk} " - f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " - f"interval_s={self.interval_s} max_step_s={self.max_step_s} " - f"max_chunks={self.max_chunks}", - flush=True, - ) - self.logged_boot = True - - - @staticmethod - def prefill_remaining(request, computed=None): - prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) - if computed is None: - computed = int(getattr(request, "num_computed_tokens", 0) or 0) - tokens = int(getattr(request, "num_tokens", prompt) or prompt) - return max(0, max(prompt, tokens - 1) - int(computed)) - - def needs_prefill_compute(self, request, computed=None): - return self.prefill_remaining(request, computed) > 0 - - def _iter_waiting(self, sched): - for name in ("waiting", "skipped_waiting"): - q = getattr(sched, name, None) - if not q: - continue - try: - for r in q: - yield r - except TypeError: - continue - - def _live_ids(self, sched): - ids = set() - for r in list(getattr(sched, "running", None) or []): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - for r in self._iter_waiting(sched): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - reqs = getattr(sched, "requests", None) - if isinstance(reqs, dict): - ids.update(reqs.keys()) - return ids - - def _prune(self, live): - for store in (self.last_service, self.arrival, self.rr_seq, - self.served_tokens): - dead = [k for k in store if k not in live] - for k in dead: - store.pop(k, None) - - - @property - def inflight_prefill(self): - return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) - - def _shape(self, decodes, prefills): - history = max((int(r.num_computed_tokens) for r in decodes), default=0) - position = max((int(r.num_computed_tokens) for r in prefills), default=0) - return (len(decodes), (history // 4096).bit_length(), - (position // 4096).bit_length()) - - def _cost_model(self): - """Fit step cost dt = a + b*n over recent prefill-bearing steps. - - Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host - time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to - a solo chunk plus a few decode rows: solo samples are pooled for the - fit. The fit is then scaled so recent mixed samples are not - underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). - Needs two distinct chunk sizes; otherwise returns None and the caller - scales linearly. v4 scaled one sample linearly, priced 1024 tokens off - 128-token samples at ~3x the real cost, and never climbed back. - """ - if self._model_cache is not None: - return self._model_cache - mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] - pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] - if len(pts) < 2 or len({n for n, _ in pts}) < 2: - return None - cnt = float(len(pts)) - sx = float(sum(n for n, _ in pts)) - sy = float(sum(dt for _, dt in pts)) - sxx = float(sum(n * n for n, _ in pts)) - sxy = float(sum(n * dt for n, dt in pts)) - den = cnt * sxx - sx * sx - b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 - a = max(0.0, (sy - b * sx) / cnt) - if mixed: - ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) - r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] - else: - r = 1.1 # decode rows not sampled yet - r = min(1.5, max(1.0, r)) - self._model_cache = (a * r, b * r) - return self._model_cache - - def _est_dt(self, tokens, shape=None): - """Conservative host-time estimate; not a guaranteed execution bound.""" - n = max(1, int(tokens)) - model = self._cost_model() - if model is not None: - a, b = model - return max(0.01, a + b * n) - samples = self.mixed_samples[-8:] - if samples: - # One size only: scale linearly with a fixed-cost floor. - return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) - return max(0.05, n / self.COLD_TOK_S) - - def _target(self, remaining, room): - """Largest rung whose estimated step fits `room` (tokens per accounted - second rise with size under a fixed per-step cost), or None.""" - rungs = {self.chunk} - if self.mixed_samples or self._cost_model() is not None: - rungs.update(self.LADDER) - if remaining is not None: - rungs = {min(n, remaining) for n in rungs} - fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] - fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] - return fitting[-1] if fitting else None - - def _credit_limit(self): - return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) - - def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - self.last_service.get(r.request_id, 0.0), - self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) - - def _promote_next(self, prefer_passed=False): - grants = (self._open_rec or {}).get("grants", {}) - candidates = self._candidates - if prefer_passed: - candidates = ([r for r in candidates if r.request_id in self._passed] - + [r for r in candidates if r.request_id not in self._passed]) - for r in candidates: - if len(self.selected | set(grants)) >= self.max_chunks: - break - if r.request_id not in self._tried: - self.selected.add(r.request_id) - if prefer_passed and r.request_id in self._passed: - # Allocation failed after this request's scheduler position - # had already passed. Carry it into the next step instead of - # resetting to the same blocked head of the fair queue. - self._carry.add(r.request_id) - - def _release(self, rid, reason="allocation"): - rec = self._open_rec - if rec is None: - return - # The running loop reports zero even when policy returned cap=0 for an - # unselected request. That request was passed, not attempted; keep it - # eligible for allocation-aware carry if a later selected request fails. - if rid not in self.selected and rid not in rec["grants"]: - return - grant = rec["grants"].pop(rid, None) - if grant: - self.credit += grant[1] - if grant[2]: - rec["borrowed"] = False - self.selected.discard(rid) - self._tried.add(rid) - if reason == "zero_progress": - # A selected request that could not make progress rotates behind - # never-served peers. Credit/gap deferrals are not admissions and - # retain their fair rank. - self.rr_n += 1 - self.rr_seq[rid] = self.rr_n - self.defer_reason = reason - self._promote_next(prefer_passed=True) - - def note_scheduled(self, request, num_new_tokens): - # Called after alignment/encoder caps; final allocation is sealed below. - if num_new_tokens <= 0: - self._release(request.request_id, "zero_progress") - - def protect_decode(self, request): - return (self.mode == "fair" and self._open_rec is not None - and self._open_rec["had_decode"] - and self.needs_prefill_compute(request)) - - def begin_step(self, sched): - now = self._now() - sid = id(sched) - tag = (sid, int(sched.current_step)) - if self.step_tag == tag: - return - if self._sched_id is not None and self._sched_id != sid: - self.inflight.clear() - self._open_rec = None - self._carry.clear() - self.credit = 0.0 - self.in_contention = False - self.last_account_mono = None - self._sched_id = sid - # An unfinished/failed schedule has dispatched nothing: refund its grants. - if self._open_rec: - self.credit += sum(g[1] for g in self._open_rec["grants"].values()) - self.step_tag = tag - self.steps += 1 - self._prune(self._live_ids(sched)) - running = list(sched.running) - waiting = list(self._iter_waiting(sched)) - prefills = list({r.request_id: r for r in running + waiting - if self.needs_prefill_compute(r)}.values()) - decodes = [r for r in running if not self.needs_prefill_compute(r)] - # The base loop checks eligibility and reserves BOTH token and input/draft - # capacity by executing these requests first. Preserve order within groups. - if self.mode == "fair": - sched.running[:] = decodes + [r for r in running - if self.needs_prefill_compute(r)] - for r in running + waiting: - self.arrival.setdefault(r.request_id, now) - self._candidates = self._rank_prefills(prefills) - self._tried = set() - self.selected = set() - self._passed = set() - live_prefills = {r.request_id for r in prefills} - self._carry.intersection_update(live_prefills) - sched._glm53_align_prefill_limit = None - self._open_rec = { - "step_id": int(sched.current_step), "t_submit": now, - "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), - "shape": self._shape(decodes, prefills), "grants": {}, - "borrowed": False, - } - # Do not forgive debt while a decoder or its outstanding work remains. - if not decodes and not any(r["had_decode"] for r in self.inflight.values()): - self.in_contention = False - self.credit = 0.0 - self._carry.clear() - self.step_mode = "legacy" if self.mode != "fair" else "solo" - self.defer_reason = "none" - if self.mode != "fair" or not prefills or not decodes: - return - if not self.in_contention: - self.credit = min(self.max_step_s, self._est_dt(self.chunk)) - self.in_contention = True - if self.inflight_prefill: - self.step_mode = "decode_only" - self.defer_reason = "async_inflight" - else: - self.step_mode = "prefill_turn" - for r in self._candidates: - if r.request_id in self._carry: - self.selected.add(r.request_id) - if len(self.selected) >= self.max_chunks: - break - self._promote_next() - self._maybe_log() - - def cap_for(self, sched, request, computed=None): - self.begin_step(sched) - sched._glm53_align_prefill_limit = None - remaining = self.prefill_remaining(request, computed) - if remaining <= 0 or self.mode == "off": - return None - peer_decode = any(r is not request and not self.needs_prefill_compute(r) - for r in sched.running) - if self.mode == "skip": - return 0 if peer_decode else None - if self.mode == "cap": - cap = self.legacy_cap if peer_decode else None - sched._glm53_align_prefill_limit = cap - return cap - if self.step_mode == "solo": - return None - rid = request.request_id - rec = self._open_rec - if rid in rec["grants"]: - return rec["grants"][rid][0] - self._passed.add(rid) - if self.step_mode != "prefill_turn" or rid not in self.selected: - return 0 - self._carry.discard(rid) - reserved = sum(g[1] for g in rec["grants"].values()) - gap_room = max(0.0, self.max_step_s - reserved) - age = self._now() - self.last_service.get(rid, self.arrival[rid]) - pick = self._target(remaining, gap_room) - if pick is None: - self._release(rid, "gap_budget") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - # Save credit for the most efficient rung that fits the step budget - # instead of spending it on a smaller chunk now (v4 did, and stalled - # at 128-token steps once one expensive sample priced 256 too high). - cap, cost = pick - borrowed = False - if cost > self.credit + 1e-9: - # A never-served newcomer gets a prompt probe; afterwards service - # ages. Either may borrow ONE step-bounded chunk globally, only - # after all shared debt is repaid: queue churn or several aged - # newcomers cannot repeatedly overdraw the decoder. - due = rid not in self.last_service or age >= self.interval_s - if (due and self.credit >= -1e-9 and not rec["grants"] - and not rec["borrowed"]): - borrowed = True - else: - self._release(rid, "credit") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - self.credit -= cost - rec["grants"][rid] = (cap, cost, borrowed) - rec["borrowed"] |= borrowed - self._tried.add(rid) - sched._glm53_align_prefill_limit = cap - return cap - - def finish_step(self, sched, scheduler_output): - """Seal final work before _update_after_schedule advances request state.""" - rec = self._open_rec - self._open_rec = None - if rec is None: - return - scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() - if n > 0} - requests = getattr(sched, "requests", {}) - prefill = {} - for rid, n in scheduled.items(): - request = requests.get(rid) - if request is not None: - amount = min(n, self.prefill_remaining(request)) - if amount > 0: - prefill[rid] = amount - reserved = 0.0 - for rid, (_, estimate, _) in rec["grants"].items(): - # Never publish/charge a tentative grant that alignment, allocation, - # encoder caps or preemption removed from the final scheduler output. - actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 - self.credit += estimate - actual_est - reserved += actual_est - if not scheduled: - # Empty schedules do not necessarily have a completion callback. - return - rec.update(output=scheduler_output, scheduled=scheduled, - prefill_tokens=prefill, reserved=reserved) - del rec["grants"] - if self.mode == "fair": - self.inflight[id(scheduler_output)] = rec - - def observe_output(self, sched, scheduler_output): - rec = self.inflight.pop(id(scheduler_output), None) - if rec is None or rec["output"] is not scheduler_output: - return # Empty, duplicate or unrelated callback; never pop another step. - now = self._now() - # Partition observed busy wall time instead of adding overlapping - # submit-to-completion latencies of queued async steps. Queue residence - # and engine idle gaps therefore cannot mint decode credit twice. - start = rec["t_submit"] - if self.last_account_mono is not None: - start = max(start, self.last_account_mono) - dt = max(0.0, now - start) - self.last_account_mono = max(now, self.last_account_mono or now) - actual = scheduler_output.num_scheduled_tokens - served = {rid: min(n, max(0, int(actual.get(rid, 0)))) - for rid, n in rec["prefill_tokens"].items() - if int(actual.get(rid, 0)) > 0} - for rid, n in served.items(): - self.last_service[rid] = now - self.rr_n += 1 - self.rr_seq[rid] = self.rr_n - self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n - n = sum(served.values()) - if rec["had_decode"]: - # Reservation was already debited; settle once, including overruns. - self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) - self.credit = min(self.credit, self._credit_limit()) - if n and dt > 0: - self.mixed_samples.append((rec["shape"], n, dt)) - self.mixed_samples = self.mixed_samples[-64:] - self.last_prefill_turn_mono = now - self._model_cache = None - elif n and dt > 0: - self.solo_samples.append((n, dt)) - self.solo_samples = self.solo_samples[-32:] - self._model_cache = None - if n and self.hist_every > 0: - print(f"[glm53-decode-floor] completed_step={rec['step_id']} " - f"contention={int(rec['had_decode'])} prefill_tokens={served} " - f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " - f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) - self._prune(self._live_ids(sched)) - - def _maybe_log(self): - if self.hist_every <= 0 or self.steps % self.hist_every != 1: - return - rec = self._open_rec or {} - remaining = sum(self.prefill_remaining(r) for r in self._candidates) - target, estimate = (self._target(None, self.max_step_s) - or (self.chunk, self._est_dt(self.chunk))) - rate = target / estimate if estimate > 0 else 0.0 - eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") - model = self._cost_model() - fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" - if model else "fit=none") - print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " - f"mode={self.step_mode} defer={self.defer_reason} " - f"inflight={self.inflight_prefill} credit={self.credit:.3f} " - f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " - f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " - f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) - - @staticmethod - def aligned_new_tokens( - start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None - ) -> int: - """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" - if policy_cap is not None and policy_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, policy_cap) - end = start + num_new - if end < prefill_end: - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end - return max(0, end - start) - - - -def _helper_text() -> str: - body = inspect.getsource(_Glm53MixedPrefill) - return ( - "\n" - + body - + "\n_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v6]\n\n" - + "def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v6]\n" - + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" - ) - - -HELPER = None # filled at apply time so tests can call _helper_text() - - -BEGIN_OLD = """ self.current_step += 1 - # NOTE(woosuk) on the scheduling algorithm: -""" -BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] - # NOTE(woosuk) on the scheduling algorithm: -""" - -OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - pooler_outputs = model_runner_output.pooler_output -""" -OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] - pooler_outputs = model_runner_output.pooler_output -""" - -RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - - # Make sure the input position does not exceed the max model len. -""" -RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" - -WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - - # chunked prefill has to be enabled explicitly to allow -""" -WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" - -RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) -""" -RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] -""" - -WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - if num_new_tokens == 0: - break -""" -WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] - if num_new_tokens == 0: - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" -FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] - with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" - -RUNNING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled because one of the following -""" -RUNNING_ZERO_NEW = """ if num_new_tokens == 0: - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - # The request cannot be scheduled because one of the following -""" - -WAITING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled. - break -""" -WAITING_ZERO_NEW = """ if num_new_tokens == 0: - # The request cannot be scheduled. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. - # Preempt the lowest-priority request. -""" -PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. - if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] - break - # Preempt the lowest-priority request. -""" - -RUNNING_ALLOC_OLD = """ if new_blocks is None: - # Cannot schedule this request. - break -""" -RUNNING_ALLOC_NEW = """ if new_blocks is None: - # Cannot schedule this request. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - req_index += 1 - continue - break -""" - -WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - break - - # KVTransfer:""" -WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break - - # KVTransfer:""" - -V4_PAIRS = ( - (BEGIN_NEW, BEGIN_OLD, 'begin'), - (OBS_NEW, OBS_OLD, 'obs'), - (RUNNING_NEW, RUNNING_OLD, 'running'), - (WAITING_NEW, WAITING_OLD, 'waiting'), - (ALIGN_NEW, ALIGN_OLD, 'align'), - (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), - (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), - (FIN_NEW, FIN_OLD, 'fin'), - (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), - (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), - (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), - (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), - (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), -) - - -def replace_once(text: str, old: str, new: str, label: str) -> str: - n = text.count(old) - if n != 1: - raise SystemExit(f"{P}: expected one {label} target, found {n}") - return text.replace(old, new, 1) - - -def _strip_helper(text: str, label: str, *, expected: str | None = None, - expected_sha256: str | None = None) -> str: - start = text.find("class _Glm53MixedPrefill:") - if start < 0: - raise SystemExit(f"{P}: {label} helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: {label} helper end not found") - end = min(candidates) - actual = text[start:end].strip() - if expected is not None and actual != expected.strip(): - raise SystemExit(f"{P}: {label} helper drifted") - if (expected_sha256 is not None - and hashlib.sha256(actual.encode()).hexdigest() != expected_sha256): - raise SystemExit(f"{P}: {label} helper drifted") - return text[:start] + text[end:] - - -def _strip_v1_helper(text: str) -> str: - start = text.find(V1_HELPER_START) - if start < 0: - raise SystemExit(f"{P}: v1 helper start not found") - if start > 0 and text[start - 1] == "\n": - start -= 1 - end_ak = text.find("class _Glm53AdaptiveK:", start) - end_cg = text.find("from vllm.compilation.cuda_graph import CUDAGraphStat\n", start) - candidates = [i for i in (end_ak, end_cg) if i > start] - if not candidates: - raise SystemExit(f"{P}: v1 helper end not found") - end = min(candidates) - return text[:start] + "\n" + text[end:] - - -def unpatch_v1(text: str) -> str: - if V1_RUNNING_NEW in text: - text = replace_once(text, V1_RUNNING_NEW, RUNNING_OLD, "v1-running") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 running insertion drifted") - if V1_WAITING_NEW in text: - text = replace_once(text, V1_WAITING_NEW, WAITING_OLD, "v1-waiting") - elif "_glm53_mixed_prefill_policy(self.running, request)" in text: - raise SystemExit(f"{P}: v1 waiting insertion drifted") - if V1_HELPER_START in text: - text = _strip_v1_helper(text) - leftover = [ - "_glm53_mixed_prefill_policy(self.running, request)", - V1_HELPER_START, - ] - for s in leftover: - if s in text: - raise SystemExit(f"{P}: v1 leftover after unpatch: {s}") - return text - - -def unpatch_v2(text: str) -> str: - pairs = ( - (V2_BEGIN_NEW, BEGIN_OLD, "v2-begin"), - (V2_OBS_NEW, OBS_OLD, "v2-obs"), - (V2_RUNNING_NEW, RUNNING_OLD, "v2-running"), - (V2_WAITING_NEW, WAITING_OLD, "v2-waiting"), - (V2_ALIGN_NEW, ALIGN_OLD, "v2-align"), - (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v2-running-mamba"), - (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v2-waiting-mamba"), - ) - for new, old, label in pairs: - if new in text: - text = replace_once(text, new, old, label) - elif MARK_V2 in new: - # Some insertions may already have been removed; fail if marker remains. - pass - if "class _Glm53MixedPrefill:" in text: - text = _strip_helper(text, "v2") - if MARK_V2 in text: - raise SystemExit(f"{P}: v2 leftover after unpatch") - return text - - -def unpatch_v3(text: str) -> str: - pairs = ( - (V3_BEGIN_NEW, BEGIN_OLD, "v3-begin"), - (V3_OBS_NEW, OBS_OLD, "v3-obs"), - (V3_FIN_NEW, V3_FIN_OLD, "v3-fin"), - (V3_RUNNING_NEW, RUNNING_OLD, "v3-running"), - (V3_WAITING_NEW, WAITING_OLD, "v3-waiting"), - (V3_ALIGN_NEW, ALIGN_OLD, "v3-align"), - (V3_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, "v3-running-mamba"), - (V3_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, "v3-waiting-mamba"), - ) - for new, old, label in pairs: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v3") - if MARK_V3 in text: - raise SystemExit(f"{P}: v3 leftover after unpatch") - return text - - -def unpatch_v4(text: str) -> str: - for new, old, label in V4_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v4") - if MARK_V4 in text: - raise SystemExit(f"{P}: v4 leftover after unpatch") - return text - - -# v5 uses the same scheduler anchors as v4 with the marker advanced; the v4 -# insertions above stay frozen so a v4 image can be unpatched exactly. -V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) -V5_HELPER_SHA256 = "d8009ae8ecdc9e731a8f2109f0d89a8f27140c83808674da54c0f6e30788cd5c" - - -def unpatch_v5(text: str) -> str: - for new, old, label in V5_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v5", expected_sha256=V5_HELPER_SHA256) - if MARK_V5 in text: - raise SystemExit(f"{P}: v5 leftover after unpatch") - return text - - -V6_PAIRS = tuple((new.replace(MARK_V4, MARK_V6), old, label) for new, old, label in V4_PAIRS) - - -def unpatch_v6(text: str) -> str: - for new, old, label in V6_PAIRS: - text = replace_once(text, new, old, label) - text = _strip_helper(text, "v6", expected=_helper_text()) - if MARK_V6 in text: - raise SystemExit(f"{P}: v6 leftover after unpatch") - return text - - -def apply_v6(text: str) -> str: - if "import os\n" not in text.split("import time\n", 1)[0]: - text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") - needle = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" - text = replace_once(text, needle, _helper_text() + needle, "helper") - for new, old, label in V6_PAIRS: - text = replace_once(text, old, new, label) - compile(text, str(P), "exec") - return text - - -def main() -> int: - if not P.is_file(): - raise SystemExit(f"missing {P}") - text = P.read_text() - original = text - if MARK_V6 in text: - # Validate existing anchors/helper instead of trusting the marker alone. - # unpatch_v6 checks every v6 insertion occurs exactly once, strips the - # helper, and rejects leftover markers. The result is discarded: a - # later overlay may legitimately sit between the helper and the - # cuda_graph import anchor, so re-applying at that fixed anchor would - # relocate the helper and fail a byte-compare on a healthy file. - unpatch_v6(text) - # The import edit is part of the applied state; the byte-compare used - # to cover it implicitly. - if "import os\n" not in text.split("import time\n", 1)[0]: - raise SystemExit(f"{P}: v6 import drifted") - compile(text, str(P), "exec") - print(f"{P.name}: {MARK_V6} already present - verified") - return 0 - if MARK_V5 in text: - text = unpatch_v5(text) - elif MARK_V4 in text: - text = unpatch_v4(text) - elif MARK_V3 in text: - text = unpatch_v3(text) - elif MARK_V2 in text: - text = unpatch_v2(text) - elif MARK in text or V1_HELPER_START in text: - text = unpatch_v1(text) - text = apply_v6(text) - if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text or MARK_V5 in text: - raise SystemExit(f"{P}: older marker left after migration") - if text != original: - P.write_text(text) - print(f"patched {P.name} ({MARK_V6})") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/fixtures/decode_floor_producers/v6_warm_deadline_1f74558d.py b/tests/fixtures/decode_floor_producers/v6_warm_deadline_1f74558d.py deleted file mode 100644 index 9b915c00..00000000 --- a/tests/fixtures/decode_floor_producers/v6_warm_deadline_1f74558d.py +++ /dev/null @@ -1,1143 +0,0 @@ -#!/usr/bin/env python3 -"""Mixed-prefill policy: skip / cap / off / fair (issue #6). - -A decode lane on this backend needs ~8 tokens (1 + DFlash2 k=7). The leftover -MNBT budget otherwise goes to a peer FLASHINFER_MLA_SPARSE_SM120 prefill -chunk. Mixed execution leaves the uniform decode FULL-graph path, and a -128-token cap is still ~10 tok/s at 80k KV. Token caps are not a millisecond -budget. - -GLM53_MIXED_PREFILL_CHUNK: - skip / -1 — do not mix prefill with decode. Starves every - waiting/running prefill while any peer is decoding; there is - no age limit. Solo prefill is unchanged. - N>0 — cap mixed prefill chunks to N tokens while a peer decodes. - The cap is fed into hybrid Mamba alignment so N < block_size - still makes sub-block progress. 128 still stalls ~10 tok/s. - 0 / off — disable the extra isolation policy. - fair — service-time mixing (TP=2 default since 2026-09-15, v5; - opt-in on TP=3/TP=4). Decode-only - steps between prefill turns; at most - GLM53_FAIR_PREFILL_MAX_CHUNKS chunks per turn (default 1). - Only prefill that contends with a decoder is charged (solo - prefill is cost-sampled, not debt). Credit accrues at SHARE - of accounted engine time. v5 fits a fixed-plus-per-token step - cost from solo and mixed samples, targets the largest ladder - rung (128..2048) whose estimated step fits - GLM53_FAIR_PREFILL_MAX_STEP_MS, and saves credit for that - rung instead of spending it on small chunks (v4 priced 1024 - tokens off 128-token samples linearly, ~3x too high, then - stalled at 128-token steps: ~70 tok/s at 20% share). A - never-served newcomer gets one prompt step-bounded probe; - afterwards a 2s age override may borrow one such chunk, only - after all shared debt is repaid. In-flight async prefill - blocks the next mixed turn. Prefills are selected by last - completed positive prefill service plus round-robin. Decode - token/input budget is allocated first by the base scheduler. - Solo prefill retains the base scheduler's limits. Timing is a - host busy-time proxy. - -Fair knobs (read at runtime; identical on every rank): - GLM53_FAIR_PREFILL_CHUNK default 256 (probe size until timing samples exist) - GLM53_FAIR_PREFILL_SHARE default 0.20 (credit accrual fraction) - GLM53_FAIR_PREFILL_MAX_INTERVAL_MS default 2000 - GLM53_FAIR_PREFILL_MAX_STEP_MS default 1000 (estimated mixed-step limit) - GLM53_FAIR_PREFILL_MAX_CHUNKS default 1 - -Versioned installer: `# [glm53-decode-floor:v6]`. v1 (no version), v2 and v5 -images are unpatched then re-patched. A v3 or v4 marker is refused without -touching the source: no authenticated producer of those intermediate helper -bodies was recovered from public history, so a canonical v3/v4 image cannot be -established, and no body is invented to fill the gap. Fail closed if anchors -drift. - -Fail-closed migration: the legacy helper site is validated *before* anything is -removed -- v1/v2/v5 against the published helper text (sha256), v6 against this -installer's own text. The frozen gate sites are inverted, one exact byte range -is removed, and the whole file is round-trip checked: re-adding the same span -and re-applying that version's frozen sites must reproduce the input -byte-for-byte. A drifted, duplicated, decorated, marker-only, v3/v4-marked or -otherwise unattested site is refused with no write. - -v6 (opt-in gate, both features OFF by default -- v5 behaviour is preserved): - GLM53_MIXED_PREFILL_WARM_TOKENS >0 admits a request whose uncached remainder - is <= this many tokens (typically one hybrid block, 3584) without waiting - for the decoders. 0 (default) disables the bypass. - GLM53_MIXED_PREFILL_MAX_WAIT_MS >0 releases a request the `skip` hold has - starved for this long under GLM53_MIXED_PREFILL_LATE_CAP (default 512) - tokens per step. 0 (default) waits forever, i.e. v1..v5 behaviour. The - release only makes the request eligible for a LATE_CAP step; it bounds - neither allocation nor compute service (capacity and the base scheduler's - own limits still decide), and the request then crawls like cap:N. - An explicit `cap` keeps its cap (only the warm bypass applies); `fair` is a - different mechanism and is not touched by either knob. -""" -from __future__ import annotations - -import hashlib -import inspect -import os -import sys -import time -from pathlib import Path - -P = Path( - os.environ.get( - "GLM53_SCHEDULER_PY", - "/usr/local/lib/python3.12/dist-packages/vllm/v1/core/sched/scheduler.py", - ) -) -MARK = "# [glm53-decode-floor]" -MARK_V2 = "# [glm53-decode-floor:v2]" -MARK_V3 = "# [glm53-decode-floor:v3]" -MARK_V4 = "# [glm53-decode-floor:v4]" -MARK_V5 = "# [glm53-decode-floor:v5]" -MARK_V6 = "# [glm53-decode-floor:v6]" - -IMPORT_OLD = """import itertools -import time -""" -IMPORT_NEW = """import itertools -import os -import time -""" - -# v1 helper + insertions (recipe f906ee9 / this overlay before v2). -V1_HELPER_START = "def _glm53_mixed_prefill_policy(running, current):" -V1_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and request.num_computed_tokens < request.num_prompt_tokens: - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V1_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self.running, request) # [glm53-decode-floor] - if mixed_cap is not None and num_computed_tokens < request.num_prompt_tokens: - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -# Frozen v2 insertions — used only to unpatch an already-v2 scheduler. -V2_BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v2] - # NOTE(woosuk) on the scheduling algorithm: -""" -V2_OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v2] - pooler_outputs = model_runner_output.pooler_output -""" -V2_RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" -V2_WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v2] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" -V2_ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v2] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -V2_RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] -""" -V2_WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v2] - if num_new_tokens == 0: - break -""" - - -class _Glm53MixedPrefill: # [glm53-decode-floor:v5] - """Bound contention using completion feedback, without synchronizing GPUs.""" - - LADDER = (128, 256, 512, 768, 1024, 1536, 2048) - COLD_TOK_S = 1300.0 - FIT_WINDOW = 24 - - def __init__(self, now=None): - self._now = now or time.monotonic - self.mode = "skip" - self.legacy_cap = 0 - self.logged_boot = False - self.logged_gate = False - self.hist_every = 50 - self._parse() - self.last_service = {} - self.arrival = {} - self.rr_seq = {} - self.served_tokens = {} - self.rr_n = 0 - self.steps = 0 - self.credit = 0.0 - self.in_contention = False - self.inflight = {} - self.mixed_samples = [] - self.solo_samples = [] - self.last_account_mono = None - self.last_prefill_turn_mono = 0.0 - self.step_tag = None - self._sched_id = None - self._open_rec = None - self.selected = set() - self._candidates = [] - self._tried = set() - self.step_mode = "solo" - self.defer_reason = "none" - self.missed_prefill = 0 - self._model_cache = None - - def _e(self, name, default): - v = os.environ.get(name) - return default if v is None or not str(v).strip() else str(v).strip() - - def _parse(self) -> None: - raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() - self.legacy_cap = 0 - if raw in ("0", "off", "no"): - self.mode = "off" - elif raw in ("skip", "-1"): - self.mode = "skip" - elif raw == "fair": - self.mode = "fair" - else: - try: - cap = int(raw) - except ValueError: - cap = 0 - if cap <= 0: - self.mode = "off" - else: - self.mode = "cap" - self.legacy_cap = cap - try: - self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) - except ValueError: - self.chunk = 256 - if self.chunk <= 0: - self.chunk = 256 - try: - self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.20")) - except ValueError: - self.share = 0.20 - self.share = min(1.0, max(0.0, self.share)) - try: - self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 - except ValueError: - self.interval_s = 2.0 - if self.interval_s <= 0: - self.interval_s = 2.0 - try: - self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) - except ValueError: - self.max_chunks = 1 - self.max_chunks = max(1, min(self.max_chunks, 16)) - try: - self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "1000")) / 1000.0) - except ValueError: - self.max_step_s = 1.0 - # v6 gate (warm bypass + deadline). Both features are OFF unless an - # operator sets the knobs: 0 disables each one and is the default, so a - # skip/cap deployment keeps exactly the v5 hold behaviour. - try: - self.warm_tokens = int(self._e("GLM53_MIXED_PREFILL_WARM_TOKENS", "0")) - except ValueError: - self.warm_tokens = 0 - if not 0 <= self.warm_tokens <= 1_000_000: - self.warm_tokens = 0 - try: - self.max_wait_ms = int(self._e("GLM53_MIXED_PREFILL_MAX_WAIT_MS", "0")) - except ValueError: - self.max_wait_ms = 0 - if not 0 <= self.max_wait_ms <= 600_000: - self.max_wait_ms = 0 - try: - self.late_cap = int(self._e("GLM53_MIXED_PREFILL_LATE_CAP", "512")) - except ValueError: - self.late_cap = 512 - if not 64 <= self.late_cap <= 8192: - self.late_cap = 512 - if self.mode == "fair" and not self.logged_boot: - print( - f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " - f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " - f"interval_s={self.interval_s} max_step_s={self.max_step_s} " - f"max_chunks={self.max_chunks}", - flush=True, - ) - self.logged_boot = True - if (self.warm_tokens or self.max_wait_ms) and not self.logged_gate: - print( - f"[glm53-decode-floor] gate v6 warm_tokens={self.warm_tokens} " - f"max_wait_ms={self.max_wait_ms} late_cap={self.late_cap}", - flush=True, - ) - self.logged_gate = True - - - @staticmethod - def prefill_remaining(request, computed=None): - prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) - if computed is None: - computed = int(getattr(request, "num_computed_tokens", 0) or 0) - tokens = int(getattr(request, "num_tokens", prompt) or prompt) - return max(0, max(prompt, tokens - 1) - int(computed)) - - def needs_prefill_compute(self, request, computed=None): - return self.prefill_remaining(request, computed) > 0 - - def _iter_waiting(self, sched): - for name in ("waiting", "skipped_waiting"): - q = getattr(sched, name, None) - if not q: - continue - try: - for r in q: - yield r - except TypeError: - continue - - def _live_ids(self, sched): - ids = set() - for r in list(getattr(sched, "running", None) or []): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - for r in self._iter_waiting(sched): - rid = getattr(r, "request_id", None) - if rid is not None: - ids.add(rid) - reqs = getattr(sched, "requests", None) - if isinstance(reqs, dict): - ids.update(reqs.keys()) - return ids - - def _prune(self, live): - for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): - dead = [k for k in store if k not in live] - for k in dead: - store.pop(k, None) - - - @property - def inflight_prefill(self): - return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) - - def _shape(self, decodes, prefills): - history = max((int(r.num_computed_tokens) for r in decodes), default=0) - position = max((int(r.num_computed_tokens) for r in prefills), default=0) - return (len(decodes), (history // 4096).bit_length(), - (position // 4096).bit_length()) - - def _cost_model(self): - """Fit step cost dt = a + b*n over recent prefill-bearing steps. - - Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host - time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to - a solo chunk plus a few decode rows: solo samples are pooled for the - fit. The fit is then scaled so recent mixed samples are not - underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). - Needs two distinct chunk sizes; otherwise returns None and the caller - scales linearly. v4 scaled one sample linearly, priced 1024 tokens off - 128-token samples at ~3x the real cost, and never climbed back. - """ - if self._model_cache is not None: - return self._model_cache - mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] - pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] - if len(pts) < 2 or len({n for n, _ in pts}) < 2: - return None - cnt = float(len(pts)) - sx = float(sum(n for n, _ in pts)) - sy = float(sum(dt for _, dt in pts)) - sxx = float(sum(n * n for n, _ in pts)) - sxy = float(sum(n * dt for n, dt in pts)) - den = cnt * sxx - sx * sx - b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 - a = max(0.0, (sy - b * sx) / cnt) - if mixed: - ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) - r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] - else: - r = 1.1 # decode rows not sampled yet - r = min(1.5, max(1.0, r)) - self._model_cache = (a * r, b * r) - return self._model_cache - - def _est_dt(self, tokens, shape=None): - """Conservative host-time estimate; not a guaranteed execution bound.""" - n = max(1, int(tokens)) - model = self._cost_model() - if model is not None: - a, b = model - return max(0.01, a + b * n) - samples = self.mixed_samples[-8:] - if samples: - # One size only: scale linearly with a fixed-cost floor. - return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) - return max(0.05, n / self.COLD_TOK_S) - - def _target(self, remaining, room): - """Largest rung whose estimated step fits `room` (tokens per accounted - second rise with size under a fixed per-step cost), or None.""" - rungs = {self.chunk} - if self.mixed_samples or self._cost_model() is not None: - rungs.update(self.LADDER) - if remaining is not None: - rungs = {min(n, remaining) for n in rungs} - fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] - fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] - return fitting[-1] if fitting else None - - def _credit_limit(self): - return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) - - def _rank_prefills(self, prefills): - return sorted(prefills, key=lambda r: ( - self.last_service.get(r.request_id, 0.0), - self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) - - def _promote_next(self): - grants = (self._open_rec or {}).get("grants", {}) - for r in self._candidates: - if len(self.selected | set(grants)) >= self.max_chunks: - break - if r.request_id not in self._tried: - self.selected.add(r.request_id) - - def _release(self, rid, reason="allocation"): - rec = self._open_rec - if rec is None: - return - grant = rec["grants"].pop(rid, None) - if grant: - self.credit += grant[1] - if grant[2]: - rec["borrowed"] = False - self.selected.discard(rid) - self._tried.add(rid) - self.defer_reason = reason - self._promote_next() - - def note_scheduled(self, request, num_new_tokens): - # Called after alignment/encoder caps; final allocation is sealed below. - if num_new_tokens <= 0: - self._release(request.request_id, "zero_progress") - - def protect_decode(self, request): - return (self.mode == "fair" and self._open_rec is not None - and self._open_rec["had_decode"] - and self.needs_prefill_compute(request)) - - def begin_step(self, sched): - now = self._now() - sid = id(sched) - tag = (sid, int(sched.current_step)) - if self.step_tag == tag: - return - if self._sched_id is not None and self._sched_id != sid: - self.inflight.clear() - self._open_rec = None - self.credit = 0.0 - self.in_contention = False - self.last_account_mono = None - self._sched_id = sid - # An unfinished/failed schedule has dispatched nothing: refund its grants. - if self._open_rec: - self.credit += sum(g[1] for g in self._open_rec["grants"].values()) - self.step_tag = tag - self.steps += 1 - self._prune(self._live_ids(sched)) - running = list(sched.running) - waiting = list(self._iter_waiting(sched)) - prefills = list({r.request_id: r for r in running + waiting - if self.needs_prefill_compute(r)}.values()) - decodes = [r for r in running if not self.needs_prefill_compute(r)] - # The base loop checks eligibility and reserves BOTH token and input/draft - # capacity by executing these requests first. Preserve order within groups. - if self.mode == "fair": - sched.running[:] = decodes + [r for r in running - if self.needs_prefill_compute(r)] - for r in running + waiting: - self.arrival.setdefault(r.request_id, now) - self._candidates = self._rank_prefills(prefills) - self._tried = set() - self.selected = set() - sched._glm53_align_prefill_limit = None - self._open_rec = { - "step_id": int(sched.current_step), "t_submit": now, - "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), - "shape": self._shape(decodes, prefills), "grants": {}, - "borrowed": False, - } - # Do not forgive debt while a decoder or its outstanding work remains. - if not decodes and not any(r["had_decode"] for r in self.inflight.values()): - self.in_contention = False - self.credit = 0.0 - self.step_mode = "legacy" if self.mode != "fair" else "solo" - self.defer_reason = "none" - if self.mode != "fair" or not prefills or not decodes: - return - if not self.in_contention: - self.credit = min(self.max_step_s, self._est_dt(self.chunk)) - self.in_contention = True - if self.inflight_prefill: - self.step_mode = "decode_only" - self.defer_reason = "async_inflight" - else: - self.step_mode = "prefill_turn" - self._promote_next() - self._maybe_log() - - def _warm_bypass(self, remaining): - """Admit a small uncached remainder at once (v6). Off while the knob is 0.""" - return self.warm_tokens > 0 and remaining <= self.warm_tokens - - def _deadline_cap(self, request, remaining): - """Late admission for a request the skip hold would starve (v6). - - Returns 0 while the deadline is disabled (0 ms = wait forever, v5 - behaviour) or not reached, else LATE_CAP tokens for this step. The - first-seen stamp is the per-request arrival the fair policy already - keeps, so chunking, preemption and requeue cannot reset the wait. - """ - if self.max_wait_ms <= 0: - return 0 - rid = request.request_id - first_seen = self.arrival.get(rid) - if first_seen is None: - first_seen = self._now() - self.arrival[rid] = first_seen - if (self._now() - first_seen) * 1000.0 >= self.max_wait_ms: - return max(1, min(self.late_cap, remaining)) - return 0 - - def cap_for(self, sched, request, computed=None): - self.begin_step(sched) - sched._glm53_align_prefill_limit = None - remaining = self.prefill_remaining(request, computed) - if remaining <= 0 or self.mode == "off": - return None - peer_decode = any(r is not request and not self.needs_prefill_compute(r) - for r in sched.running) - if self.mode == "skip": - if not peer_decode: - return None - if self._warm_bypass(remaining): - return None - return self._deadline_cap(request, remaining) - if self.mode == "cap": - cap = self.legacy_cap if peer_decode else None - if cap is not None and self._warm_bypass(remaining): - cap = None - sched._glm53_align_prefill_limit = cap - return cap - if self.step_mode == "solo": - return None - rid = request.request_id - rec = self._open_rec - if rid in rec["grants"]: - return rec["grants"][rid][0] - if self.step_mode != "prefill_turn" or rid not in self.selected: - return 0 - reserved = sum(g[1] for g in rec["grants"].values()) - gap_room = max(0.0, self.max_step_s - reserved) - age = self._now() - self.last_service.get(rid, self.arrival[rid]) - pick = self._target(remaining, gap_room) - if pick is None: - self._release(rid, "gap_budget") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - # Save credit for the most efficient rung that fits the step budget - # instead of spending it on a smaller chunk now (v4 did, and stalled - # at 128-token steps once one expensive sample priced 256 too high). - cap, cost = pick - borrowed = False - if cost > self.credit + 1e-9: - # A never-served newcomer gets a prompt probe; afterwards service - # ages. Either may borrow ONE step-bounded chunk globally, only - # after all shared debt is repaid: queue churn or several aged - # newcomers cannot repeatedly overdraw the decoder. - due = rid not in self.last_service or age >= self.interval_s - if (due and self.credit >= -1e-9 and not rec["grants"] - and not rec["borrowed"]): - borrowed = True - else: - self._release(rid, "credit") - if age >= self.interval_s: - self.missed_prefill += 1 - self._maybe_log() - return 0 - self.credit -= cost - rec["grants"][rid] = (cap, cost, borrowed) - rec["borrowed"] |= borrowed - self._tried.add(rid) - sched._glm53_align_prefill_limit = cap - return cap - - def finish_step(self, sched, scheduler_output): - """Seal final work before _update_after_schedule advances request state.""" - rec = self._open_rec - self._open_rec = None - if rec is None: - return - scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() - if n > 0} - requests = getattr(sched, "requests", {}) - prefill = {} - for rid, n in scheduled.items(): - request = requests.get(rid) - if request is not None: - amount = min(n, self.prefill_remaining(request)) - if amount > 0: - prefill[rid] = amount - reserved = 0.0 - for rid, (_, estimate, _) in rec["grants"].items(): - # Never publish/charge a tentative grant that alignment, allocation, - # encoder caps or preemption removed from the final scheduler output. - actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 - self.credit += estimate - actual_est - reserved += actual_est - if not scheduled: - # Empty schedules do not necessarily have a completion callback. - return - rec.update(output=scheduler_output, scheduled=scheduled, - prefill_tokens=prefill, reserved=reserved) - del rec["grants"] - if self.mode == "fair": - self.inflight[id(scheduler_output)] = rec - - def observe_output(self, sched, scheduler_output): - rec = self.inflight.pop(id(scheduler_output), None) - if rec is None or rec["output"] is not scheduler_output: - return # Empty, duplicate or unrelated callback; never pop another step. - now = self._now() - # Partition observed busy wall time instead of adding overlapping - # submit-to-completion latencies of queued async steps. Queue residence - # and engine idle gaps therefore cannot mint decode credit twice. - start = rec["t_submit"] - if self.last_account_mono is not None: - start = max(start, self.last_account_mono) - dt = max(0.0, now - start) - self.last_account_mono = max(now, self.last_account_mono or now) - actual = scheduler_output.num_scheduled_tokens - served = {rid: min(n, max(0, int(actual.get(rid, 0)))) - for rid, n in rec["prefill_tokens"].items() - if int(actual.get(rid, 0)) > 0} - for rid, n in served.items(): - self.last_service[rid] = now - self.rr_n += 1 - self.rr_seq[rid] = self.rr_n - self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n - n = sum(served.values()) - if rec["had_decode"]: - # Reservation was already debited; settle once, including overruns. - self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) - self.credit = min(self.credit, self._credit_limit()) - if n and dt > 0: - self.mixed_samples.append((rec["shape"], n, dt)) - self.mixed_samples = self.mixed_samples[-64:] - self.last_prefill_turn_mono = now - self._model_cache = None - elif n and dt > 0: - self.solo_samples.append((n, dt)) - self.solo_samples = self.solo_samples[-32:] - self._model_cache = None - if n and self.hist_every > 0: - print(f"[glm53-decode-floor] completed_step={rec['step_id']} " - f"contention={int(rec['had_decode'])} prefill_tokens={served} " - f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " - f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) - self._prune(self._live_ids(sched)) - - def _maybe_log(self): - if self.hist_every <= 0 or self.steps % self.hist_every != 1: - return - rec = self._open_rec or {} - remaining = sum(self.prefill_remaining(r) for r in self._candidates) - target, estimate = (self._target(None, self.max_step_s) - or (self.chunk, self._est_dt(self.chunk))) - rate = target / estimate if estimate > 0 else 0.0 - eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") - model = self._cost_model() - fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" - if model else "fit=none") - print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " - f"mode={self.step_mode} defer={self.defer_reason} " - f"inflight={self.inflight_prefill} credit={self.credit:.3f} " - f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " - f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " - f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) - - @staticmethod - def aligned_new_tokens( - start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None - ) -> int: - """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" - if policy_cap is not None and policy_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, policy_cap) - end = start + num_new - if end < prefill_end: - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end - return max(0, end - start) - - - -BEGIN_OLD = """ self.current_step += 1 - # NOTE(woosuk) on the scheduling algorithm: -""" -BEGIN_NEW = """ self.current_step += 1 - _GLM53_MIXED.begin_step(self) # [glm53-decode-floor:v4] - # NOTE(woosuk) on the scheduling algorithm: -""" - -OBS_OLD = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - pooler_outputs = model_runner_output.pooler_output -""" -OBS_NEW = """ num_scheduled_tokens = scheduler_output.num_scheduled_tokens - _GLM53_MIXED.observe_output(self, scheduler_output) # [glm53-decode-floor:v4] - pooler_outputs = model_runner_output.pooler_output -""" - -RUNNING_OLD = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - - # Make sure the input position does not exceed the max model len. -""" -RUNNING_NEW = """ if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: - num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min( - num_new_tokens, token_budget, input_budget - draft_slots - ) - mixed_cap = _glm53_mixed_prefill_policy(self, request) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request): - num_new_tokens = min(num_new_tokens, mixed_cap) - - # Make sure the input position does not exceed the max model len. -""" - -WAITING_OLD = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - - # chunked prefill has to be enabled explicitly to allow -""" -WAITING_NEW = """ threshold = self.scheduler_config.long_prefill_token_threshold - if 0 < threshold < num_new_tokens: - num_new_tokens = threshold - mixed_cap = _glm53_mixed_prefill_policy(self, request, num_computed_tokens) # [glm53-decode-floor:v4] - if mixed_cap is not None and _GLM53_MIXED.needs_prefill_compute(request, num_computed_tokens): - if mixed_cap <= 0: - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - num_new_tokens = min(num_new_tokens, mixed_cap) - - # chunked prefill has to be enabled explicitly to allow -""" - -ALIGN_OLD = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" -ALIGN_NEW = """ max_prefill_tokens = self.max_num_scheduled_tokens - long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold - if long_prefill_threshold > 0: - max_prefill_tokens = min(max_prefill_tokens, long_prefill_threshold) - _align_cap = getattr(self, "_glm53_align_prefill_limit", None) # [glm53-decode-floor:v4] - if _align_cap is not None and _align_cap > 0: - max_prefill_tokens = min(max_prefill_tokens, _align_cap) - aligned_end = end // block_size * block_size - if aligned_end > start or block_size <= max_prefill_tokens: - end = aligned_end -""" - -RUNNING_MAMBA_OLD = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) -""" -RUNNING_MAMBA_NEW = """ # Apply Mamba alignment before encoder caps. - if self.need_mamba_block_aligned_split: - num_new_tokens = self._mamba_block_aligned_split( - request, num_new_tokens - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] -""" - -WAITING_MAMBA_OLD = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - if num_new_tokens == 0: - break -""" -WAITING_MAMBA_NEW = """ num_new_tokens = self._mamba_block_aligned_split( - request, - num_new_tokens, - num_new_local_computed_tokens, - num_external_computed_tokens, - ) - _GLM53_MIXED.note_scheduled(request, num_new_tokens) # [glm53-decode-floor:v4] - if num_new_tokens == 0: - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -FIN_OLD = """ with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" -FIN_NEW = """ _GLM53_MIXED.finish_step(self, scheduler_output) # [glm53-decode-floor:v4] - with record_function_or_nullcontext("schedule: update_after_schedule"): - self._update_after_schedule(scheduler_output) -""" - -RUNNING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled because one of the following -""" -RUNNING_ZERO_NEW = """ if num_new_tokens == 0: - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - # The request cannot be scheduled because one of the following -""" - -WAITING_ZERO_OLD = """ if num_new_tokens == 0: - # The request cannot be scheduled. - break -""" -WAITING_ZERO_NEW = """ if num_new_tokens == 0: - # The request cannot be scheduled. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.mode == "fair": - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break -""" - -PREFILL_PREEMPT_OLD = """ # The request cannot be scheduled. - # Preempt the lowest-priority request. -""" -PREFILL_PREEMPT_NEW = """ # The request cannot be scheduled. - if _GLM53_MIXED.protect_decode(request): # [glm53-decode-floor:v4] - break - # Preempt the lowest-priority request. -""" - -RUNNING_ALLOC_OLD = """ if new_blocks is None: - # Cannot schedule this request. - break -""" -RUNNING_ALLOC_NEW = """ if new_blocks is None: - # Cannot schedule this request. - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - req_index += 1 - continue - break -""" - -WAITING_ALLOC_OLD = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - break - - # KVTransfer:""" -WAITING_ALLOC_NEW = """ if request.has_encoder_inputs: - self.encoder_cache_manager.free(request) - _GLM53_MIXED.note_scheduled(request, 0) # [glm53-decode-floor:v4] - if _GLM53_MIXED.protect_decode(request): - request_queue.pop_request() - step_skipped_waiting.prepend_request(request) - continue - break - - # KVTransfer:""" - - -def replace_once(text: str, old: str, new: str, label: str) -> str: - n = text.count(old) - if n != 1: - raise SystemExit(f"{P}: expected one {label} target, found {n}") - return text.replace(old, new, 1) - - -V2_PAIRS = ( - (V2_BEGIN_NEW, BEGIN_OLD, 'v2-begin'), - (V2_OBS_NEW, OBS_OLD, 'v2-obs'), - (V2_RUNNING_NEW, RUNNING_OLD, 'v2-running'), - (V2_WAITING_NEW, WAITING_OLD, 'v2-waiting'), - (V2_ALIGN_NEW, ALIGN_OLD, 'v2-align'), - (V2_RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'v2-running-mamba'), - (V2_WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'v2-waiting-mamba'), -) - -V4_PAIRS = ( - (BEGIN_NEW, BEGIN_OLD, 'begin'), - (OBS_NEW, OBS_OLD, 'obs'), - (RUNNING_NEW, RUNNING_OLD, 'running'), - (WAITING_NEW, WAITING_OLD, 'waiting'), - (ALIGN_NEW, ALIGN_OLD, 'align'), - (RUNNING_MAMBA_NEW, RUNNING_MAMBA_OLD, 'running_mamba'), - (WAITING_MAMBA_NEW, WAITING_MAMBA_OLD, 'waiting_mamba'), - (FIN_NEW, FIN_OLD, 'fin'), - (RUNNING_ZERO_NEW, RUNNING_ZERO_OLD, 'running_zero'), - (WAITING_ZERO_NEW, WAITING_ZERO_OLD, 'waiting_zero'), - (PREFILL_PREEMPT_NEW, PREFILL_PREEMPT_OLD, 'prefill_preempt'), - (RUNNING_ALLOC_NEW, RUNNING_ALLOC_OLD, 'running_alloc'), - (WAITING_ALLOC_NEW, WAITING_ALLOC_OLD, 'waiting_alloc'), -) - -V1_PAIRS = ( - (V1_RUNNING_NEW, RUNNING_OLD, 'v1-running'), - (V1_WAITING_NEW, WAITING_OLD, 'v1-waiting'), -) - -# v5 and v6 use the same scheduler anchors as v4 with the marker advanced; the -# v4 insertions above stay frozen so a v4 image can be unpatched exactly. v6 -# only changes the helper (opt-in warm bypass + deadline, both off by default), -# so an unpatch/re-patch of a v6 image is byte-identical. -V5_PAIRS = tuple((new.replace(MARK_V4, MARK_V5), old, label) for new, old, label in V4_PAIRS) -V6_PAIRS = tuple((new.replace(MARK_V5, MARK_V6), old, label) for new, old, label in V5_PAIRS) - -LEGACY_MARK = {1: MARK, 2: MARK_V2, 3: MARK_V3, 4: MARK_V4, 5: MARK_V5, 6: MARK_V6} -LEGACY_PAIRS = {1: V1_PAIRS, 2: V2_PAIRS, 5: V5_PAIRS, 6: V6_PAIRS} # v3/v4 are refused, not migrated - -NEEDLE = "from vllm.compilation.cuda_graph import CUDAGraphStat\n" -ADAPTIVE_K_HEAD = "class _Glm53AdaptiveK:" -CLASS_HEAD = "class _Glm53MixedPrefill:" - -# --------------------------------------------------------------------------- -# Canonical legacy helper registry. -# -# Every migratable version has exactly one accepted helper site, and the site is -# validated *before* anything is removed. The versions whose helper text is -# published are compared byte-for-byte (sha256 and length). v3 and v4 were -# introduced by 180725a5ce33 as markers for intermediate builds whose helper -# bodies were never recovered from public history: they are refused as -# unsupported rather than migrated on a structural guess, because an invented -# body would not be a canonical image. A site that validates is removed as one -# exact byte range, so text between the helper and its anchors can never be -# dropped by accident. -# -# provenance, public history of MiaAI-Lab/GLM-5.3-Flash-EXL3-2x-DGX-Sparks: -# v1 f3043c95bbf9 overlay/patch_scheduler_decode_floor.py:HELPER -# cross-attested byte-identical by cad980f7263:HELPER_V1, which documents -# itself as "byte-identical to the shipped v1 overlay" -# v2 9a23b2d8f061 overlay/patch_scheduler_decode_floor.py:HELPER -# v5 a9bbbadc4c72 overlay/patch_scheduler_decode_floor.py:_helper_text() -# v3/v4 no public revision carries a body (180725a5ce33 adds only the -# markers); refused, not migrated (see above) -# -# The published variants that share the unversioned marker but bake a different -# policy default (the "0" default of 83f0067/d9758a6/9a557cf/14b6a9f) are NOT -# accepted: migrating them would silently land a held default change, so they -# are refused like any other drifted site. -# --------------------------------------------------------------------------- -LEGACY_HELPER_SHA256 = { - 1: "c0c10c6385bd7e75bfc0f724378960b9d9cb943cf18c386a5a780cf03cb4c48d", - 2: "e170f2bf6d0c0fe0493a91ff54bd719036f6c4b5ea9fcf2ee1346968a4e3fbed", - 5: "c05769276df6da0a24b3b8c21e251f39ee7aa2a74ee92c6d6ce3d85c38151a53", -} -LEGACY_HELPER_LEN = {1: 784, 2: 13378, 5: 20866} - - -def _refuse(reason: str): - raise SystemExit(f"{P}: refusing to rewrite the scheduler: {reason}") - - -def _site_reasons(span: str, version: int, exact=None) -> str: - """Return '' when `span` is exactly this installer's canonical v{version} helper site. - - v1/v2/v5 are pinned to their published helper text (sha256 and length); v6 is - compared against this installer's own `_helper_text()`. A version with no - published text has no accepted site at all. - """ - if exact is not None: - return "" if span == exact else f"helper site is not this installer's v{version} helper text" - expected = LEGACY_HELPER_SHA256.get(version) - if expected is None: - return f"v{version} helper text was never published; there is no accepted site to compare" - got = hashlib.sha256(span.encode("utf-8")).hexdigest() - if got != expected or len(span) != LEGACY_HELPER_LEN[version]: - return (f"helper site is not the published v{version} helper (sha256 {got[:16]}, " - f"{len(span)} bytes; expected {expected[:16]}, {LEGACY_HELPER_LEN[version]} bytes)") - return "" - - -def _legacy_span(text: str, version: int, exact=None): - """Validate and return (start, end, span) for the unique v{version} helper site. - - Returns None when no helper definition is present at all. Refuses (no write) - on a missing, duplicated, decorated, drifted or otherwise unattested site. - """ - head_token = V1_HELPER_START if version == 1 else CLASS_HEAD - found = text.count(head_token) - if found == 0: - return None - if found != 1: - _refuse(f"{found} {head_token!r} definitions found; v{version} installs exactly one") - head = text.find(head_token) - if head < 2 or text[head - 1] != "\n" or text[head - 2] != "\n": - _refuse(f"v{version} helper is not preceded by the installer's blank-line boundary " - f"(indented or decorated site)") - start = head - 1 - reasons = [] - for anchor in (NEEDLE, ADAPTIVE_K_HEAD): - end = text.find(anchor, head) - if end < 0: - continue - span = text[start:end] - why = _site_reasons(span, version, exact) - if not why: - return start, end, span - reasons.append(why) - _refuse(f"v{version} helper site rejected: " + "; ".join(reasons)) - - -def _unpatch(text: str, version: int, exact=None): - """Invert version `version`: its frozen gate sites, then its validated helper site.""" - if version not in LEGACY_PAIRS: - _refuse(f"v{version} is not a migratable version: its helper text was never published") - for new, old, label in LEGACY_PAIRS[version]: - if new not in text: - _refuse(f"{label} insertion is missing or drifted") - text = replace_once(text, new, old, label) - site = _legacy_span(text, version, exact) - if site is None: - _refuse(f"v{version} helper site is missing") - start, end, span = site - clean = text[:start] + text[end:] - mark = LEGACY_MARK[version] - leftover = [s for s in (mark, V1_HELPER_START if version == 1 else CLASS_HEAD) if s in clean] - if leftover: - _refuse(f"v{version} marker or helper definition still present after unpatch") - return clean, start, span - - -def _round_trip(original: str, clean: str, version: int, start: int, span: str) -> None: - """Refuse unless re-applying version `version` to `clean` reproduces `original`.""" - again = clean[:start] + span + clean[start:] - for new, old, label in LEGACY_PAIRS[version]: - again = replace_once(again, old, new, label) - if again != original: - _refuse(f"v{version} round trip is not byte-identical; the image is not a canonical " - f"v{version} install") - - -def _helper_text() -> str: - body = inspect.getsource(_Glm53MixedPrefill) - return ( - "\n" - + body.replace(MARK_V5, MARK_V6) - + f"\n_GLM53_MIXED = _Glm53MixedPrefill() # {MARK_V6}\n\n" - + f"def _glm53_mixed_prefill_policy(sched, request, computed=None): # {MARK_V6}\n" - + " return _GLM53_MIXED.cap_for(sched, request, computed)\n\n\n" - ) - - -def apply_v6(text: str) -> str: - if "import os\n" not in text.split("import time\n", 1)[0]: - text = replace_once(text, IMPORT_OLD, IMPORT_NEW, "import os") - text = replace_once(text, NEEDLE, _helper_text() + NEEDLE, "helper") - for new, old, label in V6_PAIRS: - text = replace_once(text, old, new, label) - compile(text, str(P), "exec") - return text - - -def main() -> int: - if not P.is_file(): - raise SystemExit(f"missing {P}") - text = P.read_text() - original = text - if MARK_V6 in text: - # Validate existing anchors and helper instead of trusting the marker: a - # marker-only or hand-edited patch cannot round-trip and is refused - # without a write. - clean, start, span = _unpatch(text, 6, exact=_helper_text()) - if apply_v6(clean) != text: - _refuse("v6 helper drifted; unpatch/re-patch is not byte-identical") - print(f"{P.name}: {MARK_V6} already present — verified") - return 0 - for unsupported in (3, 4): - if LEGACY_MARK[unsupported] in text: - _refuse(f"v{unsupported} is refused as unsupported: no authenticated producer of that " - f"helper body was recovered from public history, so a canonical v{unsupported} " - f"image cannot be established and the source is left unchanged") - version = next((v for v in (5, 2, 1) - if LEGACY_MARK[v] in text or (v == 1 and V1_HELPER_START in text)), None) - if version is not None: - clean, start, span = _unpatch(text, version) - _round_trip(text, clean, version, start, span) - text = clean - text = apply_v6(text) - if MARK_V2 in text or MARK_V3 in text or MARK_V4 in text or MARK_V5 in text: - raise SystemExit(f"{P}: older marker left after migration") - if text != original: - P.write_text(text) - print(f"patched {P.name} ({MARK_V6})") - return 0 - - -if __name__ == "__main__": - sys.exit(main()) diff --git a/tests/fixtures/legacy_scheduler_helpers.py b/tests/fixtures/legacy_scheduler_helpers.py index 558097f9..35052e69 100644 --- a/tests/fixtures/legacy_scheduler_helpers.py +++ b/tests/fixtures/legacy_scheduler_helpers.py @@ -11,6 +11,14 @@ v2 9a23b2d8f061 overlay/patch_scheduler_decode_floor.py:HELPER v5 a9bbbadc4c72 overlay/patch_scheduler_decode_floor.py:_helper_text() +Decode-floor v7 (#283) identities, keyed by the installer's identity name: + + v1-image-d9758a6 the v1 helper baked into container image ef9f5013 (build d9758a6); with the v1 hook pairs it rebuilds that image's scheduler.py byte-for-byte + v5-main 943912c overlay/patch_scheduler_decode_floor.py:_helper_text() (main before #221) + v5-priority 9dbd6873 (#221, merged as 2151550) overlay/patch_scheduler_decode_floor.py:_helper_text() + v6-warm-deadline f89e2408 (#180 head) overlay/patch_scheduler_decode_floor.py:_helper_text() + v6-carry dc702b4f (#246 head) overlay/patch_scheduler_decode_floor.py:_helper_text() + v3 and v4 are absent on purpose: 180725a5ce33 introduced their markers as migration targets for intermediate builds whose helper bodies were never recovered from public history, so the installer refuses those markers without @@ -902,3 +910,2081 @@ def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode ''' + +HELPERS['v1-image-d9758a6'] = r''' +def _glm53_mixed_prefill_policy(running, current): + """Mixed-step prefill policy when a peer in `running` is decoding. + + None = no extra policy. 0 = skip this prefill this step. N>0 = cap. + """ + raw = os.environ.get("GLM53_MIXED_PREFILL_CHUNK", "0").strip().lower() + if raw in ("0", "off", "no"): + return None + if raw in ("skip", "-1"): + cap = 0 + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + return None + cur_id = getattr(current, "request_id", None) + for r in running: + if r is current or getattr(r, "request_id", None) == cur_id: + continue + if r.num_computed_tokens >= r.num_prompt_tokens: + return cap + return None + + +''' + +HELPERS['v5-main'] = r''' +class _Glm53MixedPrefill: # [glm53-decode-floor:v5] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) + except ValueError: + self.share = 0.30 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) + except ValueError: + self.max_step_s = 2.0 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self): + grants = (self._open_rec or {}).get("grants", {}) + for r in self._candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + self.defer_reason = reason + self._promote_next() + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + self._promote_next() + self._maybe_log() + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + return 0 if peer_decode else None + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + +_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5] + +def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5] + return _GLM53_MIXED.cap_for(sched, request, computed) + + +''' + +HELPERS['v5-priority'] = r''' +class _Glm53MixedPrefill: # [glm53-decode-floor:v5] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) + except ValueError: + self.share = 0.30 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) + except ValueError: + self.max_step_s = 2.0 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + # [glm53-prefill-priority:v1] Preserve fairness within priority tiers. + int(getattr(r, "priority", 0) or 0) if getattr(self, "priority_mode", False) else 0, + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self): + grants = (self._open_rec or {}).get("grants", {}) + for r in self._candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + self.defer_reason = reason + self._promote_next() + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + policy = getattr(sched, "policy", None) + self.priority_mode = getattr(policy, "value", policy) == "priority" + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + self._promote_next() + self._maybe_log() + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + return 0 if peer_decode else None + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + +_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v5] + +def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v5] + return _GLM53_MIXED.cap_for(sched, request, computed) + + +''' + +HELPERS['v6-warm-deadline'] = r''' +class _Glm53MixedPrefill: # [glm53-decode-floor:v6] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.logged_gate = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.20")) + except ValueError: + self.share = 0.20 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "1000")) / 1000.0) + except ValueError: + self.max_step_s = 1.0 + # v6 gate (warm bypass + deadline). Both features are OFF unless an + # operator sets the knobs: 0 disables each one and is the default, so a + # skip/cap deployment keeps exactly the v5 hold behaviour. + try: + self.warm_tokens = int(self._e("GLM53_MIXED_PREFILL_WARM_TOKENS", "0")) + except ValueError: + self.warm_tokens = 0 + if not 0 <= self.warm_tokens <= 1_000_000: + self.warm_tokens = 0 + try: + self.max_wait_ms = int(self._e("GLM53_MIXED_PREFILL_MAX_WAIT_MS", "0")) + except ValueError: + self.max_wait_ms = 0 + if not 0 <= self.max_wait_ms <= 600_000: + self.max_wait_ms = 0 + try: + self.late_cap = int(self._e("GLM53_MIXED_PREFILL_LATE_CAP", "512")) + except ValueError: + self.late_cap = 512 + if not 64 <= self.late_cap <= 8192: + self.late_cap = 512 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v5 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + if (self.warm_tokens or self.max_wait_ms) and not self.logged_gate: + print( + f"[glm53-decode-floor] gate v6 warm_tokens={self.warm_tokens} " + f"max_wait_ms={self.max_wait_ms} late_cap={self.late_cap}", + flush=True, + ) + self.logged_gate = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self): + grants = (self._open_rec or {}).get("grants", {}) + for r in self._candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + self.defer_reason = reason + self._promote_next() + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + self._promote_next() + self._maybe_log() + + def _warm_bypass(self, remaining): + """Admit a small uncached remainder at once (v6). Off while the knob is 0.""" + return self.warm_tokens > 0 and remaining <= self.warm_tokens + + def _deadline_cap(self, request, remaining): + """Late admission for a request the skip hold would starve (v6). + + Returns 0 while the deadline is disabled (0 ms = wait forever, v5 + behaviour) or not reached, else LATE_CAP tokens for this step. The + first-seen stamp is the per-request arrival the fair policy already + keeps, so chunking, preemption and requeue cannot reset the wait. + """ + if self.max_wait_ms <= 0: + return 0 + rid = request.request_id + first_seen = self.arrival.get(rid) + if first_seen is None: + first_seen = self._now() + self.arrival[rid] = first_seen + if (self._now() - first_seen) * 1000.0 >= self.max_wait_ms: + return max(1, min(self.late_cap, remaining)) + return 0 + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + if not peer_decode: + return None + if self._warm_bypass(remaining): + return None + return self._deadline_cap(request, remaining) + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + if cap is not None and self._warm_bypass(remaining): + cap = None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + +_GLM53_MIXED = _Glm53MixedPrefill() # # [glm53-decode-floor:v6] + +def _glm53_mixed_prefill_policy(sched, request, computed=None): # # [glm53-decode-floor:v6] + return _GLM53_MIXED.cap_for(sched, request, computed) + + +''' + +HELPERS['v6-carry'] = r''' +class _Glm53MixedPrefill: # [glm53-decode-floor:v6] + """Bound contention using completion feedback, without synchronizing GPUs.""" + + LADDER = (128, 256, 512, 768, 1024, 1536, 2048) + COLD_TOK_S = 1300.0 + FIT_WINDOW = 24 + + def __init__(self, now=None): + self._now = now or time.monotonic + self.mode = "skip" + self.legacy_cap = 0 + self.logged_boot = False + self.hist_every = 50 + self._parse() + self.last_service = {} + self.arrival = {} + self.rr_seq = {} + self.served_tokens = {} + self.rr_n = 0 + self.steps = 0 + self.credit = 0.0 + self.in_contention = False + self.inflight = {} + self.mixed_samples = [] + self.solo_samples = [] + self.last_account_mono = None + self.last_prefill_turn_mono = 0.0 + self.step_tag = None + self._sched_id = None + self._open_rec = None + self.selected = set() + self._candidates = [] + self._tried = set() + self._passed = set() + self._carry = set() + self.step_mode = "solo" + self.defer_reason = "none" + self.missed_prefill = 0 + self._model_cache = None + + def _e(self, name, default): + v = os.environ.get(name) + return default if v is None or not str(v).strip() else str(v).strip() + + def _parse(self) -> None: + raw = self._e("GLM53_MIXED_PREFILL_CHUNK", "skip").strip().lower() + self.legacy_cap = 0 + if raw in ("0", "off", "no"): + self.mode = "off" + elif raw in ("skip", "-1"): + self.mode = "skip" + elif raw == "fair": + self.mode = "fair" + else: + try: + cap = int(raw) + except ValueError: + cap = 0 + if cap <= 0: + self.mode = "off" + else: + self.mode = "cap" + self.legacy_cap = cap + try: + self.chunk = int(self._e("GLM53_FAIR_PREFILL_CHUNK", "256")) + except ValueError: + self.chunk = 256 + if self.chunk <= 0: + self.chunk = 256 + try: + self.share = float(self._e("GLM53_FAIR_PREFILL_SHARE", "0.30")) + except ValueError: + self.share = 0.30 + self.share = min(1.0, max(0.0, self.share)) + try: + self.interval_s = int(self._e("GLM53_FAIR_PREFILL_MAX_INTERVAL_MS", "2000")) / 1000.0 + except ValueError: + self.interval_s = 2.0 + if self.interval_s <= 0: + self.interval_s = 2.0 + try: + self.max_chunks = int(self._e("GLM53_FAIR_PREFILL_MAX_CHUNKS", "1")) + except ValueError: + self.max_chunks = 1 + self.max_chunks = max(1, min(self.max_chunks, 16)) + try: + self.max_step_s = max(0.001, int(self._e("GLM53_FAIR_PREFILL_MAX_STEP_MS", "2000")) / 1000.0) + except ValueError: + self.max_step_s = 2.0 + if self.mode == "fair" and not self.logged_boot: + print( + f"[glm53-decode-floor] fair v6 probe_chunk={self.chunk} " + f"ladder={min(self.LADDER)}..{max(self.LADDER)} share={self.share} " + f"interval_s={self.interval_s} max_step_s={self.max_step_s} " + f"max_chunks={self.max_chunks}", + flush=True, + ) + self.logged_boot = True + + + @staticmethod + def prefill_remaining(request, computed=None): + prompt = int(getattr(request, "num_prompt_tokens", 0) or 0) + if computed is None: + computed = int(getattr(request, "num_computed_tokens", 0) or 0) + tokens = int(getattr(request, "num_tokens", prompt) or prompt) + return max(0, max(prompt, tokens - 1) - int(computed)) + + def needs_prefill_compute(self, request, computed=None): + return self.prefill_remaining(request, computed) > 0 + + def _iter_waiting(self, sched): + for name in ("waiting", "skipped_waiting"): + q = getattr(sched, name, None) + if not q: + continue + try: + for r in q: + yield r + except TypeError: + continue + + def _live_ids(self, sched): + ids = set() + for r in list(getattr(sched, "running", None) or []): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + for r in self._iter_waiting(sched): + rid = getattr(r, "request_id", None) + if rid is not None: + ids.add(rid) + reqs = getattr(sched, "requests", None) + if isinstance(reqs, dict): + ids.update(reqs.keys()) + return ids + + def _prune(self, live): + for store in (self.last_service, self.arrival, self.rr_seq, + self.served_tokens): + dead = [k for k in store if k not in live] + for k in dead: + store.pop(k, None) + + + @property + def inflight_prefill(self): + return sum(bool(r["prefill_tokens"]) for r in self.inflight.values()) + + def _shape(self, decodes, prefills): + history = max((int(r.num_computed_tokens) for r in decodes), default=0) + position = max((int(r.num_computed_tokens) for r in prefills), default=0) + return (len(decodes), (history // 4096).bit_length(), + (position // 4096).bit_length()) + + def _cost_model(self): + """Fit step cost dt = a + b*n over recent prefill-bearing steps. + + Every prefill-bearing step pays a fixed cost on this kit (~0.3 s host + time for 82..256 tokens, ~2.7 s for 3584), so a mixed step is close to + a solo chunk plus a few decode rows: solo samples are pooled for the + fit. The fit is then scaled so recent mixed samples are not + underestimated (75th percentile of actual/fit, clamped to [1, 1.5]). + Needs two distinct chunk sizes; otherwise returns None and the caller + scales linearly. v4 scaled one sample linearly, priced 1024 tokens off + 128-token samples at ~3x the real cost, and never climbed back. + """ + if self._model_cache is not None: + return self._model_cache + mixed = [(n, dt) for _, n, dt in self.mixed_samples[-self.FIT_WINDOW:]] + pts = mixed + [(n, dt) for n, dt in self.solo_samples[-self.FIT_WINDOW:]] + if len(pts) < 2 or len({n for n, _ in pts}) < 2: + return None + cnt = float(len(pts)) + sx = float(sum(n for n, _ in pts)) + sy = float(sum(dt for _, dt in pts)) + sxx = float(sum(n * n for n, _ in pts)) + sxy = float(sum(n * dt for n, dt in pts)) + den = cnt * sxx - sx * sx + b = max(0.0, (cnt * sxy - sx * sy) / den) if den > 0 else 0.0 + a = max(0.0, (sy - b * sx) / cnt) + if mixed: + ratios = sorted(dt / max(1e-6, a + b * n) for n, dt in mixed) + r = ratios[min(len(ratios) - 1, int(0.75 * len(ratios)))] + else: + r = 1.1 # decode rows not sampled yet + r = min(1.5, max(1.0, r)) + self._model_cache = (a * r, b * r) + return self._model_cache + + def _est_dt(self, tokens, shape=None): + """Conservative host-time estimate; not a guaranteed execution bound.""" + n = max(1, int(tokens)) + model = self._cost_model() + if model is not None: + a, b = model + return max(0.01, a + b * n) + samples = self.mixed_samples[-8:] + if samples: + # One size only: scale linearly with a fixed-cost floor. + return max(0.01, max(dt * max(0.5, n / t) for _, t, dt in samples)) + return max(0.05, n / self.COLD_TOK_S) + + def _target(self, remaining, room): + """Largest rung whose estimated step fits `room` (tokens per accounted + second rise with size under a fixed per-step cost), or None.""" + rungs = {self.chunk} + if self.mixed_samples or self._cost_model() is not None: + rungs.update(self.LADDER) + if remaining is not None: + rungs = {min(n, remaining) for n in rungs} + fitting = [(n, self._est_dt(n)) for n in sorted(rungs)] + fitting = [(n, cost) for n, cost in fitting if cost <= room + 1e-9] + return fitting[-1] if fitting else None + + def _credit_limit(self): + return min(self.max_step_s, self._est_dt(max(self.chunk, max(self.LADDER)))) + + def _rank_prefills(self, prefills): + return sorted(prefills, key=lambda r: ( + self.last_service.get(r.request_id, 0.0), + self.rr_seq.get(r.request_id, 0), self.arrival[r.request_id], r.request_id)) + + def _promote_next(self, prefer_passed=False): + grants = (self._open_rec or {}).get("grants", {}) + candidates = self._candidates + if prefer_passed: + candidates = ([r for r in candidates if r.request_id in self._passed] + + [r for r in candidates if r.request_id not in self._passed]) + for r in candidates: + if len(self.selected | set(grants)) >= self.max_chunks: + break + if r.request_id not in self._tried: + self.selected.add(r.request_id) + if prefer_passed and r.request_id in self._passed: + # Allocation failed after this request's scheduler position + # had already passed. Carry it into the next step instead of + # resetting to the same blocked head of the fair queue. + self._carry.add(r.request_id) + + def _release(self, rid, reason="allocation"): + rec = self._open_rec + if rec is None: + return + # The running loop reports zero even when policy returned cap=0 for an + # unselected request. That request was passed, not attempted; keep it + # eligible for allocation-aware carry if a later selected request fails. + if rid not in self.selected and rid not in rec["grants"]: + return + grant = rec["grants"].pop(rid, None) + if grant: + self.credit += grant[1] + if grant[2]: + rec["borrowed"] = False + self.selected.discard(rid) + self._tried.add(rid) + if reason == "zero_progress": + # A selected request that could not make progress rotates behind + # never-served peers. Credit/gap deferrals are not admissions and + # retain their fair rank. + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.defer_reason = reason + self._promote_next(prefer_passed=True) + + def note_scheduled(self, request, num_new_tokens): + # Called after alignment/encoder caps; final allocation is sealed below. + if num_new_tokens <= 0: + self._release(request.request_id, "zero_progress") + + def protect_decode(self, request): + return (self.mode == "fair" and self._open_rec is not None + and self._open_rec["had_decode"] + and self.needs_prefill_compute(request)) + + def begin_step(self, sched): + now = self._now() + sid = id(sched) + tag = (sid, int(sched.current_step)) + if self.step_tag == tag: + return + if self._sched_id is not None and self._sched_id != sid: + self.inflight.clear() + self._open_rec = None + self._carry.clear() + self.credit = 0.0 + self.in_contention = False + self.last_account_mono = None + self._sched_id = sid + # An unfinished/failed schedule has dispatched nothing: refund its grants. + if self._open_rec: + self.credit += sum(g[1] for g in self._open_rec["grants"].values()) + self.step_tag = tag + self.steps += 1 + self._prune(self._live_ids(sched)) + running = list(sched.running) + waiting = list(self._iter_waiting(sched)) + prefills = list({r.request_id: r for r in running + waiting + if self.needs_prefill_compute(r)}.values()) + decodes = [r for r in running if not self.needs_prefill_compute(r)] + # The base loop checks eligibility and reserves BOTH token and input/draft + # capacity by executing these requests first. Preserve order within groups. + if self.mode == "fair": + sched.running[:] = decodes + [r for r in running + if self.needs_prefill_compute(r)] + for r in running + waiting: + self.arrival.setdefault(r.request_id, now) + self._candidates = self._rank_prefills(prefills) + self._tried = set() + self.selected = set() + self._passed = set() + live_prefills = {r.request_id for r in prefills} + self._carry.intersection_update(live_prefills) + sched._glm53_align_prefill_limit = None + self._open_rec = { + "step_id": int(sched.current_step), "t_submit": now, + "had_decode": bool(decodes), "had_prefill_demand": bool(prefills), + "shape": self._shape(decodes, prefills), "grants": {}, + "borrowed": False, + } + # Do not forgive debt while a decoder or its outstanding work remains. + if not decodes and not any(r["had_decode"] for r in self.inflight.values()): + self.in_contention = False + self.credit = 0.0 + self._carry.clear() + self.step_mode = "legacy" if self.mode != "fair" else "solo" + self.defer_reason = "none" + if self.mode != "fair" or not prefills or not decodes: + return + if not self.in_contention: + self.credit = min(self.max_step_s, self._est_dt(self.chunk)) + self.in_contention = True + if self.inflight_prefill: + self.step_mode = "decode_only" + self.defer_reason = "async_inflight" + else: + self.step_mode = "prefill_turn" + for r in self._candidates: + if r.request_id in self._carry: + self.selected.add(r.request_id) + if len(self.selected) >= self.max_chunks: + break + self._promote_next() + self._maybe_log() + + def cap_for(self, sched, request, computed=None): + self.begin_step(sched) + sched._glm53_align_prefill_limit = None + remaining = self.prefill_remaining(request, computed) + if remaining <= 0 or self.mode == "off": + return None + peer_decode = any(r is not request and not self.needs_prefill_compute(r) + for r in sched.running) + if self.mode == "skip": + return 0 if peer_decode else None + if self.mode == "cap": + cap = self.legacy_cap if peer_decode else None + sched._glm53_align_prefill_limit = cap + return cap + if self.step_mode == "solo": + return None + rid = request.request_id + rec = self._open_rec + if rid in rec["grants"]: + return rec["grants"][rid][0] + self._passed.add(rid) + if self.step_mode != "prefill_turn" or rid not in self.selected: + return 0 + self._carry.discard(rid) + reserved = sum(g[1] for g in rec["grants"].values()) + gap_room = max(0.0, self.max_step_s - reserved) + age = self._now() - self.last_service.get(rid, self.arrival[rid]) + pick = self._target(remaining, gap_room) + if pick is None: + self._release(rid, "gap_budget") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + # Save credit for the most efficient rung that fits the step budget + # instead of spending it on a smaller chunk now (v4 did, and stalled + # at 128-token steps once one expensive sample priced 256 too high). + cap, cost = pick + borrowed = False + if cost > self.credit + 1e-9: + # A never-served newcomer gets a prompt probe; afterwards service + # ages. Either may borrow ONE step-bounded chunk globally, only + # after all shared debt is repaid: queue churn or several aged + # newcomers cannot repeatedly overdraw the decoder. + due = rid not in self.last_service or age >= self.interval_s + if (due and self.credit >= -1e-9 and not rec["grants"] + and not rec["borrowed"]): + borrowed = True + else: + self._release(rid, "credit") + if age >= self.interval_s: + self.missed_prefill += 1 + self._maybe_log() + return 0 + self.credit -= cost + rec["grants"][rid] = (cap, cost, borrowed) + rec["borrowed"] |= borrowed + self._tried.add(rid) + sched._glm53_align_prefill_limit = cap + return cap + + def finish_step(self, sched, scheduler_output): + """Seal final work before _update_after_schedule advances request state.""" + rec = self._open_rec + self._open_rec = None + if rec is None: + return + scheduled = {rid: int(n) for rid, n in scheduler_output.num_scheduled_tokens.items() + if n > 0} + requests = getattr(sched, "requests", {}) + prefill = {} + for rid, n in scheduled.items(): + request = requests.get(rid) + if request is not None: + amount = min(n, self.prefill_remaining(request)) + if amount > 0: + prefill[rid] = amount + reserved = 0.0 + for rid, (_, estimate, _) in rec["grants"].items(): + # Never publish/charge a tentative grant that alignment, allocation, + # encoder caps or preemption removed from the final scheduler output. + actual_est = min(estimate, self._est_dt(prefill[rid], rec["shape"])) if rid in prefill else 0.0 + self.credit += estimate - actual_est + reserved += actual_est + if not scheduled: + # Empty schedules do not necessarily have a completion callback. + return + rec.update(output=scheduler_output, scheduled=scheduled, + prefill_tokens=prefill, reserved=reserved) + del rec["grants"] + if self.mode == "fair": + self.inflight[id(scheduler_output)] = rec + + def observe_output(self, sched, scheduler_output): + rec = self.inflight.pop(id(scheduler_output), None) + if rec is None or rec["output"] is not scheduler_output: + return # Empty, duplicate or unrelated callback; never pop another step. + now = self._now() + # Partition observed busy wall time instead of adding overlapping + # submit-to-completion latencies of queued async steps. Queue residence + # and engine idle gaps therefore cannot mint decode credit twice. + start = rec["t_submit"] + if self.last_account_mono is not None: + start = max(start, self.last_account_mono) + dt = max(0.0, now - start) + self.last_account_mono = max(now, self.last_account_mono or now) + actual = scheduler_output.num_scheduled_tokens + served = {rid: min(n, max(0, int(actual.get(rid, 0)))) + for rid, n in rec["prefill_tokens"].items() + if int(actual.get(rid, 0)) > 0} + for rid, n in served.items(): + self.last_service[rid] = now + self.rr_n += 1 + self.rr_seq[rid] = self.rr_n + self.served_tokens[rid] = self.served_tokens.get(rid, 0) + n + n = sum(served.values()) + if rec["had_decode"]: + # Reservation was already debited; settle once, including overruns. + self.credit += rec["reserved"] + self.share * dt - (dt if n else 0.0) + self.credit = min(self.credit, self._credit_limit()) + if n and dt > 0: + self.mixed_samples.append((rec["shape"], n, dt)) + self.mixed_samples = self.mixed_samples[-64:] + self.last_prefill_turn_mono = now + self._model_cache = None + elif n and dt > 0: + self.solo_samples.append((n, dt)) + self.solo_samples = self.solo_samples[-32:] + self._model_cache = None + if n and self.hist_every > 0: + print(f"[glm53-decode-floor] completed_step={rec['step_id']} " + f"contention={int(rec['had_decode'])} prefill_tokens={served} " + f"reserved_s={rec['reserved']:.3f} accounted_s={dt:.3f} " + f"credit={self.credit:.3f} timing=host_busy_proxy", flush=True) + self._prune(self._live_ids(sched)) + + def _maybe_log(self): + if self.hist_every <= 0 or self.steps % self.hist_every != 1: + return + rec = self._open_rec or {} + remaining = sum(self.prefill_remaining(r) for r in self._candidates) + target, estimate = (self._target(None, self.max_step_s) + or (self.chunk, self._est_dt(self.chunk))) + rate = target / estimate if estimate > 0 else 0.0 + eta = remaining / (self.share * rate) if self.share * rate > 0 else float("inf") + model = self._cost_model() + fit = (f"fit_fixed_s={model[0]:.3f} fit_us_per_tok={model[1] * 1e6:.0f}" + if model else "fit=none") + print(f"[glm53-decode-floor] step={rec.get('step_id', self.steps)} " + f"mode={self.step_mode} defer={self.defer_reason} " + f"inflight={self.inflight_prefill} credit={self.credit:.3f} " + f"remaining={remaining} target={target} est_s={estimate:.3f} {fit} " + f"eta_est_s={eta:.1f} max_step_s={self.max_step_s:.3f} " + f"missed={self.missed_prefill} timing=host_busy_proxy", flush=True) + + @staticmethod + def aligned_new_tokens( + start, num_new, prefill_end, block_size, max_prefill_tokens, policy_cap=None + ) -> int: + """Hybrid align clip. policy_cap is an intentional mixed cap, not leftover budget.""" + if policy_cap is not None and policy_cap > 0: + max_prefill_tokens = min(max_prefill_tokens, policy_cap) + end = start + num_new + if end < prefill_end: + aligned_end = end // block_size * block_size + if aligned_end > start or block_size <= max_prefill_tokens: + end = aligned_end + return max(0, end - start) + +_GLM53_MIXED = _Glm53MixedPrefill() # [glm53-decode-floor:v6] + +def _glm53_mixed_prefill_policy(sched, request, computed=None): # [glm53-decode-floor:v6] + return _GLM53_MIXED.cap_for(sched, request, computed) + + +''' diff --git a/tests/test_scheduler_decode_floor.py b/tests/test_scheduler_decode_floor.py index f5fde3bf..035e4e17 100644 --- a/tests/test_scheduler_decode_floor.py +++ b/tests/test_scheduler_decode_floor.py @@ -514,26 +514,25 @@ def load_fixture(name, path): legacy = load_fixture('legacy_scheduler_helpers', fixtures / 'legacy_scheduler_helpers.py') - def legacy_producer(version): - text = clean.replace(mod.NEEDLE, legacy.HELPERS[version] + mod.NEEDLE, 1) - for new, old, label in (mod.V1_PAIRS if version == 1 else mod.V2_PAIRS): + # Every accepted public identity, rebuilt from its byte-exact helper text at the + # installer's anchor plus that version's hook pairs (the installer authenticates + # each by sha256 and length before migrating it). + pairs = {'v1': mod.V1_PAIRS, 'v1-image-d9758a6': mod.V1_PAIRS, 'v2': mod.V2_PAIRS, + 'v5-historical': mod.V5_PAIRS, 'v5-main': mod.V5_PAIRS, 'v5-priority': mod.V5_PAIRS, + 'v6-warm-deadline': mod.V6_PAIRS, 'v6-carry': mod.V6_PAIRS} + helpers = {'v1': legacy.HELPERS[1], 'v2': legacy.HELPERS[2], 'v5-historical': legacy.HELPERS[5], + **{k: legacy.HELPERS[k] for k in ('v1-image-d9758a6', 'v5-main', 'v5-priority', + 'v6-warm-deadline', 'v6-carry')}} + + def producer(ident): + text = clean.replace(mod.NEEDLE, helpers[ident] + mod.NEEDLE, 1) + for new, old, label in pairs[ident]: text = mod.replace_once(text, old, new, label) return text - producers = {'pristine': clean, 'v1': legacy_producer(1), 'v2': legacy_producer(2)} - image_helper = load_fixture('v1_image_helper', fixtures / 'decode_floor_producers' / 'v1_image_d9758a6_helper.py') - v1_image = clean.replace(mod.NEEDLE, image_helper.HELPER + mod.NEEDLE, 1) - for new, old, label in mod.V1_PAIRS: - v1_image = mod.replace_once(v1_image, old, new, label) - producers['v1-image-d9758a6'] = v1_image + producers = {'pristine': clean, **{ident: producer(ident) for ident in helpers}} if found is not None: producers[found[0]] = image - # Byte-exact public producers of every other non-v7 identity the installer accepts. - for ident, fname in (('v5-historical', 'v5_historical_8e2a32b2.py'), ('v5-main', 'v5_main_d5713ec6.py'), - ('v5-priority', 'v5_priority_763bd30e.py'), ('v6-warm-deadline', 'v6_warm_deadline_1f74558d.py'), - ('v6-carry', 'v6_carry_0870fab9.py')): - producer = load_fixture('producer_' + ident.replace('-', '_'), fixtures / 'decode_floor_producers' / fname) - producers[ident] = (getattr(producer, 'apply_v6', None) or producer.apply_v5)(clean) def run(text, temp): target = Path(temp) / 'scheduler.py' From f5830dde4e0ab0a3faceb93c3b26d6a7a83800e0 Mon Sep 17 00:00:00 2001 From: Hector Isaza <76170012+hecisaza@users.noreply.github.com> Date: Thu, 1 Oct 2026 03:54:22 +0000 Subject: [PATCH 7/8] scheduler: refuse import drift on legacy installs; ship the migration fixture in the image; record legacy digests Address review on #309 (b852f57): - _identify() refuses a legacy install whose `import os` is missing ("import drifted") instead of repairing it during migration, matching the check the current-v7 path already had. Three refusal cases added; the refusal test now also requires a deliberate refusal message, not a crash. - Dockerfile: the fixtures COPY now ships tests/fixtures/legacy_scheduler_helpers.py to /opt/glm53/fixtures/, and the runner resolves its fixture beside itself before the checkout tree, so the mandatory build-time self-test (python3 /opt/glm53/test_scheduler_decode_floor.py) finds it. New tests/test_image_test_layout.py stages the Dockerfile's /opt/glm53 layout from its COPY lines and checks that resolution (and, with GLM53_SCHEDULER_PY_SRC set, runs the full runner there). - legacy_scheduler_helpers.py records PINNED_CLEAN_SHA256 and a SOURCE_DIGESTS entry per reconstructed legacy install on the pinned source. On that source the installation test compares each rebuild against its recorded digest and checks the comparison's sensitivity, so the installer's hook tables are checked against recorded history rather than against themselves. A wrong begin_step hook in the V5 or V6 table is now caught. Installer-only change: the installed scheduler on image ef9f5013 is still sha256 97c90a18fd0df6f9aeb1452c39801455d90756d315c6846ca0642e934fe4e269. --- Dockerfile | 2 +- overlay/patch_scheduler_decode_floor.py | 4 ++ tests/fixtures/legacy_scheduler_helpers.py | 17 ++++++ tests/test_image_test_layout.py | 61 ++++++++++++++++++++++ tests/test_scheduler_decode_floor.py | 34 ++++++++++-- 5 files changed, 114 insertions(+), 4 deletions(-) create mode 100644 tests/test_image_test_layout.py diff --git a/Dockerfile b/Dockerfile index 06958474..92a7db83 100644 --- a/Dockerfile +++ b/Dockerfile @@ -485,7 +485,7 @@ COPY tests/test_indexer_workspace.py tests/test_indexer_warmup_range.py /opt/glm COPY overlay/patch_tool_choice_none.py /opt/glm53/patch_tool_choice_none.py COPY overlay/patch_loadclone.py /opt/glm53/patch_loadclone.py COPY tests/test_loadclone.py /opt/glm53/test_loadclone.py -COPY tests/fixtures/loadclone_weight_utils.py.txt tests/fixtures/kpool_tail_seed_kernel-487ecf187.py.txt tests/fixtures/kpool_tail_seed_kernel-db1bfdd.py.txt /opt/glm53/fixtures/ +COPY tests/fixtures/loadclone_weight_utils.py.txt tests/fixtures/kpool_tail_seed_kernel-487ecf187.py.txt tests/fixtures/kpool_tail_seed_kernel-db1bfdd.py.txt tests/fixtures/legacy_scheduler_helpers.py /opt/glm53/fixtures/ COPY tests/test_tool_choice_none.py /opt/glm53/test_tool_choice_none.py COPY overlay/ablit_runtime.py /opt/glm53/ablit_runtime.py COPY overlay/patch_ablit.py /opt/glm53/patch_ablit.py diff --git a/overlay/patch_scheduler_decode_floor.py b/overlay/patch_scheduler_decode_floor.py index 1664a00a..76b9dc63 100644 --- a/overlay/patch_scheduler_decode_floor.py +++ b/overlay/patch_scheduler_decode_floor.py @@ -1388,6 +1388,10 @@ def _identify(text: str): reasons.append(str(exc).split("scheduler: ", 1)[-1]) continue _round_trip(text, clean, start, span, pairs, name) + # Every legacy helper uses `os`; the import is part of the applied state and + # must be present before a legacy identity is accepted for migration. + if "import os\n" not in text.split("import time\n", 1)[0]: + _refuse(f"{name}: import drifted") return name, clean, start _refuse(f"no authenticated identity for {marker}: " + " | ".join(reasons)) diff --git a/tests/fixtures/legacy_scheduler_helpers.py b/tests/fixtures/legacy_scheduler_helpers.py index 35052e69..aea761fe 100644 --- a/tests/fixtures/legacy_scheduler_helpers.py +++ b/tests/fixtures/legacy_scheduler_helpers.py @@ -28,6 +28,23 @@ HELPERS = {} +# Independent record of each reconstructed legacy install on the pinned scheduler +# source (image ef9f5013 with its baked v1 helper removed), i.e. helper text at the +# installer's anchor plus that version's hook insertions. The installation test +# compares its rebuilds against these, so the installer's hook tables are checked +# against recorded history rather than against themselves. +PINNED_CLEAN_SHA256 = '097be6dc0c82371f17c64d05db8f012095fcccfb958299cea8b0bc940405e3a6' +SOURCE_DIGESTS = { + 'v1': 'cd0bd6678c0b74a73e49ae78fe86517adc5ef5b136aa96686e1ee99d4a1b691c', + 'v2': 'd8acecff7eb214e64c9ce728d237656a6334af7fe4dba704d56ef438a9719412', + 'v5-historical': '437ea4c1c6eb7510d29d049d74c24dfdde7f2e0f2ebf5e0025b8a576eb51bf57', + 'v1-image-d9758a6': 'f8c8b15b5d33ec2cd4eee73d07ecea7957a7ac6692e74485f50aa3449ce1285f', + 'v5-main': '62669979a25b721a0918a0d10e1982db9d6772279e42cb10cddff5727e480266', + 'v5-priority': '82d0846d6eefcf314d754e1682767f82f00d1b55bacac68e62fe53f1d47eb2eb', + 'v6-warm-deadline': 'e479072623c18d59a062e4865e84f3b8a32e98a3e7469eced03de9c5bc0ef066', + 'v6-carry': '46c006b04df827d48416c31d0e2cf4c393e9875b493f70dbb5fd43334302aef0', +} + HELPERS[1] = r''' def _glm53_mixed_prefill_policy(running, current): """Mixed-step prefill policy when a peer in `running` is decoding. diff --git a/tests/test_image_test_layout.py b/tests/test_image_test_layout.py new file mode 100644 index 00000000..22905f1e --- /dev/null +++ b/tests/test_image_test_layout.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 +"""The Dockerfile stages the decode-floor installation runner at /opt/glm53 and runs it +at build time; this stages that same layout from the Dockerfile's COPY lines, with no +checkout-relative paths available, and checks the runner finds its overlay and fixture. +With GLM53_SCHEDULER_PY_SRC set, the full runner is executed in that layout as well.""" +from __future__ import annotations + +import os +from pathlib import Path +import re +import shutil +import subprocess +import sys +import tempfile +import unittest + +ROOT = Path(__file__).resolve().parent.parent +RUNNER = "tests/test_scheduler_decode_floor.py" + + +def staged_copies(): + """(source, destination) pairs for every COPY into /opt/glm53 in the Dockerfile.""" + pairs = [] + for line in (ROOT / "Dockerfile").read_text().splitlines(): + m = re.match(r"COPY\s+(.+?)\s+(/opt/glm53\S*)\s*$", line) + if not m: + continue + sources, dest = m.group(1).split(), m.group(2) + for src in sources: + target = dest + Path(src).name if dest.endswith("/") else dest + pairs.append((src, target)) + return pairs + + +class ImageTestLayout(unittest.TestCase): + def test_runner_resolves_overlay_and_fixture_in_image_layout(self): + pairs = staged_copies() + self.assertIn((RUNNER, "/opt/glm53/test_scheduler_decode_floor.py"), pairs) + with tempfile.TemporaryDirectory() as tmp: + for src, dest in pairs: + if src.startswith(("tests/", "overlay/")) and (ROOT / src).is_file(): + out = Path(tmp) / dest.lstrip("/") + out.parent.mkdir(parents=True, exist_ok=True) + shutil.copy(ROOT / src, out) + staged = Path(tmp) / "opt/glm53/test_scheduler_decode_floor.py" + probe = ("import importlib.util, sys; spec = importlib.util.spec_from_file_location('r', sys.argv[1]); " + "m = importlib.util.module_from_spec(spec); spec.loader.exec_module(m); " + "print(m.PATCH); print(m.FIXTURES / 'legacy_scheduler_helpers.py')") + result = subprocess.run([sys.executable, "-c", probe, str(staged)], capture_output=True, text=True, timeout=120) + self.assertEqual(result.returncode, 0, result.stderr[-800:]) + patch_path, fixture_path = result.stdout.strip().splitlines()[-2:] + self.assertEqual(Path(patch_path), Path(tmp) / "opt/glm53/patch_scheduler_decode_floor.py") + self.assertEqual(Path(fixture_path), Path(tmp) / "opt/glm53/fixtures/legacy_scheduler_helpers.py") + self.assertTrue(Path(fixture_path).is_file()) + if os.environ.get("GLM53_SCHEDULER_PY_SRC"): + run = subprocess.run([sys.executable, str(staged)], capture_output=True, text=True, timeout=900) + self.assertEqual(run.returncode, 0, run.stderr[-800:]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_scheduler_decode_floor.py b/tests/test_scheduler_decode_floor.py index 035e4e17..5c97de1e 100644 --- a/tests/test_scheduler_decode_floor.py +++ b/tests/test_scheduler_decode_floor.py @@ -3,6 +3,7 @@ from __future__ import annotations import ast +import hashlib import contextlib import importlib.util import io @@ -28,6 +29,11 @@ + ', '.join(str(p) for p in _PATCH_CANDIDATES) + ')' ) +_FIXTURE_CANDIDATES = (HERE / 'fixtures', ROOT / 'tests' / 'fixtures') # image layout, then checkout +FIXTURES = next((p for p in _FIXTURE_CANDIDATES if (p / 'legacy_scheduler_helpers.py').is_file()), None) +if FIXTURES is None: + raise SystemExit('missing fixtures/legacy_scheduler_helpers.py (tried ' + + ', '.join(str(p) for p in _FIXTURE_CANDIDATES) + ')') spec = importlib.util.spec_from_file_location('glm53_decode_floor', PATCH) mod = importlib.util.module_from_spec(spec) sys.modules[spec.name] = mod @@ -503,8 +509,6 @@ def installation_tests(): clean = image if found is None else found[1] assert mod.MARK_V7 not in clean and mod.CLASS_HEAD not in clean - fixtures = ROOT / 'tests' / 'fixtures' - def load_fixture(name, path): spec = importlib.util.spec_from_file_location(name, path) module = importlib.util.module_from_spec(spec) @@ -512,7 +516,7 @@ def load_fixture(name, path): spec.loader.exec_module(module) return module - legacy = load_fixture('legacy_scheduler_helpers', fixtures / 'legacy_scheduler_helpers.py') + legacy = load_fixture('legacy_scheduler_helpers', FIXTURES / 'legacy_scheduler_helpers.py') # Every accepted public identity, rebuilt from its byte-exact helper text at the # installer's anchor plus that version's hook pairs (the installer authenticates @@ -531,6 +535,24 @@ def producer(ident): return text producers = {'pristine': clean, **{ident: producer(ident) for ident in helpers}} + # On the pinned source, every rebuilt legacy install must match its recorded digest: + # this checks the installer's hook tables against history, not against themselves. + on_pinned = hashlib.sha256(clean.encode()).hexdigest() == legacy.PINNED_CLEAN_SHA256 + if on_pinned: + for ident in helpers: + got = hashlib.sha256(producers[ident].encode()).hexdigest() + assert got == legacy.SOURCE_DIGESTS[ident], f'{ident}: rebuilt install differs from its recorded bytes' + # Sensitivity: a wrong hook insertion must be caught by that comparison. + for ident, table in (('v5-main', mod.V5_PAIRS), ('v6-carry', mod.V6_PAIRS)): + wrong = tuple((new.replace('_GLM53_MIXED.begin_step(self)', '_GLM53_MIXED.broken_begin_step(self)'), old, label) + for new, old, label in table) + assert wrong != table, ident + text = clean.replace(mod.NEEDLE, helpers[ident] + mod.NEEDLE, 1) + for new, old, label in wrong: + text = mod.replace_once(text, old, new, label) + assert hashlib.sha256(text.encode()).hexdigest() != legacy.SOURCE_DIGESTS[ident], f'{ident}: digest check is blind to a wrong hook' + else: + print('note: scheduler source is not the pinned image source; recorded legacy digests not compared') if found is not None: producers[found[0]] = image @@ -605,11 +627,17 @@ def run(text, temp): # Only the import anchor or adaptive-K may follow the helper. 'v7-unknown-text-after-helper': installed.replace( mod.NEEDLE, 'class _Glm53OtherOverlay:\n pass\n\n' + mod.NEEDLE, 1), + # The overlay's `import os` is part of the applied state: a legacy or current + # install missing it is drift, not something to repair in place. + **{f'{ident}-import-removed': text.replace('import os\n', '', 1) + for ident, text in (('v1-image-d9758a6', producers['v1-image-d9758a6']), + ('v5-main', producers['v5-main']), ('v7', installed))}, } for case, text in refused.items(): assert text != installed and text != v5, f'{case}: mutation did not apply' result, after = run(text, temp) assert result.returncode != 0 and after == text, f'{case} was not refused without writing' + assert 'refusing to rewrite' in result.stderr, f'{case}: not a deliberate refusal: {result.stderr[-200:]}' return installed From ef5a0dd0dd2f6d2b6dcd8337d05091a2534a2dc2 Mon Sep 17 00:00:00 2001 From: Hector Isaza <76170012+hecisaza@users.noreply.github.com> Date: Thu, 1 Oct 2026 04:30:20 +0000 Subject: [PATCH 8/8] tests: migrate supplied legacy sources at their own position; record legacy hook-table digests Address review on #309 (f5830dd): - installation_tests() now falls back to legacy identification for any v7-marked source that is not the current helper (as the installer does), so the runner can execute against the accepted pre-release v7 identities. A supplied legacy install of any identity is migrated and compared against apply_v7(clean, at=) and re-applied as a no-op, instead of being compared with the fresh-install bytes; with adaptive-K between the two positions that comparison was wrong. The same position rule is covered with the public v5-main fixture, helper before and after adaptive-K. - legacy_scheduler_helpers.py records HOOK_TABLE_DIGESTS for the V1/V2/V5/V6 hook tables. They are checked unconditionally (no scheduler source needed, also as a plain pytest test), so a changed table is caught on build-shaped and deployed sources too, where the whole-source digests cannot apply. A wrong begin_step hook in the V5 or V6 table now fails on every source layout. Tests only; the installed scheduler on image ef9f5013 is still sha256 97c90a18fd0df6f9aeb1452c39801455d90756d315c6846ca0642e934fe4e269. --- tests/fixtures/legacy_scheduler_helpers.py | 11 +++++ tests/test_scheduler_decode_floor.py | 56 ++++++++++++++++++---- 2 files changed, 59 insertions(+), 8 deletions(-) diff --git a/tests/fixtures/legacy_scheduler_helpers.py b/tests/fixtures/legacy_scheduler_helpers.py index aea761fe..9507e291 100644 --- a/tests/fixtures/legacy_scheduler_helpers.py +++ b/tests/fixtures/legacy_scheduler_helpers.py @@ -34,6 +34,17 @@ # compares its rebuilds against these, so the installer's hook tables are checked # against recorded history rather than against themselves. PINNED_CLEAN_SHA256 = '097be6dc0c82371f17c64d05db8f012095fcccfb958299cea8b0bc940405e3a6' + +# Independent record of each legacy version's hook insertions (the installer's +# V1/V2/V5/V6 pair tables, as sha256 of their JSON form). Checked unconditionally, +# with no scheduler source needed, so a changed table cannot silently rewrite +# both the migration and the test input that is supposed to check it. +HOOK_TABLE_DIGESTS = { + 'v1': 'd4eb922526f9d7945ccf9e468383707a638536a86d3fa9a6fa755d790ef63339', + 'v2': '2813e9969de28e6b5aedb98bfb010914f6f3746c2fb5e52b993480be8730fa5c', + 'v5': '02def284d1ae0c5491063dd9c914bfb0ad3a9d26e8ed7784eeea850ccb1b7893', + 'v6': '3597f2869c66f6e7cea9892bc95325146293ab6bfd3b5da7031ded482e75d67c', +} SOURCE_DIGESTS = { 'v1': 'cd0bd6678c0b74a73e49ae78fe86517adc5ef5b136aa96686e1ee99d4a1b691c', 'v2': 'd8acecff7eb214e64c9ce728d237656a6334af7fe4dba704d56ef438a9719412', diff --git a/tests/test_scheduler_decode_floor.py b/tests/test_scheduler_decode_floor.py index 5c97de1e..08de5bb6 100644 --- a/tests/test_scheduler_decode_floor.py +++ b/tests/test_scheduler_decode_floor.py @@ -7,6 +7,7 @@ import contextlib import importlib.util import io +import json import os from pathlib import Path import subprocess @@ -483,6 +484,21 @@ def test_finished_decode_does_not_hold_phantom_input_reservation(self): self.assertEqual(ns['input_budget'], 0) +def hook_table_digests_check(legacy): + """The installer's legacy hook tables must match their independently recorded digests.""" + for version, table in (('v1', mod.V1_PAIRS), ('v2', mod.V2_PAIRS), ('v5', mod.V5_PAIRS), ('v6', mod.V6_PAIRS)): + got = hashlib.sha256(json.dumps([list(t) for t in table]).encode()).hexdigest() + assert got == legacy.HOOK_TABLE_DIGESTS[version], f'{version} hook table differs from its recorded digest' + + +class HookTableDigests(unittest.TestCase): + def test_legacy_hook_tables_match_recorded_digests(self): + spec = importlib.util.spec_from_file_location('legacy_scheduler_helpers_digests', FIXTURES / 'legacy_scheduler_helpers.py') + legacy = importlib.util.module_from_spec(spec) + spec.loader.exec_module(legacy) + hook_table_digests_check(legacy) + + def installation_tests(): """Installer contract for decode-floor v7 (#283; integrates #246/#221/#180). @@ -498,13 +514,18 @@ def installation_tests(): if src is None: raise SystemExit('Set GLM53_SCHEDULER_PY_SRC to the pinned scheduler source') image = src.read_text() - current = None + current = found = None if mod.MARK_V7 in image: - # An already-current v7 install: verify it exactly (as main() does), then unpatch. - clean, current_at, _ = mod._unpatch(image, 'v7', mod.MARK_V7, mod.CLASS_HEAD, mod.V7_PAIRS, - exact=mod._helper_text()) - found, current = None, image - else: + # An already-current v7 install verifies exactly (as main() does); any other + # v7-marked source falls back to legacy identification, also as main() does + # (the accepted pre-release v7 identities carry the same marker). + try: + clean, current_at, _ = mod._unpatch(image, 'v7', mod.MARK_V7, mod.CLASS_HEAD, mod.V7_PAIRS, + exact=mod._helper_text()) + current = image + except SystemExit: + pass + if current is None: found = mod._identify(image) clean = image if found is None else found[1] assert mod.MARK_V7 not in clean and mod.CLASS_HEAD not in clean @@ -535,6 +556,7 @@ def producer(ident): return text producers = {'pristine': clean, **{ident: producer(ident) for ident in helpers}} + hook_table_digests_check(legacy) # On the pinned source, every rebuilt legacy install must match its recorded digest: # this checks the installer's hook tables against history, not against themselves. on_pinned = hashlib.sha256(clean.encode()).hexdigest() == legacy.PINNED_CLEAN_SHA256 @@ -553,8 +575,6 @@ def producer(ident): assert hashlib.sha256(text.encode()).hexdigest() != legacy.SOURCE_DIGESTS[ident], f'{ident}: digest check is blind to a wrong hook' else: print('note: scheduler source is not the pinned image source; recorded legacy digests not compared') - if found is not None: - producers[found[0]] = image def run(text, temp): target = Path(temp) / 'scheduler.py' @@ -585,6 +605,15 @@ def run(text, temp): assert current == mod.apply_v7(clean, at=current_at), 'the installed v7 is not canonical' again, repeat = run(current, temp) assert again.returncode == 0 and repeat == current and 'already present' in again.stdout + if found is not None: + # A supplied legacy install (any accepted identity, including the pre-release v7 + # ones) migrates in place: expect the v7 helper at the legacy helper's position. + ident, found_clean, found_at = found + result, after = run(image, temp) + assert result.returncode == 0 and f'from {ident}' in result.stdout, (ident, result.stderr) + assert after == mod.apply_v7(found_clean, at=found_at), f'{ident}: migration did not keep the helper position' + again, repeat = run(after, temp) + assert again.returncode == 0 and repeat == after and 'already present' in again.stdout, ident # Composition with patch_adaptive_k, which inserts its class at the same anchor: # the helper may sit before it (decode-floor applied first, as deployed) or after @@ -603,6 +632,17 @@ def run(text, temp): for layout, text in (('helper-after-adaptive-k', helper_after), ('helper-before-adaptive-k', helper_before)): again, repeat = run(text, temp) assert again.returncode == 0 and repeat == text and 'already present' in again.stdout, layout + # A legacy install migrates in place in either layout: the v7 helper takes the + # legacy helper's position (public v5-main fixture, so no private source needed). + for layout, at in (('legacy-after-adaptive-k', None), ('legacy-before-adaptive-k', adaptive_at)): + legacy_src = (with_adaptive[:at] + helpers['v5-main'] + with_adaptive[at:] if at is not None + else with_adaptive.replace(mod.NEEDLE, helpers['v5-main'] + mod.NEEDLE, 1)) + for new, old, label in pairs['v5-main']: + legacy_src = mod.replace_once(legacy_src, old, new, label) + result, after = run(legacy_src, temp) + assert result.returncode == 0 and 'from v5-main' in result.stdout, (layout, result.stderr) + expected = mod.apply_v7(with_adaptive, at=at) if at is not None else mod.apply_v7(with_adaptive) + assert after == expected, f'{layout}: migration did not keep the helper position' v5 = producers['v5-priority'] refused = {