Repository navigation
perf(offload): port the slow-tier admission and lookup-skip gates to main - #2283
Closed
PerryZhang01 wants to merge 9 commits into
Closed
PerryZhang01 wants to merge 9 commits into
PerryZhang01 wants to merge 9 commits into
Conversation
Saving is unconditional, which is only right when the slow tier has customers. It does not always have them. With a KV pool large enough to hold the working set, vLLM's own HBM prefix cache answers nearly everything and the slow tier returns almost nothing -- 0.03% of the tokens queried, on MiniMax-M3 at 46647 blocks -- while the full save tax is still paid, costing 10.4% throughput against the same build with offload off. Shrink the pool to 16612 and that build gains 21.7%. The difference is not the implementation; it is whether anyone ever reads what was written. Nothing new has to be measured to tell those apart. The load path already reaches that verdict per request, from two numbers it has in hand: how much of the prompt HBM already held, and how much LMCache claims. Fold the verdict into an estimate, and let the estimate decide what fraction of saves to admit. The estimate lives on the worker rather than the scheduler because they are different processes -- the scheduler's own verdict, in the EngineCore process, cannot reach the save it would gate. The worker sees the same two numbers, so the loop closes locally. The loop is circular: saving less can only lower the observed payoff. A floor is what stops that being a ratchet to zero -- it bounds the worst-case tax still paid and keeps enough content in the tier for a workload that starts cold and later develops reuse to climb back to saving everything. Admission is deterministic rather than sampled, so a short benchmark window measures the same rate as a long one. Off unless asked for (OFFLOAD_SAVE_ADMISSION=1): it is a heuristic over a workload property, and the operating points it has been measured at are two. Only the dense family gates on it; DSV4 records the payoff but never acts on it, because its save owns PAGE/SLOT sidecar and snapshot cleanup whose contract cannot be validated on this hardware. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit 5744ea8e4b35fea4fd3cca8c7e3d95b576084488)
…orker The admission estimate was built on the worker, from the load path's own ``lmc <= hbm`` test, and at the operating point admission exists for it sees almost nothing: when HBM satisfies a request the scheduler declines the load and never emits a worker task at all. A 46647-block calibration run made that plain -- 8084 stores, 4 retrieves, and not one admission report, because the estimate never left its warm-up. The policy would have stayed dormant exactly where it was supposed to act, and worse, any sample it did collect was drawn only from loads that survived the scheduler, so the payoff would read near 1 when the tier was useless. The verdict is reached in the EngineCore process and the save that acts on it runs in the worker, so it has to travel. It travels as a tally on the step's metadata, which already crosses that boundary: only the ratio matters, so two counters are enough and no per-request field is needed. Folding a step's tally applies the decay in closed form, which treats the step as a uniform rate -- the order of a few verdicts inside one step is arbitrary anyway -- and converges on the rate the scheduler reports. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit ea54548b76d9950c29ef772d58be13293353875b)
…ires The tally sat in build_connector_meta, past two earlier exits, and a 46647-block probe showed it never ran: 226 lookups, 193 declined loads, zero verdicts recorded. The load is abandoned at whichever gate first proves it pointless, and with a large KV pool that is never the last one. Count at the gates instead. ``need <= 0`` in get_num_new_matched_tokens means LMCache held the prefix and HBM held it too, so no LoadSpec is ever built. Requests that get past it are still judged where they were. A lookup miss is deliberately not counted: a novel prefix says nothing about whether saving pays, and counting it would bias a cold workload into throttling itself. The same probe also priced the verdict rate -- about three per second at concurrency 25 -- which the estimate's defaults were wrong for. At the old alpha the estimate needed several hundred verdicts to come down from its optimistic start, so a ten-minute window would have been spent almost entirely climbing. 1/16 converges in about forty. What the probe showed about the workload is worth recording: at 46647 blocks every single declined load was `too_small`, LMCache holding exactly one chunk more than HBM already had. The tier is not merely under-read there; it is never more than 128 tokens ahead. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit 1abd06680fc0f3812eab917764a0b5988dda00e5)
The tally reset itself on every build_connector_meta, but not every metadata object reaches the worker -- the plugin forwards only the steps that carry inner metadata. Every verdict on a step that was not forwarded was simply lost. Measured on M3 at 46647 blocks, the scheduler reached about 193 verdicts in a minute and the workers saw 16 to 31 of them: roughly seven in eight went missing, which left the estimate sitting near its optimistic start and the policy dormant. Cumulative counters are immune. The worker subtracts the totals it saw last time, so a dropped step costs a little latency and nothing else, and a scheduler that restarts its counters resynchronises instead of folding a negative delta. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit eee12e54fea534729e9b8154ea2cf5f514312666)
…site _decide_load_after_alloc has four callers -- two in the scheduler mixin's parking and free paths, two in build_connector_meta -- and the tally sat in build_connector_meta. On M3 at 46647 blocks that saw seven verdicts where the run produced about two hundred: the parking path reaches the decision first for almost every request, and the loop that was instrumented never ran for them. Counting inside the decision is the only placement that cannot miss a caller. With this the estimate behaves as designed. A ninety-second window at 46647 blocks: 199 verdicts, none of them a hit the slow tier could serve, payoff converged to 0.0, and 135 of 269 saves declined. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> (cherry picked from commit a107fca2c50c30d870f96cdf3ce76ebe49bc9432)
… and time the ones that do The scheduler calls `get_num_new_matched_tokens` for every waiting request whose `num_computed_tokens` is still zero, so a request that loses the budget race is re-examined on every step until it is admitted. Only the first examination reaches a lookup -- the rest read the memoised result -- but the prompt copy sat above that branch and so was paid every time. Worse, it was paid twice: slicing a `ConstantList` already yields a list, and the `list()` around it copied that again. Measured at 68k tokens: 0.229 ms per call, of which 0.168 ms was the redundant second copy. Timing is added in the same place because save admission cannot reach this call. Admission throttles bytes written; the lookup hashes the whole prompt in the EngineCore process and then blocks on a round trip to every lookup rank, and a request pays that whether or not anything is ever saved. At 46647 blocks the slow tier paid off 4 times in 2305 verdicts and admission recovered only 28% of the loss, so the remainder has to be somewhere admission does not reach; this is the first candidate, and it is now measurable instead of assumed. `perf_counter` is nanoseconds against a round trip measured in hundreds of microseconds, so the tally stays on by default. Reporting cadence is `OFFLOAD_LOOKUP_REPORT_EVERY` (default 256). Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit cc010972127d90615177ed020e441ed03e0f87b1)
…hange A lookup blocks the scheduler on a ZMQ round trip to every worker rank while those ranks are running the model. Measured on M3 at 46647 blocks: 4.4 ms per call early in a run, 14-17 ms once prompts reach 100k tokens, 22.4 s across an 853 s window -- 2.6% of wall clock. The hashing the round trip nominally pays for is 0.4 ms of that, so the cost is queueing behind the forward pass for the GIL, and the lever is how many lookups are issued rather than how big they are. Most of them cannot matter. A load reaches the worker only if `_decide_load_after_alloc` clears it, and that wants at least `_min_load_tokens` past the chunk boundary above the HBM frontier. The slow tier can never offer more than the prompt, and a full-prompt hit gives back a token so the last one is recomputed, so `num_prompt - 1 - hbm_cached` bounds what any answer is worth. Below the floor the answer is irrelevant: the request prefills from HBM plus recompute either way. On the vLLM plugin path that bound is available before the lookup -- `num_computed_tokens` is the prefix-cache frontier and the plugin pushes it onto the seq first -- and it only tightens as the request runs, since a frontier moves forward. At 46647 blocks HBM answers 92.8% of tokens and the slow tier 0.03%, so nearly every lookup there is already decided. On ATOM's native scheduler the frontier is not known yet (block_manager sets it after), so the bound degenerates to the prompt length and this fires only for prompts too short to load at all -- correct, just rarely useful. A skipped request keeps the save floor a hit would have left it. Those tokens came out of the HBM prefix cache, so the request never computed them and has nothing new to write there; dropping the floor instead would make it re-save its whole prompt, which is the regression `_clear_pending_load` already documents at 22% of bytes written D2H. Off by default, like save admission: it trades a cache hit the floor would have refused for the round trip, which is right only once the floor is trusted. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> (cherry picked from commit db9910576edabc2cdc520ccf6ae89afb27a848f5)
`OFFLOAD_LOOKUP_SKIP_UNLOADABLE` skips a lookup when HBM already covers all
but a sub-floor tail, so no answer could clear `_decide_load_after_alloc`.
The bound is exact and monotone-safe. The problem is the side effect: the
skip also leaves the save floor at the HBM frontier, asserting that LMCache
already holds that prefix. Nothing checked, because nothing asked.
Whether that assertion is harmless depends on the pool. Measured on M3-MXFP8,
TP=4, talos radix c25, strict eval window:
46647 blocks skip +6.1 points the tier answers nearly every lookup
and HBM held the same tokens anyway
16612 blocks skip -3.8 points the tier is really serving; the
suppressed write is a hit lost later
So the sign is a runtime property, not something an environment variable can
be set to in advance. This gates the skip on what the tier is actually worth:
the share of issued lookups whose answer clears the load floor
(`hit - hbm_cached >= _min_load_tokens`), as an EMA. Below the cut the skip
fires; above it the round trip is paid.
Not the admission payoff, which would be the obvious input and is biased by
the skip itself -- the skip removes exactly the requests most likely to be
judged unusable (`slow_tier_unused` 1779 -> 226 while `aligned_large_hit`
holds at 1293 -> 1313). Gating an action on a metric that action inflates is
positive feedback. The actionable rate is fed only by lookups that ran.
Two details are load-bearing, both of them measured rather than reasoned:
* The cut is 1%, not a comfortable 10%. A cold 16612 tier needs ~1536 probes
to become useful and reads 0-17% throughout that fill -- indistinguishable
from 46647's steady state at a 10% cut. A gate there opens during the fill
and starves the tier it was filling: 1280 skips, -1.91% against the
never-skip control. What does separate them is an order of magnitude:
46647 returns 4 actionable answers in 2048 probes, 16612 trickles 1.6-7%
from its second block. The margin is not generous -- the lowest in-fill
EMA observed was 0.0181, 1.8x the cut.
* An open gate must keep sampling. Every skipped request is one the estimator
never hears about, so a gate fed only by issued lookups can never learn the
tier became useful. One skippable request in 32 pays for its lookup anyway.
Visible in the logs: at each skip milestone the probing arm has issued more
lookups than the arm without it (12/18/26 against an arithmetic 8/17/25).
One setting, opposite and correct behaviour at both pool sizes:
s46ad2 3,179,591 tpm/GPU -0.59% vs the skip arm, +5.50% vs no-skip
s16ad2 2,267,176 tpm/GPU -0.26% vs the no-skip arm, +2.58% vs skip
Arms of identical behaviour at 46647 span 1.23% across a few hours on this
host, so "lands on" is resolved to about a point; the effect being reproduced
is 4-5x that.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
(cherry picked from commit 8fba0d6e9f7e463b028bf158a775cf219a60ed9e)
The eight commits before this one add two optional gates, eight environment
knobs and two default-on log records to a subsystem whose README is its
contract -- 1300 lines with an env table, a diagnostics list, a test index
and a glossary. None of it was documented. A migration that lands code into
a documented subsystem without entering the document is not finished.
The word needed disambiguating before anything else could be written. This
README already used "admission" for two unrelated things, both predating
this branch: bounded DSV4 staging-row admission, and the worker save
admission that `OFFLOAD_MAX_PENDING_SAVES` bounds. Slow-tier admission is a
third, and differs in kind -- the first two bound concurrency and always
run, while this one declines work outright and is off by default. The
glossary now says so rather than leaving a reader to infer which is meant.
Placed where the evidence already lives: the gates read the reload verdict
table in "When Does a Reload Actually Happen?", so the mechanism section
sits directly under it. Both regimes are stated with their measurements,
M3's two pool sizes and GLM-5.2's four on/off cells, because the sign of
both switches is a workload property and a reader setting them from the
table alone would have no way to know that.
Three things are documented because they mislead otherwise:
- `save_share` is the realized share `saved / (saved + skipped)`, so it
reads 1.0 by construction whenever nothing was refused. It is not the
share the gate computed and carries no information about the estimate.
Read `payoff`.
- The payoff dips far below target during ramp -- 4e-05 on GLM-5.2's
first report against a 0.10 target, settling near 0.9 -- so a threshold
between the transient and the steady state throttles during ramp.
- `warmup`, `alpha` and the probe cadence are defined in event counts, so
their wall-clock constants drift with the verdict rate.
Two new diagnostics entries cover the shapes this workload actually
produced: admission enabled with `skipped` stuck at 0, and the skip enabled
but structurally unreachable. Both look like a switch that failed to take
and neither is.
Documentation only; no code or behavior changes.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Contributor
Author
|
Superseded by #2282 — folding the fix and the optimization into a single PR as requested. All nine commits from this branch were cherry-picked onto #2282 unchanged (no conflicts; the two parts touch disjoint files), and #2282's description now covers both. Closing here to avoid a duplicate; the branch is left in place. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this is
Nine commits from the MiniMax-M3 offload line, rebased onto current
main.They add two optional, default-off gates to the LMCache offload path, plus
the telemetry that decides whether either is worth turning on, plus an
offload/README.mdsection documenting both.Depends on nothing outstanding — #2231 is already merged, and this applies
cleanly on top of
1c4caef48.Why
Saves and lookups on the offload path are unconditional today. That is the
right behaviour when the slow tier has customers, and pure cost when it does
not. Nothing in the connector currently measures which regime it is in, so
there is also no way to answer the question for a given workload without
guessing.
The first thing these commits add is therefore not a gate but a verdict
tally: every reload decision is already classified (
aligned_large_hit,slow_tier_unused,too_small, …), and[OFFLOAD-ADMISSION]now reports therunning payoff whether or not either gate is enabled. That part is free and
always on, and it is what makes the two gates decidable rather than a matter of
taste.
The two gates
OFFLOAD_SAVE_ADMISSION=1— treat each reload verdict as one sample(
aligned_large_hit= the tier paid off, anything else = it did not) and run anEMA over them. At or above
OFFLOAD_ADMISSION_TARGET(0.10) every save isadmitted; below it, the admitted share falls linearly to
OFFLOAD_ADMISSION_FLOOR(0.25), carried by a credit accumulator so therealized share tracks the computed one without randomness.
OFFLOAD_LOOKUP_SKIP_UNLOADABLE=1— skip a lookup whose answer cannotchange the outcome. When even a total hit (
num_prompt - 1) would leave theloadable gap
num_prompt - 1 - hbmbelowOFFLOAD_MIN_LOAD_TOKENS(8192), theverdict is
too_smallno matter what the tier says, so the round trip is purecost.
Skipping is not free, and the commits are explicit about the two ways it
bites:
holds this prefix — an assertion nothing checked, because nothing asked.
Dropping the floor instead does not mean "write the prefix once", it means
rewriting the whole prompt every request: measured, stored bandwidth
0.272 → 2.448 GB/s and throughput −44.7%.
exactly the ones HBM already covers, so what survives over-reports payoff:
at 16612 blocks
slow_tier_unusedfalls 1779 → 226 whilealigned_large_hitholds at 1293 → 1313, roughly doubling the estimate the controller sees. The
skip therefore records the verdict it provably would have reached
(
note_slow_tier_verdict(False)) instead of dropping the sample, stands downwhenever the tier is measurably load-bearing (actionable-rate EMA ≥ 1% after
64 lookups), and keeps one forced probe every 32 skippable lookups so the
estimate stays fed after the gate closes.
What is measured, and where it does nothing
On MiniMax-M3 the regime decides the sign. At 46647 blocks the tier answered
0.03% of queried tokens and the skip won 6.1% throughput; at 16612 blocks the
same switch lost 3.8 points, because there it starved a tier that was being
read.
On GLM-5.2 both gates are identity transforms, and I could not make either of
them do anything. Concurrency 16, 600 s, pool 64, the full 2×2 plus an
interleaved repeat of the admission cell:
is ~26.9k tokens against an 8192 floor, so no lookup is ever skippable.
OFFLOAD-LOOKUP-SKIPis zero in every cell.(356 / 512 / 522 admitted,
skipped=0in all three).span 2.7% end to end, so anything under ±2.7% here is "not measurable", not
"no effect".
The same code is also live (default-off) in a GLM-5.2 AgentX plugin-backend run
at concurrency 32 finished today, where the tier supplied 0.0529% of prompt
tokens and neither gate triggered either.
So this PR does not claim a throughput win. It claims: the gates are correct,
they are off by default, they are documented together with the regime that makes
each one right or wrong, and the telemetry that tells you which regime you are in
is now always available. On the one workload where they help, they help; on
GLM-5.2 they correctly do nothing.
Two cautions that are in the README, not just here
warmup,alphaand the probe cadence are defined in event counts, sotheir wall-clock constants drift with the verdict rate — a change elsewhere
that changes the rate silently retunes them.
before settling near 0.9), so a threshold sitting between the transient and
the steady state throttles during ramp rather than at the working point.
One thing not to misread in the telemetry:
save_shareis the realized sharesaved/(saved+skipped), soskipped=0pins it to 1.0. It is not independentevidence that the gate computed 1.0.
Overlap with in-flight PRs — please note the merge order
Related, and independent of this one: #2282 fixes the LMCache world-size
collapse on the plugin path. That fix is what makes the slow tier readable at
all there; the gates in this PR only decide whether a save is worth doing.
Neither depends on the other and they touch different files.
"Admission" already meant two things in this subsystem before this branch, and
this adds a third. The glossary entry in the README commit spells all three out:
staging-row admission bounds DSV4 temp rows; worker save admission is the
OFFLOAD_MAX_PENDING_SAVESrunning-plus-queued bound; slow-tier admission(
OFFLOAD_SAVE_ADMISSION, this PR) decides whether a save is worth doing at all.The first two bound concurrency and always run; the third declines work outright
and is off by default.
That matters here because #2201 (
fix(offload): bound Dense/M3 saves with safe retirement) and #2246 (fix(offload): bound DSV4 save admission) are open onthe second meaning and touch three of the same files —
atom/kv_transfer/offload/_offload_common.py,atom/kv_transfer/offload/README.mdandtests/test_dense_offload_connector.py— and #2201 addsatom/kv_transfer/offload/dense/save_admission.py, which is not this gate.The changes are independent in behaviour but will conflict textually. I am happy
to rebase this one behind either of them; say which order you want.
Tests
tests/test_slow_tier_admission.py(new, 255 lines) and additions totests/test_dense_offload_connector.py(+233) cover the EMA and creditaccumulator, the skip's precondition and stand-down, the forced probe cadence,
the verdict-recording on the skip path, and the save-floor behaviour.
.github/scripts/run_unit_tests.shwas run onorigin/main(1c4caef48) and onthis branch, in the same container. The failure sets were compared test-by-test
from the JUnit reports, not by count: identical — the same 28 failures and
the same 4 collection errors on both, with no test newly failing and none newly
passing. main: 28 failed / 7015 passed / 103 skipped / 4 errors. This branch:
28 failed / 7047 passed / 103 skipped / 4 errors — the +32 are the new tests in
tests/test_slow_tier_admission.pyandtests/test_dense_offload_connector.py.The pre-existing 28/4 are main's own and are untouched by this PR.
🤖 Generated with Claude Code