[Bugfix][GLM-5.3-Flash] Address kpool tail blocks by the padded indexer stride in the NVIDIA prefill seed kernel - #57477
Merged
ZJY0516 merged 2 commits intoSep 20, 2026
Conversation
…er stride in the NVIDIA prefill seed kernel The kpool tail cache aliases each indexer layer's cache: the tail tensor is placed at the indexer tensor's offset with the indexer's page as block stride (38016 B for GLM-5.3-Flash, block_size 1152), so tail block b occupies the first 2048 B of indexer block b. The NVIDIA `_kpool_tail_seed_kernel` still computed `base = (blk * 2 * KPOOL + t % KPOOL) * HEAD_DIM`, i.e. a dense 2048 B block stride. Every prefill therefore left its own tail block unseeded and wrote raw bf16 K/gate rows into indexer block `blk * 2048 // 38016`, corrupting fp8 keys and fp32 scales of an unrelated request. Blocks with low ids (the first prefix-cached prompt after startup, typically the serving system prompt) are hit by every request whose tail block id is below ~num_blocks / 18, and the damage persists for as long as the prefix stays cached. Address the tail through `tail.stride(0)` / `tail.stride(1)` like the decode kernel and the AMD seed kernel already do, and run the existing padded-stride regression test on every platform instead of ROCm only. Signed-off-by: Jared Wen <w13431838023@gmail.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
ZJY0516
marked this pull request as ready for review
September 18, 2026 03:51
ZJY0516
requested review from
AndreasKaratzas,
WoosukKwon,
mgoin,
tlrmchlsmth,
yewentao256 and
zyongye
as code owners
September 18, 2026 03:51
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ZJY0516
approved these changes
Sep 20, 2026
|
✅ @JaredforReal, CI is now available for this PR.
|
Member
|
/ci run |
|
✅ Triggered Buildkite CI #90031 for commit |
nick-oconnor
added a commit
to nick-oconnor/vllm
that referenced
this pull request
Sep 21, 2026
…kpool tail seed poisoning fix (vllm-project#57477) The 09-18 0.30 boot lost the full 1M context (auto-fit cut max_model_len to 516,096) and the branch carried a silent KV-cache corruption. Root causes, both now attributed and fixed: - Context: the indexer prefill gather workspace is sized in tokens while this indexer's KV is pool-granular, so it reserved max_model_len * 40 * 132 B = 5.16 GiB/GPU instead of 1.29 GiB. Upstream PR vllm-project#55222 divides by index_kpool at the call site; carried here. Consumed 82.27 -> 78.40 GiB/GPU, available KV 3.81 -> 7.68 GiB, full 1,048,576 context fits again (1,064,361 KV tokens, 1.02x). - Corruption: the NVIDIA _kpool_tail_seed_kernel addressed tail blocks densely while the tail tensor aliases the indexer tensor with the indexer's padded block stride (38016 vs 1024 elements here), so every prefill scribbled raw K / gate scores into unrelated indexer blocks and left the request's own tail unseeded. Fixed upstream by vllm-project#57477, in the new base. Also records that the 4.19 GiB CUDAGraph estimate is not the context regression and must not be disabled.
nick-oconnor
added a commit
to nick-oconnor/vllm
that referenced
this pull request
Sep 21, 2026
Add SM120-GLM53-FLASH.md documenting the rebuild onto upstream main (the fork retired once native GLM-5.3-Flash support landed upstream): the kept and dropped ocnr changes with their upstream status, the deployed production serve args (replacing the overlay-era max-num-seqs 10 + MTP config; verified by the c=1 bench and auto-fit boot logs, MTP overlay numbers kept as a variant note), and the boot-verify checklist that must pass before deploying a build. Updated through the 0.30 re-cut: records the 1M-context regression (the indexer prefill gather workspace sized in tokens; fixed by carried vllm-project#55222) and the silent kpool tail-seed KV-cache poisoning (fixed upstream by vllm-project#57477, in the base), and that the 4.19 GiB CUDAGraph estimate is not the context regression and must not be disabled. Signed-off-by: nick-oconnor <nick@ocnr.org>
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.
Purpose
Fix silent indexer-cache corruption for GLM-5.3-Flash on NVIDIA: the kpool prefill tail-seed kernel addresses tail blocks with a dense stride while the tail cache is stored with the indexer's padded block stride.
How the tail cache is laid out.
_get_kv_cache_groups_glm5_nextpads theKpoolTailSpecpage to the indexer page (page_size_padded = idx_page; 38016 B for the defaultblock_size=1152: 288 pool rows x 132 B), andget_kv_cache_config_from_groupsplaces every indexer layer's tail tensor at the same offset as that layer's indexer tensor withblock_stride = idx_page. Tail blockbis therefore the first 2048 B ([2, kpool=4, 128]bf16) of indexer blockb; the two never collide because both groups draw distinct block ids from one pool.create_kv_cache_viewsbuilds the view withas_strided, sotail_kv_cache.stride(0) == 19008elements instead of the dense 1024.The bug.
_kpool_tail_seed_kernelinvllm/models/glm5next/nvidia/ops/kpool_compress.pycomputedso for every prefill it (a) left the request's own tail block unseeded and (b) wrote 2048 B of raw bf16 K / gate rows at byte
blk * 2048of the layer's indexer region, i.e. into indexer blockblk * 2048 // 38016at offsetblk * 2048 % 38016— fp8 key rows or the fp32 scale region of whatever request owns that block. The decode kernel (TAIL_BLOCK_ELEMS/KPOOL_HEAD) and the AMD seed kernel already address the tail throughstride(0)/stride(1); only the NVIDIA seed kernel did not, and the existing regression test for exactly this layout wasskipif(not is_rocm).Why it shows up as "the cached system prompt is broken". Only indexer blocks with id below
num_blocks * 2048 / 38016(~5% of the pool) can be hit, and those are the ids handed out first after startup, i.e. the blocks of the first prefix-cached prompt, typically a long-lived system prompt. Every later request whose tail block id falls in[19, ~19 * k]overwrites part of that prompt's indexer blocks1..k; prefix-cached blocks are never recomputed, so the damage persists and accumulates, on every TP rank identically. The MLA KV itself is untouched (all writes stay inside the layer's indexer region), so the effect is a corrupted sparse top-k selection for any context longer thanindex_topk(2048). The predicted signature is therefore: with prefix caching on, requests that hit a long-lived cached prompt degrade over time while the same prompt sent with acache_salt(fresh, high block ids) stays correct even on repeated cache hits, and all TP ranks are affected identically. The reproduction below confirms exactly this.Fix. Pass
TAIL_BLOCK_ELEMS = tail.stride(0)andKPOOL_HEAD = tail.stride(1)to the NVIDIA seed kernel (same as the AMD kernel), assert the expected view layout inkpool_seed_tail_cache, and runtest_prefill_seed_honors_padded_tail_block_strideon every platform.Origin. The tail cache has aliased the indexer tensor with a padded, strided view since the tail spec was introduced on the GLM-5.3-Flash development branch (07-28); the prefill seed was a torch fancy-index scatter (
tail[block_ids, 0, offsets] = k, stride-correct) until it was replaced by this Triton kernel with dense addressing (08-17, "perf: indexer eager path tail seed and scatter"). The kernel reached main unchanged with #53906 and is unchanged in the current main, so every NVIDIA build since #53906 is affected.Not a duplicate. #56059 adds a
NUM_TAIL_BLOCKSout-of-range guard to the same kernel for a ROCm fault (#56037); it does not change the addressing, and the two changes compose. No other open PR touches the NVIDIA seed kernel's addressing.Test Plan
Kernel-level repro (SM103): build the runtime tail view with
create_kv_cache_views(KpoolTailSpec(block_size=4, num_kv_heads=2, head_size=128, dtype=bf16, page_size_padded=38016), ...)over a sentinel-filled 256 x 38016 B buffer, seed one 6-token request whose tail block is 200 (and 18), then inspect which bytes changed.End-to-end byte-level probe (GLM-5.3-Flash TP4 on 4 x GB300,
FLASHINFER_MLA_SPARSE,enforce_eager, prefix caching on,num_gpu_blocks_override=800so the block pool wraps every ~130 requests): an offlineLLMserves a 14.5k-token system prompt of 420 "vault -> access code" records, thencollective_rpcsnapshots the indexer K cache pages and the tail cache of the prompt's blocks on rank 0. Eight rounds of 50 unrelated filler prompts (200-2000 tokens each) follow; after every round one lookup question re-hits the cached prompt (keeping it hot in the LRU) and the cached pages are diffed against the snapshot taken right after the first request. Finally seven lookup questions are asked against the hot cached prompt. Same script, same seed, on main and on this branch.Request-level reproduction against a running server (OpenAI API only, no GPU introspection): same 14.5k-token registry prompt, 7 lookup questions, then batches of 100 unrelated filler prompts (150-1500 tokens, unique prefixes, concurrency 16) each followed by the 7 lookups (cache hit), then the same prompt with
cache_salttwice. Server: main, TP4,FLASHINFER_MLA_SPARSE,enforce_eager,--num-gpu-blocks-override 800so the block pool cycles quickly (with the default pool the same script needs several thousand fillers).Test Result
Kernel-level repro, tail view
stride = (19008, 512, 128, 1)elements:[409600, 411648)= indexer block 10, fp8 rows 230..245[36864, 38912)= indexer block 0 scale region + block 1 rows 0..6[7603200, 7605248)= block 200, bytes 0..2048[684288, 686336)= block 18, bytes 0..2048pytest tests/kernels/test_kpool_decode_update_batched.py: on main the un-skippedtest_prefill_seed_honors_padded_tail_block_stridefails (tail[3, 0, 2]is still the sentinel); with this PR 29 passed, 1 skipped (ROCm-only preshuffle test).End-to-end probe. Runtime views match the analysis: indexer
k_cache[7200, 32, 132]uint8 stride(4224, 132, 1)(800 blocks x 9 DeepGEMM pages) andtail_cache[800, 2, 4, 128]bf16 stride(19008, 512, 128, 1)share onedata_ptr. The system prompt occupies indexer blocks 1..13; the 38 fully cached pages of blocks 1..5 were watched.Final seven lookups against the hot cached prompt: main 1/7 correct, with garbled answers such as
vault harbor-harbor-7andraven-saffron-203? No, wait; this PR 7/7 correct. The modified regions on main are contiguous 2048 B runs (e.g. page 9 offsets 0..2816), i.e. one misplaced tail seed per filler request, and on main the very first request already writes 2048 B into indexer block 0 (the null block); neither happens with this PR. In a shorter server-based run on main (default pool of 9259 blocks, only 120 fillers, no wrap-around) the cached prompt was still answered correctly, which is consistent with the sparse per-request damage: each filler misplaces one 2048 B chunk, so a hot prompt degrades as the free list keeps cycling through the low block ids, while acache_saltcopy of the prompt (fresh, high block ids) is never hit.Request-level reproduction on main (correct lookups out of 7;
/metricsconfirmed 96768/101978 prompt tokens served from the prefix cache in every re-hit pass):cache_salt(fresh blocks)After the fillers the model no longer even reads the question correctly (
vault velvet-tundra-143forvelvet-tundra-67,they haven't specified which vault,the human turn is empty): the cached prompt is unusable while the salted copy of the very same prompt is fine.AI assistance (Claude Code) was used for the investigation and the patch; every changed line was reviewed and the tests above were run by the submitter.
🤖 Generated with Claude Code