Repository navigation
[Engine] Opt-in routed prompt stream: skip streamed experts no token routes to (follow-up read time -18%, same output) - #4283
Open
bitaccesscomau wants to merge 1 commit into
Conversation
…reamed experts no token of the chunk routes to; test A prompt chunk of 1,024 tokens or more streams every non-resident expert of every layer through the ring, and the issuer thread enqueues each copy as soon as the ring has room. The ring is longer than a layer's entries, so a layer's copies are all on the copy stream before its router has run: an expert no token routes to is copied anyway. With the switch on, for a streamed chunk on one GPU with the issuer thread: - the issuer keeps at most STRATA_PREFILL_ROUTED_STREAM_DEPTH copies (default 16) queued ahead of the link: it waits until the copy issued that many copies earlier has landed (cudaEventQuery on its slot's copied event). The ring's own rule is unchanged, so the stream still runs as far ahead of the compute as before; - after each layer's host grouping the compute thread publishes which experts at least one token routes to; - the issuer drops an entry it has not issued yet whose layer is published and whose expert has no token. Entries already issued stay. An entry staged through the Stager's pinned buffers is never dropped (its job queue is in order). Every routed entry is still copied, in the same order, into the same ring slot (k % ring), with the same slot waits and copied event; the consumer's walk is unchanged and never reads an unrouted entry's slot, and the copy stream stays in order, so a group's single wait on its last copy still covers the group. The prompt's products are the same bits. After the read the borrowed slots are refilled from RAM blob by blob into the same slots, independent of what the ring held. On main at the release defaults (RTX 3090, PCIe 3.0 x16, 250 W; the same binary with the switch off and on, 4 pairs): a 1,790-token follow-up read on an 18.8K-token conversation took 17.9% less time, fresh reads of 2K, 4K and 32K tokens ran 20%, 8% and 5% faster, and the streamed experts fell by 23%, 16% and 4%. In a deterministic lane the replies are byte-identical to main at 1K to 32K with the switch off and on, and the switch dropped 42% of the streamed experts at a 1,046-token read and 3% at 31K tokens. Unset, nothing new runs: no gate, no publication, and the default issuer loop is not edited. The fused path, the peer prompt ring, the inline issuer (STRATA_PREFILL_ISSUER=0) and the staged walk are untouched. With STRATA_TRACE=1 each chunk prints the ring entries it copied and dropped. The issuer walk and the route gate are header-only (include/strata/prefill/routed_stream.hpp); routed_stream_test runs the walk itself against a simulated in-order link and a consumer that walks the layers as prefill.cpp does (every routed entry issued into its own slot and read with its own bytes, only unrouted entries of published layers dropped, never more than the depth queued, never past the ring), and on three threads with real atomics.
This was referenced Oct 27, 2026
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.
Problem and result
A prompt chunk of 1,024 tokens or more streams every non-resident expert of every layer through the pinned prompt ring
(
src/prefill/prefill.cpp), and the issuer thread enqueues each copy as soon as the ring has room. The ring holds moreentries than a layer has, so a layer's copies are all on the copy stream before the layer's router has run: an expert
that no token of the chunk routes to is copied over the PCIe link anyway. On an RTX 3090 with the Swift IQ3_XXS pack,
the switch below skipped 6,379 of the 15,308 streamed experts of a 1,046-token read (42%), 3,347 of 15,481 on a
1,979-token follow-up turn (22%) and 2,054 of 67,372 at 31K tokens (3%).
With
STRATA_PREFILL_ROUTED_STREAM=1(off by default, read once):STRATA_PREFILL_ROUTED_STREAM_DEPTHcopies (default 16, clamped to 1 to 256) queued ahead ofthe link, so the copies are not committed far ahead of it; the ring's own rule is unchanged, so the stream still runs
as far ahead of the compute as it does today;
already issued stay, and an entry staged through the Stager's pinned buffers is never skipped (its job queue is in
order).
Every routed entry is still copied, in the same order, into the same ring slot (
k % ring), with the same slot wait andthe same
copiedevent, and the consumer's walk is unchanged (it never reads an unrouted entry's slot), so the prompt'sproducts are the same bits. With the variable unset the engine runs today's code: no gate, no publication, and the
default issuer loop is not edited.
Measured on
mainon one RTX 3090 at the release defaults, the same binary with the switch off and on, 4 pairs: theread of a follow-up turn (1,790 new tokens after 18,763 reused) took 17.9% less time, 3,214 to 2,640 ms (pair ratios
1.224, 1.212, 1.209, 1.225); fresh prompt reads ran 19.9% faster at 1,979 tokens, 8.2% at 3,963, 2.1% at 6,011 and 5.3%
at 32,105 (8,192-token chunks), and the conversation's 18,768-token first turn took 6.0% less time. Decode after the
counting prompts (the same greedy text in every arm) moved by -0.8% to +0.5%; the agent turns decoded 3.5% slower, on
texts that differ in every arm, with 60.6% of their drafts accepted against 62.3% (every row is in the table below).
In a deterministic lane the replies, the prompt and reuse counts, the engine's
session-state fingerprint and the prompt-path residuals are byte-identical to
mainfrom 1K to 32K tokens with theswitch unset and set. An earlier 4-pair run on the same machine with other settings and another base gave a follow-up
read 17.7% shorter.
Related discussion
Scope proposal: #4282. Issues and pull requests, open and closed, were searched again for duplicates on 2026-10-27,
including those filed since 2026-10-26 (
STRATA_PREFILL_ROUTED_STREAM,ROUTED_STREAM,unrouted,routed-only,stream_all,streamed experts,experts_streamed,route gate,no token routes,prompt ring,issuer,copy queue depth,read-ahead experts,skip experts prefill,expert copies routing,PCIe expert copies prompt):none skips unrouted entries of the streamed walk.
mainase258e6ec):STRATA_PREFILL_RING=8stages only the routed experts of a large chunk, but gives upthe streamed ring's read-ahead (a ring of 8 slots). This keeps the full ring and its read-ahead and skips only what the
routing shows is unused.
copiedevents and deduplicated slot waits, Windows HIP): the same issuer loop. Two textualconflicts in
prefill.cpp(the start of the issuer thread and the trace after its join; keep both) and one semanticone: with [Engine] prefill: opt-in sparse copied events and deduplicated slot waits (Windows HIP prompt reads +36% at 4K, +15% at 32K) #4005's sparse events on, the consumer waits on counters (
a_rec,rec_of) that only [Engine] prefill: opt-in sparse copied events and deduplicated slot waits (Windows HIP prompt reads +36% at 4K, +15% at 32K) #4005's loop writes,and the routed walk would not advance them. Whichever merges second needs one rule: the routed walk off while sparse
events are on (one more condition in
rs_on), or the walk recording through [Engine] prefill: opt-in sparse copied events and deduplicated slot waits (Windows HIP prompt reads +36% at 4K, +15% at 32K) #4005'srecord. I will rebase ontowhichever lands first.
STRATA_PF_PEER_CPU): both insert a few lines right after the host grouping's count loop; keep both.This switch is off when a peer prompt ring is in use.
CMakeLists.txtand the paragraph indocs/DETAILS.mdsit at the same place;keep both.
stream): the Stager's entries, which this switch never skips.
STRATA_PREFILL_PIPE_FLOOR, the smallest pipelined prompt chunk (default 512, as before) #3821 (the pipefloor): other levers on the same reads.
Merge-tested against all 517 open pull requests on 2026-10-27, on this revision: 516 tested, and #4178's head could not
be fetched (its four files include none of this PR's). The conflicts this branch adds are the three above.
Approach
The change is in
src/prefill/prefill.cpp(the chunk's stream setup, the issuer thread, the host grouping, the join)plus a header-only helper,
include/strata/prefill/routed_stream.hpp, which holds the route gate and the issuer walk sothe walk itself can be tested without a GPU.
per-expert counts (
m.cnt) right after the count loop, then stores a per-layer atomic flag with release semantics; theissuer reads the flag with acquire before the row. One writer per layer per chunk, reset before the issuer thread
starts, destroyed after it is joined.
k % ring, as today. An entry iscopied only when both today's ring rule (
k < a_consumed + ring) and the depth rule hold: the copy issueddepthcopies earlier has landed (
cudaEventQueryon its slot'scopiedevent). While the walk waits for either, itre-checks the skip, so an entry is skipped exactly when its layer's routing has landed before it was issued.
a_issuedadvances past every entry, copied or skipped, so the consumer'swait_issuedis unchanged.on copies already issued, so there is no cycle. A skipped entry touches no slot, event or
stage_liveflag, and theslot it would have used is still released by the consumer's own
usedevent. The copy stream stays in order, so agroup's single wait on its last copy still covers the group.
a_stopis checked in the wait (early exits).blob, into the same slots, as today, whatever the ring held.
ring > 8, 1,024 tokens or more) on one GPU with the issuer thread. The fused prompt path(
STRATA_PF_FUSED) publishes no routing, so its layers skip nothing; the peer prompt ring (--peer-device) and theinline issuer (
STRATA_PREFILL_ISSUER=0) keep the switch off. Host code only: both HIP builds compile the new callsthrough the existing compat headers; the Metal build compiles
prefill.cpp, but itsPrefill::runhands the promptto the Metal window before
run_impl, where the walk is; SYCL builds its ownprefill.cpp.STRATA_TRACE=1prints, per chunk, the ring entries copied and skipped; one start-up line names thedepth when the switch first acts.
tests/prefill/routed_stream_test.cpp(GPU-free, underSTRATA_BUILD_TESTS) runs the sameissue_walktemplate the engine uses against a simulated in-order link and a consumer that walks the layers as
prefill.cppdoes:8 fixed and 38 random simulated chunks plus 6 three-thread walks with real atomics. On every one: each routed entry is
copied once into slot
k % ringand read with its own bytes; only unrouted, arena-copied entries of published layersare skipped; never more than
depthcopies are queued; no copy goes past the ring;a_issuedcovers every entry.Validation
Revision
ab73738f: this branch, one commit onupstream/main4b9e05a2(v0.2.6, engine 0.1.51); base4b9e05a2,built the same way. The binaries below were built from
56294b33, which differs fromab73738fonly indocs/DETAILS.md(these numbers) and the commit message:CMakeLists.txt,prefill.cpp, the header and the test arethe same blobs. The change was first written and measured on a branch with two other opt-in commits under it; on
mainits own lines are unchanged and only the placement of the test target and the DETAILS paragraph moved. Everything below
except the last block of details was built and measured on
main.System (one machine; no claim for others):
GB/s host to device), driver 580.178.04, CUDA 13.0 (nvcc 13.0.88). The card's tensor rate is unlocked by a kernel
module (FP16 GEMM with FP32 accumulation measured 65.7 TFLOPS before and 100.5 TFLOPS after, at 250 W), active since
2026-10-26 and the same in every run here.
pinned to 16 physical cores and the GPU passed through; host RAM DDR4-2133 in four channels.
free -g), no swap, Ubuntu 24.04.5 LTS, kernel 6.8.0-146-generic, gcc 13.3.0, CMake 4.4.3,Ninja 1.13.2; model files on a virtual (QEMU) disk.
-DSTRATA_ENABLE_CUDA=ON -DSTRATA_BUILD_TESTS=OFF -DSTRATA_Q6K_EXPERTS=ON -DSTRATA_MMQ_KQUANTS=ON -DCMAKE_CUDA_ARCHITECTURES=86 -DCMAKE_CUDA_COMPILER=<cuda-13.0>/bin/nvcc -DCUDAToolkit_ROOT=<cuda-13.0> -DSTRATA_GGML_DIR=<llama.cpp at setup's pin 3cf03257>(the CMake cache equals the installed engine's apart from theSTRATA_PLE_FIXTURE_DIRpath).serve/server.py --engine strata --config <config> --port 8080(v0.2.6, unchanged; run through a wrapper that onlyrecords the server's per-request log lines) with
strata --serve --pack <pack> --native <gguf> --ple-gguf <gguf> --expert-profile <expert-profile.bin> --expert-cache auto --prefill auto --spec 4 --spec-min-p 0.5 --mtp <mtp dir> --max-context 262144 --kv int8 --kv-resident 32768. In the benchmark noSTRATA_*variable but the switch; 10,320 expert slots and 23,738 MiB of VRAM in use at ready in every arm, switch on or off.
Passed:
CUDA build, the line under System:
cmake --build <build> --target strata strata_kqfor base and head, rc 0 (165 stepseach), 248 warning lines each with the same set (none in
prefill.cpp); both reportstrata 0.1.51; the switch'sname is in the head binary and not in the base one. The test build (
-DSTRATA_BUILD_TESTS=ON) rc 0 with no warning,and
ctest -Rover the 12 GPU-free unit tests ofCMakeLists.txtplusrouted_stream_test: head 13 of 13, base 12of 12.
Metal on macOS 15.7.7 (Apple M3, Apple clang 17.0.0), the steps of the
Metal buildworkflow (metal-cpp 26 at itspinned SHA-256, llama.cpp
3cf03257,cmake -S metal -B build-metal -G Ninja -DCMAKE_BUILD_TYPE=Release -DMETAL_CPP_DIR=<metal-cpp> -DSTRATA_GGML_DIR=<llama.cpp>thencmake --build build-metal --target strata): base andhead rc 0, the same 76 object files (
prefill.cppamong them) and the same 10 warnings.HIP and SYCL, compile only (x86-64 Linux containers, no AMD or Intel GPU): base built in full, then this commit
checked out in the same tree and rebuilt (
ninja -k 0, every target).docs/AMD_HIP.md's line (-DSTRATA_ENABLE_HIP=ON -DSTRATA_ENABLE_CUDA=OFF -DCMAKE_HIP_ARCHITECTURES=gfx1100) plus-DSTRATA_BUILD_TESTS=ON, inrocm/dev-ubuntu-24.04:7.14.1-full(HIP7.14.60850, AMD clang 23.0.0git, gcc 13.3.0, CMake 3.28.3): base rc 0 (428 steps), head rc 0 (15 steps rebuilt,
prefill.cppamong them).ctest -R 'routed_stream_test|share_rules_test|host_limits_test|prefill_plan_test': 4 of4 pass.
docs/AMD_HIP.md's line (-DSTRATA_HIP_GFX906=ON -DCMAKE_HIP_ARCHITECTURES=gfx906, ROCm clang as the Cand C++ compiler) plus
-DSTRATA_BUILD_TESTS=ON, same image:prefill.cppcompiles and the same 4 tests pass. Onetarget fails on base and head alike:
src/core/dflash2.cpp(cudaDeviceGetStreamPriorityRangeandcudaStreamCreateWithPriorityhave no mapping in the gfx906 compat headers), sostratadoes not link in thisconfiguration on
maineither.strata-sycl-devrecipe (docker build --pull -t strata-sycl-dev sycl/tools, thensycl/tools/build.shwithout AOT; oneAPI DPC++ 2026.1.0, CMake 4.2.3), configured with
STRATA_GGML_DIRat the pinned llama.cpp: base rc0 (232 steps), head rc 0 (50 relinks, nothing recompiled: the SYCL tree compiles none of the changed files).
hipError_tnodiscard warnings in
prefill.cpp, at lines 3080 to 3089: the uncheckedcudaStreamWaitEvent, twocudaMemcpyAsyncandcudaEventRecordof the routed walk's copy, which repeats the default loop's copy (its samefour calls already warn; base reports 235 (wave32) and 227 (gfx906) of these in
prefill.cpp). No other new warning.routed_stream_teston macOS with Apple clang 17:c++ -std=c++20 -O2 -Wall -Wextra -Iinclude tests/prefill/routed_stream_test.cpp -o rst && ./rstprintsrouted_stream_test: OKwith no warning, and the same with-O1 -g -fsanitize=threadgives no ThreadSanitizer report.Default path byte-identical, and the switch output-preserving, at 1K to 32K. One locked batch, a fresh server and
engine per arm, four arms in this order: base (
main), head with the switch unset, head withSTRATA_PREFILL_ROUTED_STREAM=1, base again (the control). Deterministic lane in every arm (STRATA_PREFILL_CPU_SHARE=0,--pcie-frac 0.32,--adapt-every 0) plusSTRATA_TRACE=1,STRATA_STATE_HASH=1andSTRATA_PREFILL_DUMP_R(the prompt-path residuals of every 64th position, fp32); temperature 0, thinking off, 64 tokens. Prompts: four
consecutive slices of the 128K identity text (documentation, each with a label and a one-sentence question), a
two-turn conversation (14,277 tokens, then a follow-up turn of 1,979 new tokens), and the 31,473- and 126,970-token
identity prompts. Compared per step: the reply's SHA-256, the engine's
prompt N tokens = R reused + K readandgeneratedcounts, itsSTATE_HASHline (GDN, PLE history, index tails and pools, K/V cells, the stale cells past thelive length, MTP K/V) and the residual dump's SHA-256. Up to 32K every field is identical in all four arms; the head
with the switch on equals the head with it unset on all eight steps, 128K included. No error line in any engine or
server log; the switch's start-up line appears in its own arm only.
f7f174354f1175ca9eba857dd98257b17311ac3a371f213de0ede8e747f668301fc92b9e1fc92b9e1fc92b9eThe base streamed the same experts as the head with the switch unset at every step. The 128K row is under Failed.
Speed on
main, release defaults (above), the head binary with the switch unset and set, one locked batch in theorder OFF ON ON OFF, then ON OFF OFF ON (pairs: arms 2 and 1, 3 and 4, 5 and 6, 8 and 7), a fresh server and engine per
arm. Requests per arm, in this order: counting prompts of 6,011 and 1,979 tokens; six tiny chats (17 to 19 prompt
tokens, 64 tokens); counting prompts of 955, 3,963 and 32,105 tokens; a 3,828-token natural-text prompt twice (384
tokens); a 3,971-token prompt with thinking on, temperature 1.0 and seeds 4242, 4243, 4244 (256 tokens); an 18,768-token
agent prompt with 33 tools and its second turn (1,790 new tokens), thinking on, the server's default sampling with a
fixed seed (384 tokens). Counting prompts: a filler with a unique nonce in every repetition, 2 runs each (the second
run's read is reported; each run is preceded by two 1-token calibration requests), 256 tokens, temperature 0, thinking
off, 0 tokens reused. Engine-side numbers: prompt tokens over the engine's prompt ms, the engine's decode tok/s.
Off and On are the means of the four arms per side. Follow-up turn: median pair ratio 1.218, percentile bootstrap 95%
interval 1.209 to 1.225 (20,000 resamples of the four ratios). The switch's code runs in every read above but the
955-token one and the tiny chats; the streamed experts fell by 23% at 1,979 tokens, 16% at 3,963, 9% at 6,011, 4% at
32,105 and 22% on the follow-up turn (table below). At these settings the replies differ between arms of the same
binary and switch (the adaptive CPU share and PCIe split): five of the six tiny chats' replies differ among the four
switch-off arms, and the natural-text, sampled and agent requests produce a different text in every arm, so those
decode rows compare different texts.
Every run of the timing batch
Engine prompt ms per run (counting prompts: run 1 / run 2):
Engine decode tok/s per run:
Experts streamed per request (the server's per-request count), four arms per side:
Unpaired prompt times from the identity arms (deterministic lane, one run each, not a benchmark), base / unset / on /
base again, ms: 1,046 tokens 2,948 / 2,933 / 2,030 / 2,963; 1,965 tokens 3,213 / 3,132 / 2,560 / 3,100; follow-up turn
3,049 / 3,049 / 2,486 / 2,957; 31,473 tokens 15,669 / 15,499 / 14,639 / 15,286; 126,970 tokens 60,480 / 60,615 / 57,034 /
60,734.
The earlier run: production settings and another base, 4 pairs
The same change before the rebase (identical lines in
prefill.cpp, the header and the test), on v0.2.6 with two otheropt-in commits under it (#4119's
STRATA_SELECT_TWO_PASSand a local asynchronous checkpoint,STRATA_CKPT_ASYNC, noton
main), measured on the same machine on 2026-10-26 with settings that are not the release defaults, in both arms:STRATA_TOPK_SPLIT=1,STRATA_CKPT_POOL=1,STRATA_SELECT_TWO_PASS=1,STRATA_CKPT_ASYNC=1,--conversation-cache-min-tokens 2048and--conversation-cache-mib 8192. The same requests as above, two batches (OFFON ON OFF, then ON OFF OFF ON). Absolute numbers are not comparable with the table above; compare the ratios.
Pairs: arms 2 and 1, 3 and 4 of the first batch, 1 and 2, 4 and 3 of the second. Follow-up turn: percentile bootstrap
95% interval 1.176 to 1.269.
Failed:
main, the samebinary) gave different replies to the 126,970-token prompt (
47f66830and1fc92b9e) and differentSTATE_HASHlines after decode, with the same prompt-path residuals and counts; the first base arm logged another decode
expert-cache hit count (33,733 against 32,948 in the other three arms). So byte identity at 128K is not established
for the default path. The head with the switch unset and with it set both equal the second base arm byte for byte,
and the prompt-path residuals are identical in all four arms.
above 1.0, and no greedy decode or prompt-read tier with a median pair ratio below 0.98) held except for one tier: the
tiny chats' decode, 0.9797. Those chats (17 to 19 prompt tokens) never reach the switch's code, their replies differ
between arms of the same binary at these settings, and the same tier was 1.008 in the earlier run.
Not run:
window reads the prompt); on HIP it would run where the CUDA build runs it, and that is unverified on hardware.
main(no 128K or 256K pairs; in the unpaired identity arms the switch read 126,970 tokens in5.9% less time, one run each), any depth other than 16, and the switch together with
STRATA_PREFILL_STREAM_MIN.there, so the identity runs pin both. Task-quality or long-context recall checks (the replies were byte-identical).
by construction).
Limits
on
main42% of the streamed experts at 1,046 tokens, 14% at 4,867, 6% at 14,277 (two chunks), 3% at 31,473 and 2% at126,970.
host under memory pressure would skip less.
cudaEventQuerypoll's host cost (one driver call per copy) was notmeasured on its own; a host whose CPUs are busy could see it.
has no guard, so no slack is left inside the loan. Where a slot does have slack past its blob, those bytes hold
prompt-dependent leftovers today as well, but this PR does not prove that no kernel reads them.
runs pin both.
(see Related discussion).