Repository navigation
perf(ArcInfer/ArcGraph): device decode loop — 2 host syncs per token to 0, opt-in - #217
Merged
heydryft merged 4 commits intoAug 21, 2026
Merged
Conversation
…h, zero host syncs CUDA-graph capture already cut launch APIs per token 3,961 -> 3.8 and bought only ~8%. The reason is not that op count stopped mattering: replay() launches the graph and then blocks on cudaStreamSynchronize (graph.rs:362), the caller samples on the host, and the host-side argmax D2H (sampler.rs:1479) re-synchronizes anyway. Capture removed the cost of *issuing* the work and left the serialized host round trip between steps completely intact. Removing the sync alone buys nothing, so this removes the sync and the host sample together, behind an opt-in. Per step, all on one stream: cuGraphLaunch -> CudaSampler (on device, token to an i32 device buffer, never copied down) -> arc_graph_step_commit (scatters the token into the PINNED U32 buffer the captured graph reads, advances the device position, publishes into a pinned+mapped host ring). Nothing blocks. Fault detection replaces the sync rather than dropping it, with three signals that need no blocking call: non-blocking cudaStreamQuery (which NVIDIA documents as also returning errors from previous asynchronous launches), an idle-but-short ring (every step retired yet fewer tokens landed — a fault with no timeout heuristic in it), and a sticky device fault word the commit kernel latches before it writes anything. A spin backstop covers a kernel that hangs without ever erroring. The driver is expressed against a DeviceOps trait that deliberately has no synchronize method, so it is structurally incapable of blocking the stream, and it compiles without CUDA — 20 unit tests run against a scripted fake GPU on any host. Two mutations were checked: dropping the device-fault read and skipping one launch per burst each fail exactly the tests that name them. Aliasing: the logits tensor aliases storage the next launch overwrites. Safe here only because every access is stream-ordered and none is a host read; the ring the host drains is a separate pinned allocation no launch touches. The invariant that lets the aliasing tensor be returned at all is enforced at runtime, not just documented. Default OFF (ARC_GRAPH_DEVICE_LOOP). Greedy-only: the device sampler runs Splitmix64 against the host's Isaac64, so any stochastic mode would draw a different, equally valid token and silently break seeded reproducibility. Every refusal falls back to the existing path; one failure latches it off for the process. Kept as a pure insertion, away from replay(), to merge cleanly with PR #205. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7
The device-loop branch sits earlier in the chain than the runner.replay(bs) arm that PR #205 extended with ARC_GRAPH_VERIFY_ALWAYS. Without this it would take every step and that switch would silently verify nothing — a clean diagnostic that never ran. Also cross-references #205's BOOBY TRAP note from replay_handles, since the same aliasing applies to the handles it returns. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7
A burst admits no new request for its whole length (default 4 decode steps). Harmless at batch 1 — the only shape admit() accepts — and not harmless above it, which is why admit() refuses batch_size != 1 outright rather than treating it as a tuning question. Stated in the module docs so it survives past the return it was first noted in. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7
Code Metrics Report━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ Language Files Lines Code Comments Blanks ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ C Header 5 305 210 52 43 CSS 2 1181 1036 34 111 CUDA 81 29475 20161 6275 3039 Dockerfile 1 39 22 8 9 JavaScript 16 3546 2676 482 388 Jinja2 7 694 656 5 33 JSON 75 4896 4893 0 3 Makefile 1 6 5 0 1 Metal Shading Lan| 33 12224 9431 1142 1651 PowerShell 1 300 227 30 43 Python 147 15371 12677 824 1870 Shell 42 10304 6858 2769 677 Plain Text 4 3801 0 2479 1322 TOML 33 1498 1294 54 150 YAML 3 25 23 2 0 ───────────────────────────────────────────────────────────────────────────────── HTML 4 2687 2604 43 40 |- CSS 2 543 479 37 27 |- JavaScript 1 1233 1215 12 6 (Total) 4463 4298 92 73 ───────────────────────────────────────────────────────────────────────────────── Jupyter Notebooks 4 122 83 23 16 |- Markdown 1 60 30 22 8 |- Python 1 122 113 1 8 (Total) 304 226 46 32 ───────────────────────────────────────────────────────────────────────────────── Markdown 213 48932 0 38081 10851 |- BASH 72 1655 1203 331 121 |- C 4 19 19 0 0 |- CUDA 2 84 56 16 12 |- JSON 19 779 779 0 0 |- PowerShell 1 1 1 0 0 |- Python 23 1008 787 113 108 |- Rust 68 2063 1727 78 258 |- TOML 6 207 164 0 43 |- YAML 5 41 36 5 0 (Total) 54789 4772 38624 11393 ───────────────────────────────────────────────────────────────────────────────── Rust 681 340988 293252 18096 29640 |- Markdown 504 32562 471 28063 4028 (Total) 373550 293723 46159 33668 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ Total 1353 516771 363188 99077 54506 ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ |
The CUDA lane's whole purpose, on its first run against this branch: three
jobs, one error, in the one block a macOS `cargo check` structurally cannot
reach.
error[E0596]: cannot borrow `*seq` as mutable, as it is behind a `&` reference
--> mistralrs-core/src/pipeline/sampling.rs:557:12
557 | && seq.sampler().is_greedy_trivial()
`Sequence::sampler()` takes `&mut self` (sequence.rs:885) even though it only
clones an Arc, so the `&Sequence` parameter could never call it. `sample_sequence`
already holds `&mut Sequence` and was reborrowing it down to `&`, which is why
the call site compiled and the body did not.
This is rustc's own suggested fix. Nothing here mutates the sequence, and the
`&&` chain still borrow-checks: `sampler()` returns an owned `Arc<Sampler>`, so
the mutable borrow ends before `seq.recognizer` is read. Verified by compiling
the cuda body on a host with no CUDA — cfg removed, the three `arc_cuda_graph`
calls stubbed — which puts the real `Sequence` through borrowck and passes.
The `#[cfg(not(feature = "cuda"))]` twin moves with it so the two signatures
cannot drift.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7
This was referenced Aug 21, 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.
Rebase of
arcgraph/device-decode-loop-on-205ontorelease/openrouter-ready. The branch it was cut on (arcgraph/capture-lane-to-master) landed as #205, so only the three device-loop commits are carried here; the twelve capture-lane commits below them are already on master and were dropped.What this does
CUDA-graph capture already cut launch APIs per token 3,961 → 3.8 and bought ~8%. That is not op count ceasing to matter — it is that
replay()launches the graph and then blocks oncudaStreamSynchronize(arc-cuda-graph/src/graph.rs:362), the caller samples on the host, and the host-side argmax D2H (mistralrs-core/src/sampler.rs:1479) re-synchronizes anyway. Capture removed the cost of issuing the work and left the serialized host round trip between steps completely intact. There are two host syncs per decode step, not one, and removing either alone buys nothing. This removes both, together, behind an opt-in.Per step, all on one stream, no host in the middle:
cuGraphLaunch— the captured forward.CudaSampler(arc-cuda-graph/src/sampling_cuda.rs) — argmax on device, token id to ani32[batch]device buffer, never copied down. This is its first caller outside its own crate; it was already complete and tested bit-exact against the host argmax, and had zero users.arc_graph_step_commit— scatters that token into the pinnedu32buffer the captured graph already baked at capture time (so the next launch consumes it), bumps the device position, and publishes the token into a pinned+mapped host ring.The host drains that ring with plain volatile loads — no CUDA API call, no stream sync, no D2H. Draining is a host memory poll, categorically different from a stream sync: with a sync the host cannot enqueue step N+1 until step N retires; here every step is already enqueued and the host only watches them land.
Fault detection replaces the sync rather than dropping it
A fault during graph execution is asynchronous —
cuGraphLaunchreturns SUCCESS and only a later API call surfaces it, which is exactly whyreplay()used a full sync. Dropping the check outright would turn an illegal access into a silently poisoned context and an output buffer of plausible garbage. Three signals cover it with no blocking call:cudaStreamQuery— non-blocking. NVIDIA's CUDA Runtime API reference states: "Note that this function may also return error codes from previous, asynchronous launches." A sticky async fault surfaces on a call that never blocks.cudaSuccessfrom the query means every enqueued step retired. If all of them retired and fewer tokens than steps reached the ring, the work definitively did not do what it was asked. No timeout heuristic in it.max_spinsis a last-resort liveness backstop for a kernel that hangs without ever erroring, not the primary detector.Aliasing
The logits tensor
replay_handlesreturns aliases graph-owned storage the next launch overwrites — the same BOOBY TRAPreplaydocuments. It is safe here for one specific reason, stated rather than assumed: every access is enqueued on the same stream, in order, and none of them is a host read. Stream ordering guarantees the sampler starts after launch N wrote the logits and finishes before launch N+1 overwrites them. The ring the host drains is a separate pinned allocation no launch touches.normal.rsenforces that invariant at runtime, not just in a comment.Layering
Everything above the CUDA calls is expressed against a
DeviceOpstrait that deliberately has nosynchronizemethod — the driver cannot block the stream because it is not given a way to. The resultingBurstDriveris plain Rust, compiles without thecudafeature, and runs 20 unit tests against a scripted fake GPU on any host. Two mutations were checked: dropping the device-fault read and skipping one launch per burst each fail exactly the tests that name them.Default OFF
ARC_GRAPH_DEVICE_LOOP(arc-cuda-graph/src/device_loop.rs:135, gate at:170). Unset → off. Read by value, not by presence:0,false,off,noand empty are explicit refusals, so an operator can pin it off in an environment that sets it upstream. Greedy-only — the device sampler runs Splitmix64 against the host's Isaac64, so any stochastic mode would draw a different, equally valid token and silently break seeded reproducibility;Sampler::is_greedy_trivial()is the precondition and mirrors the conditions insidesamplerather than approximating them. Every refusal falls back to the existing path, and one failure latches it off for the process. The loop also stands aside forARC_GRAPH_VERIFY_ALWAYS, which would otherwise verify nothing while looking like a clean diagnostic.KNOWN LIMITATION — a burst holds the scheduler
A burst enqueues
burstdecode steps back-to-back and does not return to the engine until they have all landed. For that whole window the engine admits no new request, so a request arriving mid-burst waits up toburstdecode steps (default 4) before it is even considered.At batch 1 — the only shape
admit()accepts — that is harmless: there is one sequence, nothing to interleave with, and the alternative is a host round trip per token. It is not harmless above batch 1, where a burst would straddle sequences with different stop conditions and different admission urgency, and where the wasted tail after one sequence's EOS is multiplied by the batch.admit()refusesbatch_size != 1outright rather than treating it as a tuning question. Anyone widening this past batch 1 owns that problem first.Status of the claims
The mechanism is implemented and the host-side driver is unit-tested. No number in this PR is a measurement. The "2 → 0 host syncs per step" is a property of the code as written, and the throughput consequence is derived, not observed. This branch was written on macOS, which has no nvcc, so its
#[cfg(feature = "cuda")]blocks andsrc/cuda/decode_loop.cuhad never been compiled by anything before this PR — closing that gap is the purpose of opening it.Still blocking a hardware run: capture must succeed, which needs
ARC_CAPTURE_STREAM=1 ARC_V4_CAPTURE_PROBE=1 ARC_GRAPH_DEVICE_LOOP=1together.ARC_CAPTURE_STREAMaccepts only the literal"1".ARC_V4_PAGED_ATTN=1is mutually exclusive with the capture probe — paging shadows the only graph arm V4 has.Rebase notes
arcgraph/capture-lane-to-master, which merged as ArcGraph: land the CUDA-graph capture lane — 3,961 -> 3.8 launch APIs/token, 10,114 -> 131 allocs/step #205 (fdf0b6c19). The twelve capture-lane commits were dropped rather than replayed; only the three device-loop commits are carried, and they apply torelease/openrouter-ready(re-cut on master8ae209008) with no conflicts. The diff against the pre-rebase branch is the borrow fix below and nothing else.cargo check -p arc-cuda-graph --testsclean,cargo clippy --no-deps -p arc-cuda-graph --tests --examples -- -D warningsclean, and the 21device_loopunit tests pass against the scripted fake GPU.The CUDA lane found exactly what it exists to find
The branch was written on macOS, so its
#[cfg(feature = "cuda")]blocks had never been compiled by anything. The first CI run (e385e5f) took all three CUDA jobs red on one error, inside the CUDA-gated block:Sequence::sampler()takes&mut self(sequence.rs:885) even though it only clones anArc, so a&Sequenceparameter could never call it.sample_sequencealready holds&mut Sequenceand was reborrowing it down to&, which is why the call site compiled and the body did not. Fixed with rustc's own suggestion —seq: &mut Sequenceon both thecudaandnot(cuda)twins, so the signatures cannot drift. Nothing mutates the sequence, and the&&chain still borrow-checks becausesampler()returns an ownedArc<Sampler>, so the mutable borrow ends beforeseq.recognizeris read. Verified on a CUDA-less host by removing thecfgand stubbing the threearc_cuda_graphcalls, which puts the realSequencethrough borrowck.On that first run both nvcc jobs got past
Compile arc-cuda-graph CUDA kernelsbefore failing at the laterarc-enginestep, sosrc/cuda/decode_loop.cu— the newarc_graph_step_commitkernel — already compiled clean under nvcc for sm_80 and sm_90.With the fix in (
2ad07fc), all three lanes are green:cargo check (cuda, workspace)nvcc compile (sm_80)nvcc compile (sm_90)That is the whole deliverable this PR was opened for: every
#[cfg(feature = "cuda")]block on this branch, and the new.cu, now compile for A100 and Hopper. It says nothing about whether they run.Two things worth a look before anyone spends GPU time on this
1. Parked tokens are never cleared when a sequence finishes.
clear_pending_tokensdocuments that it "MUST be called when a sequence finishes, is preempted, or the path falls back", and the unit testclearing_pending_tokens_prevents_leaking_them_into_the_next_sequencecovers the function — but the only call sites are the fall-back paths innormal.rsand the not-eligible arm ofsampling.rs. Nothing calls it on completion.eos_token_idis passed as-1(normal.rs,DeviceLoopConfig::new(vocab, -1, capacity)), which disables EOS truncation, so a burst always parks its full length. If a sequence stops on any token but the last of a burst, the remainder stays parked, and the next greedy-trivial sequence's firstsample_sequence— including its prefill step — takes one viatake_pending_token(). That is the "+1-step divergence" failure mode the module docs name, arriving by a different route.2. Two doc line-citations drift by 22 lines against this branch's own insertion.
sampler.rs:1479(the argmax D2H) andsampler.rs:1490(logprob: 0.0) are correct on the base but land at:1501and:1512here, becauseis_greedy_trivialadds 22 lines above them. Left as the author wrote them rather than silently re-numbered.Numbers
Unchanged and still not measurements: "2 host syncs per step to 0" is a property of the code as written, and the throughput consequence is derived. The only new fact this PR produces is whether the CUDA-gated code compiles.