Repository navigation
Mamba align: retire state blocks across null gaps (backport vllm #55450), peak pool pressure of long prefills halved - #101
Conversation
…null-gaps hunk (the file now matches the object the tests ran against)
…ve under MTP/EAGLE, Mamba align null-gap retirement) The v0.29.0 tag has neither upstream fix (#52771/#52807 for the offload tier, #55450 for align-mode Mamba retirement), so a 0.29 pin that left them out would drop what syv-ai#100 and syv-ai#101 add. Both ride here as patch files exported from the fork's commits (cpuchip/vllm qwen38/0.29, ab6843f55 and 708d18c64): the offload patch applies to 0.29 unchanged, the Mamba patch's field init and free-path pop are re-anchored around 0.29's _num_checkpoint_blocks. syv-ai#100's serve oracle comes along as bench/replay_offload_serve.py; the launcher hunk from syv-ai#100 was already carried. Ledger rows and docs/vllm-0.29.md updated.
|
Merging. The backport is correct, the mechanism is real, and I reproduced the win on the box — but on your configuration, not on this box's production one, so I'm recording that split here rather than in the PR title. What I checked before touching the GPU: the base remover in vllm/v1/core/single_type_kv_cache_manager.py ends its backward walk at GPU, one 3090 at 250 W, production stopped, file snapshotted and restored bit-identically (md5 matched both ends). Two paired boots per arm, pool verified constant from the launcher log in each; sampling Your configuration (SPEC=mtp CTX=long PREFIX_CACHE=1), pool 176,020 tokens in both arms, peak baseline -> patched: Counter, same arms: baseline grows one retained state block per ~1664 computed tokens (4,5,6,...,16 at computed=0..19968), per-request maxima 16/12/11/5. Patched: every one of the 12 requests max=5. That 16 -> 5 bound is the cleanest statement of what your patch does, and I'd rather the PR body said that than "roughly halved" — at the prompt lengths a 24 GiB card can reach, the pool reduction is -40.5% on the one row with no cached prefix and -21 to -23% where a large cached prefix dominates the pool. I could not get near your 480k / 71-block regime. Production configuration here (SPEC=dflash2 CTX=fast PREFIX_CACHE=1 DFLASH_TOKENS=15 INT8_ACT=int8 PREFILL_ATTN=int8), pool 57,669 tokens in both arms: every peak bit-identical to four decimals across 20 requests, cached_tokens identical, generated text identical on all of them, zero engine errors. The counter says why — retained state maxes at 16 with and without the patch in CTX=fast, so there is no null-gap leak to fix on that path. Since PREFIX_CACHE=1 is what turns align on, this code goes live in production the moment it merges; my answer to your implicit "is that safe" is yes on the evidence I have, and the honest caveat is that "no measurable harm over 20 requests" is not "harmless under every workload". Answering my own earlier asks so they don't sit open:
Two things I did not settle, neither blocking: whether Follow-up I'd like from you in a separate PR, not a condition of this merge: a line in docs/gotchas.md next to the align/prefix-cache note recording the null-gap retirement, the 16 -> 5 bound, that the win is mtp/long-shaped, and that the patch file carries single-space blank context lines a whitespace trim will destroy. Good find — the null-gap/high-water-mark reasoning was right, the dependency ordering was right, and the diff did exactly what the description said it did. |
#101 added the patch file and #54 added patches/series within twenty seconds of each other, so neither PR's checks saw the other's change and main landed with the two disagreeing. It applies to the same file as mamba-align-checkpoint-order.patch and must follow it, which the last line satisfies. Series check is green again at 32 patches, 0 offsets.
docs/pr-standard.md is the maintainer's bar, read off his review comments on syv-ai#90, syv-ai#91, syv-ai#93, syv-ai#99, syv-ai#100, syv-ai#101 and syv-ai#54: verified on the box; title, body, lead commit and diff say one thing; what is unproven stated; promised docs in the same PR; sibling launchers swept; a series line for every patch; knobs registered in envs.py; benches guard their hardware; no claim wider than its measurement; credit and asks specific. scripts/pr-gate/pr_gate.py checks what text can check and prints the rest; reply_gate.py does the same for a public comment; retracted.txt is the list of phrases we have withdrawn. Run against the three held PRs it reports every hold he placed and one more of the same class (syv-ai#90 reads VLLM_SPEC_ATTN_DEBUG unregistered).
…ment beside item 44 Two follow-ups asked for in review. syv-ai#91's merge comment: item 12 covered spec-decode-attn.patch and its query-cap trap and said nothing about the large-pool limit; it now records the int32 cast, that a pool over about 2,383 blocks was silently reading the wrong memory before it, the synthetic test, and that the int8 prefill kernel had the same pattern. syv-ai#101's merge comment: item 44 gains the null-gap retirement, the 16 to 5 retained-block bound measured on the reference 3090, that the win is mtp/long shaped and inert on the dflash2 fast profile, and the single-space blank context lines in that patch file that a whitespace trim destroys.
…ment beside item 44 Two follow-ups asked for in review. syv-ai#91's merge comment: item 12 covered spec-decode-attn.patch and its query-cap trap and said nothing about the large-pool limit; it now records the int32 cast, that a pool over about 2,383 blocks was silently reading the wrong memory before it, the synthetic test, and that the int8 prefill kernel had the same pattern. syv-ai#101's merge comment: item 44 gains the null-gap retirement, the 16 to 5 retained-block bound measured on the reference 3090, that the win is mtp/long shaped and inert on the dflash2 fast profile, and the single-space blank context lines in that patch file that a whitespace trim destroys.
…ve under MTP/EAGLE, Mamba align null-gap retirement) The v0.29.0 tag has neither upstream fix (#52771/#52807 for the offload tier, #55450 for align-mode Mamba retirement), so a 0.29 pin that left them out would drop what syv-ai#100 and syv-ai#101 add. Both ride here as patch files exported from the fork's commits (cpuchip/vllm qwen38/0.29, ab6843f55 and 708d18c64): the offload patch applies to 0.29 unchanged, the Mamba patch's field init and free-path pop are re-anchored around 0.29's _num_checkpoint_blocks. syv-ai#100's serve oracle comes along as bench/replay_offload_serve.py; the launcher hunk from syv-ai#100 was already carried. Ledger rows and docs/vllm-0.29.md updated.
… fork; the syv-ai#91 and prefill int64 casts Rebased over syv-ai#54 (patches/series), syv-ai#91, syv-ai#99 and syv-ai#101. patches/series lists the 0.29 set: the two files 0.29.0 retired (vllm-pr54282-draft-gumbel-salt, xgrammar-spec-terminated) leave the list and offload-mtp-serve.patch (syv-ai#100) joins it after mamba-align-retire-null-gaps. The Dockerfile, README loop, verify.sh and check_vllm_series.sh apply in series order, at --fuzz 0 as before, and the int8 knob check asks the live registry. Every patch file is re-exported from cpuchip/vllm qwen38/0.29 at c68ae347f. Content changes: the split-KV verify kernel casts the block id to int64 before the stride multiply (syv-ai#86/syv-ai#91, the same three lines the 0.28 series has since e79596b), and the int8 prefill kernel casts at its three block-table loads (the syv-ai#86 pattern in a second kernel; the 0.28 fix is syv-ai#109). triton-prefill-attn-int8.patch is prefill-attn-int8.patch here as on main. Three files (dflash2-prewarm, dflash2-z-adaptive-emitted, prefill-attn-int8) carried the whole diff a second time as stray lines above the first file header, a leftover of the import from the old patch files; the fork commit bodies are prose only now and the preambles follow.
…eries) and applied at --fuzz 0 The two kvarn/ files were the last hand-cut patches in the tree: cut before the syv-ai#100 and syv-ai#101 backports, they landed after the series by fuzz (four hunks) and only the tree identity check could tell. The fork branch now carries the three KVarN commits last, in the order kvarn/install.sh applies them, and the files are exported from those commits, so they apply to the series-patched tree with exact context. install.sh applies at --fuzz 0 and stops on a rejected hunk instead of relying on the port(kvarn-v2) marker count, which on 0.29 covers one file of seven.
…ment beside item 44 Two follow-ups asked for in review. #91's merge comment: item 12 covered spec-decode-attn.patch and its query-cap trap and said nothing about the large-pool limit; it now records the int32 cast, that a pool over about 2,383 blocks was silently reading the wrong memory before it, the synthetic test, and that the int8 prefill kernel had the same pattern. #101's merge comment: item 44 gains the null-gap retirement, the 16 to 5 retained-block bound measured on the reference 3090, that the win is mtp/long shaped and inert on the dflash2 fast profile, and the single-space blank context lines in that patch file that a whitespace trim destroys.
…ecipe Adopt the KVarN subsystem upstream keeps in kvarn/ (dropped with the WSL2 lane at the fork), on request: - patches/kvarn-0.28.0.patch and kvarn-v2-runner-0.28.0.patch (verbatim from upstream kvarn/) plus kvarn-files-0.28.0.patch (create-only port of the eight kvarn/files/vllm modules), so the whole subsystem flows through the existing patch machinery -- patch_vllm.py and the Dockerfile loop -- with no separate install step. Build order: after int4-kv-per-token-head (the only shared file) and before the mamba-align patches; verified in a Linux venv. - patches/mamba-align-retire-null-gaps.patch (upstream PR syv-ai#101 / vllm#55450), which the 0.28.0 sync left pending a 3090 A/B against our VLLM_MAMBA_ALIGN_KEEP_CHECKPOINTS=1 default: adopted with the KVarN lane, which the v2-runner patch was cut against. - recipes/w4a16-k4v2-dflash2.sh: kvarn_k4v2_g128 KV, dflash2 drafter, TP=2. The upstream single-card dflash2+huge profile adapted: --block-size 128 and --prefix-match-unit 128 (tile = block = hash unit), 2 seqs, 245760 max-len, 2048 batched tokens, KVARN_POOL_MEM_FRAC=0.15, VLLM_SPEC_DECODE_ATTN=0; deviations disclosed in the header. - README/AGENTS: the seventh recipe, the 34-patch set, the k4v2 invariants, KVarN out of the dropped-lane list. Validated: docker build (pristine Linux vllm 0.28.0 venv, 34/34 patches in alphabetical order + compileall), in-image KVarN wiring check with the port(kvarn-v2) marker audit, and the patch_vllm.py reset->apply->stamp cycle in the container. Provenance: upstream bae2023.
Backport of vllm-project/vllm#55450 (merged 2026-09-11, not in 0.29.0) onto the 0.28.0 tree, on top of
mamba-align-checkpoint-order.patch.What it fixes. In
--mamba-cache-mode align, a long prefill leaves null gaps between the Mamba state snapshots that await retirement. The base block remover walks a request's blocks backward and stops at the first null block, so every state older than the gap stays allocated until the request ends.MambaManagernow overrides the remover for align mode: skip nulls instead of stopping, and remember how far retirement has reached so the prefix is not rescanned. States come back at request end either way; the cost is peak pool pressure during long prefills, not a permanent leak, which is why an idle server shows nothing.Unit evidence, both boxes. Upstream's
tests/v1/core/test_single_type_kv_cache_manager.pyat v0.28.0 with the PR's test additions, against this tree's manager (fork series applied): before,test_mamba_retirement_crosses_null_gapsfails and the 3584-token bounds tests fail with71 == 7 + 1and71 == 7 + 2, upstream's own numbers; after, 7 of 7 new tests and 18 of 18 in the file. Identical on a WSL2 4090 (fork image) and a native 3090 venv.Real-path evidence.
SPEC=mtp CTX=long PREFIX_CACHE=1, one long prompt per row,vllm:kv_cache_usage_percsampled every 0.25 s during the prefill, pool verified identical between arms before comparing percentages.WSL2 4090,
MAX_LEN=65536, default prefill attention (noPREFILL_ATTN=int8), pool 148,501 tokens both arms as computed by the launcher; so the two tables share the long profile's KV dtype, drafter and a 65,536 ceiling, and differ in int8 prefill attention, the card, and the pool:Native 3090,
MAX_LEN=65536(so the long profile's 150k ceiling was not in play),PREFILL_ATTN=int8inherited from that box's production in both arms, pool 161,051 tokens both arms as computed by the launcher forCTX=long(not a chosen pin). Both arms share all three, so the A/B is internally valid; the absolute percentages are not comparable to the 4090 rows above and would not be the numbers at a true 150k context:Settle is 0.0% on every arm. The un-retired states cost about one extra pool token per prompt token, and the ratio widens with prefill length (0.60 to 0.50 of the shipped peak across that ladder), so on a pinned single-card pool this is the difference between a 45k-token prompt fitting alongside other work or not. Prefix reuse is unaffected: the same 512..32k reuse ladder is byte-identical between arms on both boxes.
One note for anyone editing the patch file: it contains blank context lines that are a single space; a trailing-whitespace trim turns it into a malformed patch that
git apply --checkrejects.