Repository navigation
[Engine] Opt-in routed prompt stream: do not copy the streamed experts no token of the chunk routes to #4282
Description
Area
[Engine] (speed, kernels, models, the server and API)
What do you want to do?
Read prompts and conversation turns of about 2K to 32K tokens faster where the experts stream over PCIe, without
changing any output.
A prompt chunk of 1,024 tokens or more streams every non-resident expert of every layer through the prompt ring
(src/prefill/prefill.cpp), and the issuer thread enqueues each copy as soon as the ring has room. The ring holds more
entries than a layer has, so a layer's copies are on the copy stream before its router has run, and an expert that no
token of the chunk routes to is copied anyway. On an RTX 3090 with the Swift IQ3_XXS pack, the switch below skipped 42%
of the streamed experts of a 1,046-token read, 22% of a 1,979-token follow-up turn and 3% of a 31K read (only entries
not yet issued when their layer's routing was known can be skipped).
Proposed behaviour, opt-in with STRATA_PREFILL_ROUTED_STREAM=1 (off by default): the issuer keeps at most a few copies
(default 16) queued ahead of the link, each layer hands the issuer the set of experts its tokens route to as soon as its
router has run, and the issuer skips an entry it has not issued yet when no token routes to it. Every routed expert is
still copied, in the same order, into the same ring slot, so the prompt's results are the same bits. With the variable
unset nothing new runs.
Measured on a working branch (one commit on main 4b9e05a2, v0.2.6; one RTX 3090 24 GB on PCIe 3.0 x16 at 250 W,
release defaults, the same binary with the switch off and on, 4 pairs): the read of a 1,790-token follow-up turn on an
18.8K-token conversation took 17.9% less time (3,214 to 2,640 ms; pair ratios 1.209 to 1.225), and fresh reads of 2K,
4K, 6K and 32K tokens ran 19.9%, 8.2%, 2.1% and 5.3% faster. Decode after those prompts moved by -0.8% to +0.5%; in six
64-token chats, which never reach the switch's code, by -2.6% (pair ratios 0.89 to 1.06); in two agent turns by -3.5%,
on texts that differed in every arm (60.6% of drafts accepted against 62.3%). Replies are byte-identical to
main at 1K to 32K in a deterministic lane, with the switch off and on.
Tradeoffs:
- About 200 lines of host code (57 in
prefill.cpp, 142 in a new header holding the walk and the gate) and a 381-line
GPU-free test; one commit, no lines removed. - One GPU with the threaded issuer only: the peer prompt ring and the inline issuer keep today's walk, the fused prompt
path publishes no routing (so nothing is skipped there), and experts staged through the Stager's pinned buffers are
never skipped. - The gain is largest for short streamed chunks (2K, 4K) and small for long ones (6K, 32K); there is none below the
1,024-token stream floor. - To keep the queue short the issuer polls one event per copy (
cudaEventQuery); that host cost was not measured on its
own.
Is this in scope as an opt-in switch? The code, tests and measurements are in a pull request opened with this issue,
which links here; I will rework it or close it if the scope should be different.
What do you do today instead?
Nothing: every streamed expert is copied. STRATA_PREFILL_RING=8 (#225, on main as e258e6ec) stages only the routed
experts but gives up the streamed ring's read-ahead; this proposal keeps the ring and its read-ahead and skips only what
the routing shows is unused.
Activity