Qwen 3.8 long context: paged KV with SSD tier, Quest sparse decode, MTP composition - #20
Merged
Conversation
…t metadata layout
…nder full budget kvPagedPolicy=on hands full-attn KV to KVPageStore; decode runs the paged split-KV kernel over sinks+recent+top-k pages with lag-one Quest scores; page seals compute min/max metadata on the token CB. Prefill writes stream into the identity-mapped pool, existing kernels unchanged. MTP defers to plain decode in paged mode (v1). E2E parity: paged==dense greedy streams.
…I flags Pool smaller than context: sealed pages spill (write-behind, layer-major sparse file) and evict under LRU; decode fetches selection misses with the live selection pinned so its own fetches cannot evict it. Beyond-RAM chat appends run the blocked prefill: chunk KV scatters into free slots, the sealed past streams via one sequential pread per window through a 2-stage ring overlapped with flash-update dispatches (FP32 carry state), tail pages fold causally from the pool. In-chunk min/max metadata avoids a refetch storm at first decode. CLI: --kv-paged on/off/auto, --kv-topk, --kv-pool-pages; auto pool sizing from RAM.
Shared Qwen38PagedKVRuntime (extracted from the runner) drives one cursor, pin set, and score state for both plain tokens and spec rounds. The verify pass scatters draft rows through the page store, runs per-position paged attention over the round's pinned selection (table extended across span page-crossings), and scores sealed pages with the always-committed bonus query. Cursor rewind on rejected drafts un-seals pages (KVPageStore.rewind; stale spills unreachable). canRunRound now requires position >= 1 — a round at 0 has no prior hidden to seed the drafter and anchored RoPE at -1. Auto pool headroom 20->22 GiB: measured 2x decode loss from a 4 GiB pool squeezing weights on the 24 GiB host. Byte-identity of spec vs plain decode preserved under paging (Qwen38PagedMTPTests); full suite 1078 green.
|
You have reached your Codex usage limits for security reviews. Please try again later. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b33191abde
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Merged
2 tasks
A speculative round reuses one page table across its verify rows where plain paged decode reselects per token, and accepted rows are emitted without their own Quest score pass — under a sparse selection the speculative stream could drift from the plain paged stream (Codex P1 on PR #20). Rounds now run only while the selection through the round's span, plus one position of margin, provably covers the entire context; past that point decode falls back to plain paged tokens. The margin guarantees the token feeding the first sparse selection's lag-one scores is always plain-decoded, so byte-identity holds across the crossover. KVPageSelector.coversEntireContext is conservative about score staleness: gap pages are picked by top-k only when scored, so any unscored just-sealed page must sit inside the recent window. Tests: selector coverage boundaries (agrees with select()), a rounds-stop-at-the-boundary gate test (fails without the gate: rounds kept running 83 -> 100 in the sparse tail), and a sparse-budget e2e byte-identity run across the crossover. The suite's MTP-off reference runners now disable MTP per-instance instead of setenv, which raced runners constructed concurrently by other suites.
The write-behind pwrite pair asserted success from the spill queue, so ENOSPC, a short write, or an I/O error terminated the whole process (Codex P2 on PR #20). Writes now run through a complete-write loop (resuming short writes and EINTR); the first failure is recorded and thrown by the next spill-file read — fetch or readSpilledSpan — so the foreground operation fails cleanly instead of reading garbage. reset() clears the recorded failure along with the rest of the state, since every page is rewritten before it can be read again. The error box is deliberately separate from the store: spill closures must not retain the store itself, because its deinit synchronizes on the spill queue and dropping the last reference there would deadlock.
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
1 participant
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.
Qwen 3.8 long context: paged KV cache with SSD tier, Quest sparse decode, and MTP composition
Qwen 3.8-27B's full-attention KV cache costs 64 KiB/token — 16 GiB at the model's full 262,144 context, which neither fits beside 14 GiB of weights on a 24 GB machine nor could be read per-token even if it did (dense attention at 262k = 16 GiB of reads ≈ 181 ms/token from RAM). This PR makes long contexts practical at full FP16 KV precision (no quantization) by treating capacity and bandwidth as separate problems:
Controls
Correctness
PagedAttentionParityTests).Qwen38PagedKVParityTests).Qwen38BlockedPrefillTests).Qwen38PagedMTPTests).Measured (M5 MacBook Pro 24 GB, real 27B checkpoint)
--kv-paged off† Measured before the exactness gate (see review updates below): at the default budget the needle's context exceeds the exhaustive-selection window, so MTP now hands those decodes to plain paged tokens (≈ the 5.9 tok/s plain rate). Raise
--kv-topkpast the context length to keep speculative rounds running exactly.Auto-pool sizing was tuned during the sweep: a 4 GiB pool beside 14 GiB of weights pushed the 24 GB host into memory compression (~2x decode loss); the default now leaves 22 GiB of headroom (2 GiB pool ≈ 32k resident tokens/layer) and full-262k settings decode at full speed.
Test plan
Review updates (2026-08-26)
Two commits addressing the Codex review:
KVPageSelector.coversEntireContext, conservative about unscored just-sealed pages); past that point decode falls back to plain paged tokens. Within the gated regime the round's single page table equals the per-token selection, and the margin guarantees the token feeding the first sparse selection's lag-one scores is always plain-decoded — so byte-identity with MTP-off paged decode holds everywhere, including across the crossover. New tests: coverage-boundary units, a rounds-stop-at-the-boundary test (fails 83→100 rounds without the gate), and a sparse-budget e2e byte-identity run.KVPageStoreError.ioFailedinstead of trapping the process from the background queue.reset()clears the recorded failure. The error box deliberately does not retain the store (spill closures retaining it would deadlock deinit's queue sync).