Repository navigation
perf(qwen-flash): 4096-token prefill chunks with an overlapped n-gram gather - #470
Conversation
|
I tested this locally against the PR's base using matched production builds on gfx1151. Vision functionality passes in AR and MTP, and large MTP contexts work through 238,884 tokens. The roughly 1.32 GiB increase in GPU allocations is confirmed and fits this 128 GB machine. However, I found a repeatable MTP decode slowdown after switching from text to vision, including with the default automatic cache budget:
Reproduction sequence: C1, context 262144, MTP with seven drafts, thinking off, greedy sampling, cache lookup disabled; send two identical 4,648-token text requests, then repeated 128×128 image/color requests, followed by arithmetic requests. ABBA comparisons produce identical outputs. I'd investigate and resolve this before merging. Correctness looks good, but the performance regression remains unexplained. All local validation results: correctness, performance, vision, large contexts and memoryEnvironment and method Candidate Performance controls use matched production binaries and workload histories. ABBA means base/PR/PR/base. HTTP slowdown flags use 5% plus 3 ms per request/phase; medians below summarize a shape without clearing individual flags. Functional suites record candidate correctness; they are not all matched timing qualifications. Builds and correctness
Actual draft execution was checked: AR has zero drafts; MTP functional cache/batch checks execute 708 drafts, separate text long-context checks 117, and vision suites 214 (142 accepted). Numerical-logit SHA256: Served text prefill and deep decode C1, context 262144, MTP/7 drafts, greedy, thinking off, fixed 1 GiB cache, uncached inputs, warmed execution shapes, ABBA.
All outputs match; no paired prefill phase is flagged. The gains partially reproduce the claim, with smaller gains on these public prompts.
The focused longer outputs match and execute 140 drafts each. They find no sustained text-only deep-decode regression; the original short-output flag remains retained. Standard AR benchmark Production, depth zero, warmed, three repetitions per point in ABBA. Values are means of the two round means per build; all output hashes match and actual draft counts are zero.
Concurrent serving C2, context 131072, MTP, ABBA: one 512-token code stream is already decoding when a second 26,618-token prefill arrives.
Small text prompts before vision C1/MTP, context 262144, fixed 1 GiB cache, ABBA, warmed, four measured repetitions per shape/build round. All generated messages match. Positive latency change means slower.
Arithmetic and tail each have one isolated decode flag: 71.51→80.89 ms and 76.64→87.35 ms. Those points remain inconclusive/unqualified. Initial image timings and focused transition controls Initial ABBA image measurements, fixed 1 GiB cache; messages match. Decode is flagged in both pairs for every image shape.
The two initial 1024-pixel prefill/TTFT flags are not reproduced in focused controls; retain them as inconclusive. The decode slowdown is reproduced. Focused ABBA, same 4,648-token preceding text history and fixed 1 GiB cache, every image shape warmed, followed by text controls:
Messages match throughout; the 128-token answers execute 122 drafts each. The roughly 9% longer code-answer decode shows that the finding also affects a substantial output. Additional ABBA using the default automatic cache budget:
One warmed PR image sample is fast; the slowdown is not asserted for every request. In a separate run that activates vision before large text prefill, PR decoding is normal (77.95 ms first image, 182.03 ms at 238K depth). That different history does not qualify the text-to-vision transition above. Root cause remains unisolated. Actually occupied large contexts C1, context limit 262144, MTP/mmproj, fixed 1 GiB cache; one matched base/PR sequence, not an ABBA speed qualification. All inputs are fully prefilled and all outputs match: text marker
Faster prefill does not clear the slower decode phases. Memory Global GTT+VRAM and process/system memory sampled every 200–250 ms. GPU allocations consume shared RAM on UMA; GPU and process RSS counters overlap and should not simply be added.
Cache-budget and setup limitations
For the last case, 11,845,269,388 retained bytes plus the requested 6,679,072,616-byte final checkpoint exceed the 12,278,530,048-byte budget. Zero-prefill replay is not claimed at that occupancy/budget. Raw commands, requests, timings, memory traces, corrected setup failures and outputs are retained locally. These checks do not reproduce the author's private prompts, second machine, SATA/network storage, every request-arrival history, 392–491-step numerical trace or separate d32K benchmark matrix. Vision checks establish functionality, not a new upstream encoder-accuracy qualification. No full model sweep was run. |
|
Thanks for the detailed numbers. I tried hard to reproduce this on both of my machines (Ryzen AI Max+ 395, 128 GB, Linux 7.0) and it does not show up. Each test ran the PR's base and this PR (code identical to Median server
Outputs are identical in every run. I also checked whether the PR's extra 1.32 GiB could cross a TTM page limit: with The one slow state I did see came from a test setup that left the n-gram reads short of page cache (a buffered loop device). There, decode slowed by 15–60 % on both builds, so decode is sensitive to n-gram read latency, but it did not separate the base from the PR. Since your vision-first run is normal, this looks like a state or ordering effect on your machine rather than a fixed cost. These would help narrow it down:
If it bisects to a commit, I'll fix it here before merge. |
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
5bddc52 to
5c6ebb4
Compare
|
Pushed fixes and refreshed the benchmarks and project README.
Numerical checks retain byte-identical logits; AR/MTP, batching, images, cancellation and cache replay checks pass. Small queue/snapshot timing flags remain inconclusive. |
|
@Nickybuk always a please to see your PRs! When is the next one coming? 😂 I run this locally and can confirm the performance improvement on large prompts. On small batches up to pp2048 I don't see any regressions with respect to main |
…-org#485) # Conflicts: # src/cli/serve/continuation_cache.cpp # src/cli/serve/http_server.cpp # src/cli/serve/openai_chat.cpp # src/models/qwen/chat_template.cpp # src/models/qwen38_flash_next/engine.cpp # src/models/qwen38_flash_next/kernels/rocm/executor.cpp # src/models/qwen38_flash_next/kernels/rocm/executor.hpp
This raises the Qwen3.8-Flash-Next prefill chunk from 2,048 to 4,096 tokens and keeps the larger chunk's n-gram gather off the critical path. Outputs stay byte-identical. On private repository prompts served by
gufo serve, each 24K–102K-token prompt's median prefill is 6–11% faster on two machines, and the gain grows with length. Prompts under ~8K tokens are not slower (−0.0% to +0.6% per prompt in dedicated rechecks on both machines), and decode is unchanged. Peak memory rises 1.3 GiB at 262K context.There are four code commits, plus one for the docs:
Why the chunk size changed outputs
The indexer budget (
indexer_top_k) is 2,048 tokens:So the same prompt got different logits depending on where the chunk boundaries fell. At a 2,048-token chunk no batch crosses the budget, so serving was unaffected. Any larger chunk changed outputs: main's
gpu_probe --batch 4096and--batch 8192dumps differ from--batch 2048on 4/4 repository documents.Fix.
Executor::Attentionsplits the launch at the budget. Pre-budget queries always take the dense tiles. The sparse part goes first, because it is the launch that can be refused.Why 4096 needs the n-gram changes
The PLE at layer 1 needs a gather of n-gram rows for every token in the chunk: 32,768 row lookups at 2,048 tokens and 65,536 at 4,096, served by random 4 KiB
O_DIRECTreads except for repeats and cache hits. Today only layer 0 overlaps that gather. At 4,096 the gather doubles while the overlap stays at one layer, so the host waits longer for every chunk. Each part of this PR addresses that in turn:min(32, cores)readers; this PR uses 128. Readers block inpread, so the core count is not a useful cap. A gather with fewer reads than readers, such as a decode step's, now wakes one reader per read instead of the whole pool.Prefill()takesmin(max_input_tokens, remaining, PrefillCapacity())per step. It now passes those tokens toSession::Sync, which passes them toForward. Once a chunk's own rows have arrived, the executor starts reading the next chunk's rows into a pageable buffer. The gather then overlaps a whole chunk's forward instead of one layer.The rows are a pure function of the n-gram history and the tokens. The next batch uses them only if both match exactly, from any session, and copies them into the one pinned buffer that uploads and captured graphs read. Any other gather first waits for a prefetch in flight and keeps its rows claimable, so a decode step interleaved by the scheduler does not discard them. A mismatch falls back to the normal read, so outputs cannot change, and nothing is prefetched past
max_batch.Correctness
Upstream
gpu_probedecode-mode logit dumps (--prompt-file,--decode-file,--dump), compared bysha256sum:--batch 2048--batch 2048--batch 4096--batch 8192gufo servewith MTP, greedy outputs for 3 cases × 2 reps are identical on both builds and both machines: one hash per case and rep.Speed: private repository prompts
gufo serve llm --context 262144 --sessions 1 --speculative mtp, 12 prompts from private repositories, 8 interleaved ABBA rounds (4 per build, each prompt twice per round), so 8 samples per prompt and build. Prefill tok/s as median [min–max] per prompt.Machine A
Machine B
Machine B's drive loses most of its random-read rate under sustained random reads, on every build. That is where the low minimums in its table come from, and the 7,765-token +15.1% is that noise: half of main's samples for that prompt were slow ones. The recheck below is the figure to use for short prompts.
Short prompts, rechecked
The five short prompts again, same builds and server flags, in a separate run with more samples. Each round sends every short prompt 4 times with no long prompt before it, then three pairs: the 102K-token prompt followed by the 8,447-token prompt. Interleaved ABBA, 4 rounds per build on each machine: 16 samples per short prompt and 12 per pair. Token counts differ from the table above by up to 20 because each request carries a unique prefix. Median prefill tok/s:
Short prompts are within −0.0% to +0.6% on both machines, including the 8,467-token prompt straight after the 102K-token one.
Speed:
gufo benchgufo bench -m … -p 4096 -n 32 -d 0,32768 -r 6, 8 interleaved ABBA rounds of 6 repetitions, 4 rounds per build. Mean of the round means, tok/s:Every round mean is within 0.5% of its build's mean. The round spreads over six repetitions are ± 11–17 tok/s at d0 and ± 0.7–2.3 at d32K on both builds. tg32 is given as the range of the round means.
A served decode benchmark (
--sessions 1 --context 262144, MTP, 4 interleaved ABBA rounds, 2 per build, 2 reps per case per round; mean tok/s):Each case has the same output hash on every build and machine.
Memory, threads and concurrent sessions
Peak GTT+VRAM sampled every second, plus threads and RSS of
gufo, over a fixed request sequence:--sessions 1 --context 262144, main--sessions 2 --context 131072, mainThe extra memory is the 4096-token scratch. The prefetch adds one pageable row buffer of at most 40 MiB, allocated on first use. The 96 extra threads are idle readers blocked on a condition variable. Over the 12-prompt run the peak differs less than this, and on B it is lower on this PR (106.1 vs 107.6 GiB), because the prompt snapshot budget is sized from the memory free at startup.
Concurrent sessions. While another request decodes, the scheduler bounds prefill steps by
--prefill-chunk(default 512), not bykPrefillChunkTokens, so other streams' gaps do not change. Measured with one session streaming a code answer while another prefills a 26.6K-token repository prompt (--sessions 2 --context 131072, ABBA, 2 reps per round):The typical worst gap is 0.67 s on both builds. One rep on B had a 1.14 s gap; the other three reps of this PR on B were 0.67–0.68 s.
A request that arrives during a long prefill. Any other n-gram gather first waits for a prefetch in flight. So a request that arrives while another session prefills in 4,096-token chunks can wait for that session's prefetch before its own first gather, and if the next step is not the predicted chunk, the prefetched rows go unused.
Who this helps, and what it does not change
kPrefillChunkTokenscan go back to 2048 without touching the rest. That combination was not measured.--prefill-chunk);Overlap with open PRs
This merges cleanly with main and with every open PR except #459, where both add a row to the same table in
EXPERIMENTS.md. #459 also changes the prefill GEMM kernels; if it lands first, I'll rebase and re-measure.Method
nix buildof the exact branch and of main 33d1b20, with the same flags.🤖 Generated with Claude Code