Skip to content

perf(ArcInfer/ArcGraph): device decode loop — 2 host syncs per token to 0, opt-in - #217

Merged
heydryft merged 4 commits into
release/openrouter-readyfrom
arcgraph/device-decode-loop-rel
Aug 21, 2026
Merged

heydryft merged 4 commits into
release/openrouter-readyfrom
arcgraph/device-decode-loop-rel

Conversation

@heydryft

@heydryft heydryft commented Aug 21, 2026 •

Copy link
Copy Markdown
Contributor

Rebase of arcgraph/device-decode-loop-on-205 onto release/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 on cudaStreamSynchronize (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:

  1. cuGraphLaunch — the captured forward.
  2. CudaSampler (arc-cuda-graph/src/sampling_cuda.rs) — argmax on device, token id to an i32[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.
  3. arc_graph_step_commit — scatters that token into the pinned u32 buffer 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 — cuGraphLaunch returns SUCCESS and only a later API call surfaces it, which is exactly why replay() 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.
  • Idle-but-short — cudaSuccess from 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.
  • A sticky device fault word — the commit kernel range-checks the token and the position before it writes anything, and latches a code into a pinned word that travels back with the ring.

max_spins is a last-resort liveness backstop for a kernel that hangs without ever erroring, not the primary detector.

Aliasing

The logits tensor replay_handles returns aliases graph-owned storage the next launch overwrites — the same BOOBY TRAP replay documents. 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.rs enforces that invariant at runtime, not just in a comment.

Layering

Everything above the CUDA calls is expressed against a DeviceOps trait that deliberately has no synchronize method — the driver cannot block the stream because it is not given a way to. The resulting BurstDriver is plain Rust, compiles without the cuda feature, 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, no and 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 inside sample rather 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 for ARC_GRAPH_VERIFY_ALWAYS, which would otherwise verify nothing while looking like a clean diagnostic.

KNOWN LIMITATION — a burst holds the scheduler

A burst enqueues burst decode 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 to burst decode 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() refuses batch_size != 1 outright 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 and src/cuda/decode_loop.cu had 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=1 together. ARC_CAPTURE_STREAM accepts only the literal "1". ARC_V4_PAGED_ATTN=1 is mutually exclusive with the capture probe — paging shadows the only graph arm V4 has.


Rebase notes

  • Cut on 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 to release/openrouter-ready (re-cut on master 8ae209008) with no conflicts. The diff against the pre-rebase branch is the borrow fix below and nothing else.
  • Local, on a Mac: cargo check -p arc-cuda-graph --tests clean, cargo clippy --no-deps -p arc-cuda-graph --tests --examples -- -D warnings clean, and the 21 device_loop unit 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:

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()
    |            ^^^ `seq` is a `&` reference, so it cannot be borrowed as mutable

Sequence::sampler() takes &mut self (sequence.rs:885) even though it only clones an Arc, so a &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. Fixed with rustc's own suggestion — seq: &mut Sequence on both the cuda and not(cuda) twins, so the signatures cannot drift. Nothing mutates the sequence, and the && chain still borrow-checks because sampler() returns an owned Arc<Sampler>, so the mutable borrow ends before seq.recognizer is read. Verified on a CUDA-less host by removing the cfg and stubbing the three arc_cuda_graph calls, which puts the real Sequence through borrowck.

On that first run both nvcc jobs got past Compile arc-cuda-graph CUDA kernels before failing at the later arc-engine step, so src/cuda/decode_loop.cu — the new arc_graph_step_commit kernel — already compiled clean under nvcc for sm_80 and sm_90.

With the fix in (2ad07fc), all three lanes are green:

lane verdict
cargo check (cuda, workspace) pass — 5m27s
nvcc compile (sm_80) pass — 14m1s
nvcc compile (sm_90) pass — 12m43s

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_tokens documents that it "MUST be called when a sequence finishes, is preempted, or the path falls back", and the unit test clearing_pending_tokens_prevents_leaking_them_into_the_next_sequence covers the function — but the only call sites are the fall-back paths in normal.rs and the not-eligible arm of sampling.rs. Nothing calls it on completion. eos_token_id is 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 first sample_sequence — including its prefill step — takes one via take_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) and sampler.rs:1490 (logprob: 0.0) are correct on the base but land at :1501 and :1512 here, because is_greedy_trivial adds 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.

heydryft and others added 3 commits August 21, 2026 16:42
…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
@github-actions

Copy link
Copy Markdown
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant