fix(#178): complete chunked prefill — KV budget cap, FlashInfer workspace, ragged attention - #193
Open
drunkcoding wants to merge 3 commits into
Conversation
added 3 commits
September 2, 2026 18:13
…city The logical KV allocator (memory_manager) and the physical paged store (memory_coordinator) were budgeted independently, so the logical block count could exceed the physical store. _bind_layered_paged_kv_store then failed its capacity gate and swallowed the error, leaving block_store unbound and crashing _schedule_chunked_prefill with 'layered paged KV store is not bound'. Cap the logical count to the physical store before binding (PagedKVCache.resize_num_blocks), guard per-sequence decode capacity (ensure_sequence_capacity), and surface the real reason via logging instead of a bare except.
Chunked append-attention split-KV scratch (tmp_s/LSE/partial-O) grows with query rows and KV pages and overflowed the 128 MiB default on chunk 2. Raise the shared prefill+decode workspace to 256 MiB.
Concurrent requests with heterogeneous query lengths padded every sequence to q_len, feeding padding tokens to the paged backend. Pack valid tokens per query_lengths before the backend and scatter results back to the padded layout.
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
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.
Completes chunked prefill (#178). Stacks on #189 (
fix/serving-streaming-and-flashinfer-dtype); merge #189 first.Fixes
_bind_layered_paged_kv_storethen failed its capacity gate and swallowed the error, leavingblock_storeunbound and crashing_schedule_chunked_prefill("layered paged KV store is not bound"). Now caps logical→physical before binding (PagedKVCache.resize_num_blocks), guards per-sequence decode capacity, and logs the real reason instead of a bareexcept.q_len, feeding padding tokens to the paged backend; now packs valid tokens perquery_lengthsand scatters results back.Verification (Qwen3-30B-A3B, sm120, FlashInfer 0.6.18)
benchmarks/serving/chunked_prefill_latency.py) passed withoutput_parity=True(178_ab.json).Known pre-existing limitation (out of scope)
Whole-prefill hangs on prompts >
max_tokens_per_step(2048) — the exact case chunked prefill solves; tracked separately.