fix(scheduler): hybrid radix hit corruption — GDN snapshot copy-on-donate + admission barriers - #287
Conversation
…onate + admission barriers Under the default radix prefix cache, hybrid (GDN) models intermittently serve corrupted answers on cache-HIT requests that are co-batched with in-flight decode: ~10% wrong under concurrent load in ground-truth probing (empty or garbled output, special-token leaks into content; with a multimodal tower, blank visual grounding). Cold prefill is always clean, and failures cluster in sticky per-branch windows that self-heal when the branch is re-donated. Root structure: chunk-commit and finish donate the request's snapshot slot into the radix tree by OWNERSHIP TRANSFER, so one linear-state slot id is shared between the tree and the request's still-in-flight pipeline. A late write through that alias poisons every subsequent hit of the branch. Stream-level wait_stream ordering does not close the window (full-device syncs at the donate/restore points shrink the failure rate from ~10% to ~3% but not to zero); severing the alias does. Fix: - copy-on-donate: the tree receives a private clone of the snapshot slot (written once behind the donate barrier, read-only for its tree lifetime); the request keeps its own slots -- no replacement alloc at chunk commit, both ping-pong slots and the live slot freed at finish. - full-device barriers before hit-admission restores and donation bookkeeping (request-level events, sub-millisecond; they also guarantee the clone's source is quiescent). Cost: one ~MB slot copy per donation and one device sync per prefill commit / hit admission; decode throughput is unchanged. Validated on a v0.1.2-based deployment serving GLM-5.3-Flash-NVFP4 (34 KDA layers on the same hybrid_radix/linear_state pool path as qwen3.5 GDN) on an RTX PRO 6000 Blackwell: a deterministic hammer that alternates identical prompts with abandoned or saturating concurrent streams reproduced the corruption at 3/30 (images) and 2/24 (pure text) before the fix, and runs 0/60 and 48/48-clean after; an overnight ground-truth soak (~2,200 probes: exact-match arithmetic, needle retrieval, multi-turn recall, media A/B alternation with per-request content forensics) measured the pre-fix hit corruption at ~10% and confirmed payloads and cache keys correct on every failure, isolating the fault to the hybrid hit path.
|
Correction after deeper forensics — the corruption evidence in this PR's description does not implicate upstream. Slot-level checksum instrumentation overturned the ownership-aliasing theory. The true root cause of every corrupted hit we observed was in our downstream model integration: our KDA attention op (GLM-5.3 support, #270) never implemented the ×64-boundary track-snapshot write that the hybrid-radix design expects ( With the writer implemented on our side, restore-side checksums show real boundary state on every hit and 90 adversarial hammer rounds run clean — with or without the changes in this PR. What remains of this PR on its own merits, with no corruption claim attached:
Happy to re-scope the PR to the guard + optional copy-on-donate, or close it — maintainers' call. Apologies for the initial misattribution; the ground-truth probes and checksum forensics that led here are documented in the linked repository. |
|
Ran this on a Qwen3.8-Flash-Next TP=2 deployment (36 GDN layers on the same
So the corruption did not reproduce here with this probe, and the copy-on-donate clone plus the barriers cost nothing measurable on decode or TTFT. Two differences from your setup that may matter for reproducing: my hit requests are short (one-shot answers, 16 tokens) so the donated slot is rarely still in flight, and the probe alternates only one cached prompt. If you can share the hammer script I can run it on this box, since a second machine confirming the fix would help it land. Setup: 2x RTX 6000 Ada (48 GiB, sm_89, PCIe Gen4, no NVLink), 2x Xeon Gold 6526Y, 503 GiB RAM, CUDA 13.3, torch 2.11+cu130, sgl_kernel 0.4.5; Update 2026-09-05, merged into the deploy branch. Tried on 2 x RTX 6000 Ada (sm_89) serving Qwen3.8-Flash-Next (RadixArk NVFP4) at TP=2, offload backend, fp8 KV pool of 8 x 262,144 tokens, merged onto my deploy branch (main af71ba4 + #385/#386/#389/#392/#354 and ten other open PRs), tests run on the box, then put in production. Merged with conflicts against my image-aware cache key in scheduler/cache.py (three donate sites, mechanical). Two things needed on top:
On this box the corruption does not reproduce with or without the PR: 0 mismatches in 50 hit requests (control, sequential aborts, 4-way concurrent aborts, 5k-token conversations) on a GDN + vision deployment. Server-side decode is unchanged (444-450 tok/s at 8 running, same as without the PR); one GDN slot here is 55 MiB, so the clone per donation is a 55 MiB device copy. Kept in production. |
…orruption -- GDN snapshot copy-on-donate + admission barriers Upstream FlashML-org#287 at abc13da, merged onto deploy/chatdnp for the PR sweep. Conflicts in scheduler/cache.py: the three donate sites take the PR's private clone, keyed on _key_ids(req) (the image-aware cache ids of the deploy branch) instead of req.input_ids. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
…ollows ForwardOutput (FlashML-org#224) The donate and hit-admission barriers of PR FlashML-org#287 call torch.cuda.synchronize on the cache manager's device; the scheduler unit tests run that path on a CPU device, so guard the call on device.type. The drained-forward fake in test_abort_inflight_prefill built a bare tuple, which PR FlashML-org#224's scheduler no longer unpacks: build a ForwardOutput instead. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
… contract The tree now holds a private clone of the GDN snapshot and the request keeps its own slots, so the two tests that asserted ownership transfer (tree slot == request slot, ping-pong slot replaced) assert the clone contract instead. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt
…ge costs Two things made a screenshot expensive out of all proportion to what it says. **An image prompt could not share a prefix at all.** ``CacheManager`` matched image requests against the empty prefix and never inserted them, on both sides and on purpose: every image expands to the same placeholder id, so a match keyed on the raw ids could hand one image's KV to another. Correct, but total -- one screenshot anywhere in an agent conversation made every later turn re-prefill the whole context. Measured here at 163k tokens: ~40 s before the first token, every turn, for as long as the image stays in scope. The key can be fixed instead of abandoned. ``tokenize`` now derives ``cache_ids``: ``input_ids`` with the i-th placeholder run replaced by ids from that image's blake2b digest, above every vocabulary (>= 2**30) and carrying the position within the run. The cache manager keys on it through ``_key_ids`` and the two exclusion branches go away. Which run belongs to which image is read off the prompt -- placeholder runs come out in image order -- with a single combined hash as the fallback when the counts disagree; that is coarser but never wrong. Verified on both families that have a tower here: Qwen3-VL and gemma-4 keep runs separate for 1, 2 and 3 images, adjacent or not. **A 4K screenshot cost 8,184 soft tokens.** Qwen3-VL is native-resolution and the checkpoint's processor only downscales above 16.7 Mpx, i.e. never; at 32x32 pixels per soft token a 4K frame lands just past ``--max-prefill-length``, where a multimodal prompt is terminal (it cannot be split). The processor's resize bounds are now clamped on the way in, ~1 Mpx by default, ``FREETOKEN_IMAGE_MAX_PIXELS`` to raise it. gemma-4 is untouched by this: ``Gemma4ImageProcessor`` has no ``size`` and caps itself at 280 soft tokens. One interaction needed a decision. A hit that ends INSIDE the image span is unusable -- the model scatters ``mm_embeds`` against the image tokens of one forward, so the run has to sit entirely on one side of the boundary -- and it happens on every prompt ending with its image, because the match always leaves the last token out. ``match_req`` caps the key at the span start there. Without it the adder finds a cut span it cannot move and declines transiently forever: nothing about the next pass changes the cache, so the request never runs and the scheduler spins. Measured on Ornith-1.5-35B-A3B-NVFP4, one 4090, bf16 KV 262,144: 4K image 8,184 -> 967 prompt tokens, answer unchanged 16,140-token prompt + image 3.2 s cold -> 1.2 s repeated (cached 16,128) short prompt + image cached 960 / 965 on the repeat different image, same text cached 0, "Green." vs "Blue." -- no false hit gemma-4, 4K 264 soft tokens before and after the clamp tests/tokenizer 65, tests/scheduler 88, tests/kvcache/radix 142 -- all passed, unchanged from this branch without the commit. Cold system-test runs against the same branch without it: the five text cases are byte-identical, the two image cases differ in wording and agree in answer, which is the clamp changing the resolution the tower sees. Two cold runs of this build are byte-identical to each other. It does not need FlashML-org#287. Assisted-by: Claude Opus 5
…ike text 037f102 narrowed the rule from "the whole prompt in one chunk" to "the image span in one chunk", which is what a 196-token sprite in a 166k-token turn needs. The span is [first image token, last+1) because ``mm_embeds`` is one concatenated tensor scattered in one forward -- so it grows with the TEXT between two screenshots, not just with the pictures. An agent conversation reaches the limit by talking: 400 prompt with images needs 10392 contiguous tokens in one prefill chunk (the image tokens span [160334, 170726) and cannot be split) Nothing configurable moves that. Cheaper images (~490 tokens each after the clamp) only buy more turns before the gap between the first and last one exceeds a chunk, and raising --max-prefill-length OOMs long before it helps: a 32k chunk's activations do not fit beside a 5 GiB KV pool on a 24 GiB card (measured -- it took the worker down twice today). So the concatenated tensor stops being scattered whole. ``_merge_multimodal`` takes the rows belonging to the placeholders inside ITS OWN forward -- the ones an earlier chunk or a prefix-cache hit already consumed sit in front of the window -- and the adder chunks an image prompt exactly like a text one. ``Req.mm_scatter`` and the whole pull-back / reject path go away with it, ~90 lines. Both families that carry a tower here are converted; the approach is gdevenyi's, from FlashML-org#386 (28fd56d). The span cap 09ea814 put in ``match_req`` goes too. It existed because a hit landing inside a placeholder run left half the run cached and half to forward, which the all-in-one-forward scatter could not represent; the window skips the cached half instead. Without the cap a prompt that ends with its image keeps its prefix -- 20,800 of 20,840 tokens on the repeat here, 6.0 s -> 1.2 s, and a different image at the same position still misses (answered "Green" where the cached one answers "Blue"). Measured on Ornith-1.5-35B-A3B-NVFP4, one 4090, --max-prefill-length left at its 8192 default: 2 images with 9k of text between span ~19k 10,186 tokens, 3.2 s (was a 400) 6 images with 9k between each span ~50k 55,360 tokens, 18.1 s (was a 400) A(blue) 9k B(green), and reversed "Blue, Green" / "green blue" -- read across the boundary, in order tests/tokenizer 58, tests/scheduler 90, tests/kvcache/radix 142: all passed. Twelve tests pinning the removed rule are gone and three cover the window (a span wider than a chunk now admits; a chunk scatters only its own rows; a chunk holding no placeholder scatters nothing). The ``_NoSwa`` stub gained the ``page_size`` the reservation math has been reading, which is what had six of these failing on this branch already. A cold system-test run is character-identical to the same branch without this commit, all seven cases. Assisted-by: Claude Opus 5 Re-verified on this branch (no FlashML-org#337/FlashML-org#354/FlashML-org#287 under it): tests/tokenizer 58, tests/scheduler 88, tests/kvcache/radix 142 all passed; a cold system-test run is character-identical to the same change on the daily branch, all seven cases; the two shapes that used to 400 (spans of ~19k and ~50k tokens) answer at the 8192 default.
2 participants
Problem
Under the default
--cache-type radix, hybrid (GDN) models intermittently serve corrupted answers on cache-HIT requests that are co-batched with in-flight decode. Measured with deterministic ground-truth probes on a production deployment: ~10% of hit requests wrong under concurrent load — empty or garbled output, special-token leaks into content (<|assistant|>mid-text), and with a multimodal tower, answers describing a blank image. Cold prefill is always clean. Failures cluster in sticky per-branch windows (a poisoned branch keeps failing for a while, then self-heals when re-donated).Every hybrid GDN model on the
hybrid_radix/linear_statepool path is exposed; nothing about the trigger is model-specific.Root structure
_cache_req_hybriddonates the request's snapshot slot into the radix tree by ownership transfer (chunk commit donates the frozen ping-pong slot; finish donates the live slot "zero-copy"). That leaves one linear-state slot id shared between the tree and the request's still-in-flight pipeline. A late write through any retained alias poisons every subsequent hit of that branch — which is exactly the observed sticky-window signature.Stream-order alone does not close the window: adding full-device syncs at the donate and restore points shrank the failure rate from ~10% to ~3% but not to zero. Severing the alias does.
Fix
_clone_slot_for_tree), written once behind the donate barrier and read-only for its tree lifetime. The request keeps its own slots — no replacement alloc at chunk commit; both ping-pong slots and the live slot are freed at finish._restore_linear_states) and donation bookkeeping (_cache_req_hybrid): request-level events, sub-millisecond, and they guarantee the clone's source is quiescent.Cost: one ~MB slot copy per donation plus one device sync per prefill commit / hit admission. Decode throughput unchanged in before/after measurement.
Validation
On a v0.1.2-based deployment serving GLM-5.3-Flash-NVFP4 (34 KDA layers riding the same
hybrid_radix/linear_statepath as qwen3.5 GDN; support PR #270) on an RTX PRO 6000 Blackwell 96 GB:The hammer scripts are small and self-contained; happy to attach them or port them into
tests/if useful (they pair naturally with the regression-harness proposal in #281).