Repository navigation
[3/3] kv: prefix caching under tensor parallelism - #281
Gaurav-Shah05 wants to merge 11 commits into
Conversation
Under TP every rank has to skip the same tokens, and a length can only be agreed before the step is scheduled, so the match has to be settled by the time the request is ingested rather than when its inputs are prepared. The cost is that two requests sharing a prefix, both ingested before either commits a page, now both miss.
The two tests that build their condition inside the probe installed it after the ingest, so with the probe moved they would have passed while checking nothing; both now patch first, and the hashing one fails if it hashed nothing at all.
A follower's index is its own — a generated page is keyed on each rank's postprocess path, whose timing against the next commit is that rank's — so the only length every rank can skip is one rank 0 works out from all of them. This carries the ranks' answers up; nothing reads them yet.
… step Rank 0 cannot take a minimum over answers it has not received, and a step that ran before them would skip what one rank alone matched, leaving the ranks a different number of tokens into the same request. There is no timeout: a reply that never comes is a dead rank, which every other lockstep path would hang on too, so the wait only says so after five seconds.
…tes to it One span runs on every rank, so the length they skip is the smallest any of them matched. The leader works it out on the first schedule that names the request and sends it with that batch; past there it lives on each rank's own stream, which is what leaves a retry after a refused admit skipping what the first attempt did.
Every rank now opens its own index and offers what it matched; rank 0's minimum is what they all skip, so the reason the cache stayed shut above one rank is gone. Ingest no longer settles the length above one rank: there it is this rank's offer, and `agree_prefix` is what settles it.
… set The set only ever holds requests that are still alive, so dropping each one at remove keeps it as small as the live set; it had borrowed the no-speculation set's bound, which named nothing about it.
merceod
left a comment
merceod
left a comment
There was a problem hiding this comment.
Ran this on two H200s with orpheus_tp2.yaml and it hangs requests. In a 40 request test three of them never got scheduled and remained there until the client's 600 sec timeout, and a rerun with a few log lines added gave 5. Three things, then the same issues from #280.
-
The leader gets one NewRequest per partition it serves, and in orpheus_tp2.yaml rank 0 also runs snac_decoder, so _add_new_request runs twice for the same rid. The follower answers the first one. If its reply lands between the two, _record_prefix_match pops the wait and the second _await_prefix_match arms it again, and nobody is left to answer. Doing it only on the first NewRequest for the rid would fix it (check per_request_info before add_request), or skip it when the replies are already complete. Either way it must not re-arm once _agree_prefix has consumed the replies, since the scheduler checks tp_prefix_waiting on every step, not just the first.
-
The lease is taken and pinned at ingest per rank, before the group has agreed. The indexes drift (the description says so), so rank 0 might pin 6 pages for a request while rank 1 pins 4. Until the first schedule trims them, any other request's _alloc in that window has two fewer evictable pages on rank 0 and can fail there while it succeeds on rank 1. _handle_allocation_failure says the allocator state has to be symmetric or the lockstep breaks, and post_warmup_validate exists for the same reason. I have a two-manager unit repro where the same step is refused on one rank and admitted on the other. It needs a full pool, so it's an overload-time failure, but overload is where main degrades gracefully today and this would hang instead. Two ways out that I can see: don't retain before agreement (probe at ingest, report the length, take the lease at agreement and re-check it), or agree as soon as the last reply enters _record_prefix_match and push it to the followers right away, so the window is one hop rather than "until the first schedule". The first one is the safe one I think.
-
In _trim_lease, agreed is set from the length handed in even when this rank's lease is shorter. A rank told a length above its own offer would skip pages it never had. Clamping keep to len(lease or ()) makes the rest safe by construction (I think). The position-zero keying, the leaves heap and the index buffer size from your PR #280 are present here unchanged. The refused-admit one does not since the probe now runs at ingest.
One more thing to be aware of: sentence-length Orpheus prompts are under a page, so every request in my runs logged "matched 0 of its 0 pages" on both ranks. The TP hit part itself only got tested at unit level.
|
Best to close this PR for now. The ranks stay in step because every allocator reaches the same decision from the same inputs. That holds while each allocated page belongs to one live request, since every rank then holds the same number of pages. The index holds pages no live request needs, and the ranks can end up holding different ones, so the description's "fit is total minus live" doesn't survive a cache. I tried three designs of my own (carrying the index on the schedule, reducing the matched length across the group, subtree eviction), then both of yours on two managers. Option 1 fixes the window before agreement. After agreement the ranks still differ, 6 free or evictable pages against 7, and a 7-page admit goes through on one rank only. Option 2 gives the same result. Every path ends up needing the two indexes to be identical, which per-rank allocators can't guarantee. What's missing is admission control, which M* has never had. Reserving the whole prompt used to hide that, and the cache takes the reservation away. It's the stall you hit on #280, and it's why the ranks disagree here. Filed #308 with the cases and numbers. @merceod let me know what you think. |
What does this PR do?
Part of #210, stacked on #280. Under tensor parallelism every rank has to skip the same number of tokens, so this PR makes the ranks agree and lifts the refusal #280 put in at
world_size > 1.Indexes are per rank and not symmetric
Each rank keys its generated pages from its own postprocess host copy, and when that copy lands relative to the next commit is the rank's own business, so one rank can hold a page another has not keyed yet. Allocation stays symmetric because a request fits or not based on total minus live, which is the same on every rank. That settles the thread's open question: the indexes differ, so the ranks need a minimum.
Ranks agree on a length, never on pages
Every rank probes at ingest. Followers send their matched length to rank 0 (
WorkerMessageType.TP_PREFIX_MATCH, the only message that goes follower to leader). Rank 0 takes the minimum, puts it on the first schedule that names the request asScheduleTPNode.prefix_matched, and every rank cuts its lease down to it before the engine stage. A request still waiting on a reply is not scheduled. Atworld_size == 1nothing waits.Cost
The probe moves from the engine stage to
ingest_requestfor every deployment. Two same-prefix requests ingested before either has committed a page now both miss, where a step-time probe would have let the second one hit. The benchmarks will give more insight into whether that matters.How was it tested?
32 new tests.
pytest test/modular -qatb625da99gives1 failed, 1132 passed, 18 skipped.Orpheus at tp_size = 2 on two H100s: 22 requests, 44 per-rank log lines, and both ranks admitted the same length every time. I could not compare outputs on against off, because Orpheus at tp2 gives different outputs even within one process; this PR only changes how many pages a rank skips, never what a page holds. Turn-two TTFT medians came out at 208 ms cached, 204 ms uncached and 206 ms on main, inside the run-to-run spread.
Checklist
ruff check .passes