Skip to content

fix(#178): complete chunked prefill — KV budget cap, FlashInfer workspace, ragged attention - #193

Open
drunkcoding wants to merge 3 commits into
fix/serving-streaming-and-flashinfer-dtypefrom
fix/chunked-prefill-kv-budget-and-workspace
Open

fix(#178): complete chunked prefill — KV budget cap, FlashInfer workspace, ragged attention#193
drunkcoding wants to merge 3 commits into
fix/serving-streaming-and-flashinfer-dtypefrom
fix/chunked-prefill-kv-budget-and-workspace

Conversation

@drunkcoding

Copy link
Copy Markdown
Contributor

Completes chunked prefill (#178). Stacks on #189 (fix/serving-streaming-and-flashinfer-dtype); merge #189 first.

Fixes

  1. KV block-budget reconciliation — the logical KV allocator and the physical paged store were budgeted independently, so the logical count could exceed the store; _bind_layered_paged_kv_store then failed its capacity gate and swallowed the error, leaving block_store unbound 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 bare except.
  2. FlashInfer workspace — chunked append-attention split-KV scratch overflowed the 128 MiB default on chunk 2; raised the shared prefill+decode workspace to 256 MiB.
  3. Ragged query packing — concurrent requests with heterogeneous query lengths padded every sequence to q_len, feeding padding tokens to the paged backend; now packs valid tokens per query_lengths and scatters results back.

Verification (Qwen3-30B-A3B, sm120, FlashInfer 0.6.18)

  • Chunked prefill runs end-to-end: short + 1000-tok (4 chunks) + 4096-tok (16 chunks) + decode.
  • Whole-prefill non-regressing (store binding now also covers it).
  • A/B (benchmarks/serving/chunked_prefill_latency.py) passed with output_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.

drunkcoding 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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant