From f46354a1d283458adc1c7030b95e957365c9410b Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:23:45 +0100 Subject: [PATCH 1/4] =?UTF-8?q?perf(ArcInfer/ArcGraph):=20device=20decode?= =?UTF-8?q?=20loop=20=E2=80=94=20N=20steps,=20N=20cuGraphLaunch,=20zero=20?= =?UTF-8?q?host=20syncs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- arc-cuda-graph/src/cuda/decode_loop.cu | 112 ++ arc-cuda-graph/src/device_loop.rs | 1778 +++++++++++++++++++++++ arc-cuda-graph/src/ffi.rs | 7 + arc-cuda-graph/src/graph.rs | 35 + arc-cuda-graph/src/lib.rs | 13 + mistralrs-core/src/pipeline/normal.rs | 207 +++ mistralrs-core/src/pipeline/sampling.rs | 65 + mistralrs-core/src/sampler.rs | 22 + 8 files changed, 2239 insertions(+) create mode 100644 arc-cuda-graph/src/device_loop.rs diff --git a/arc-cuda-graph/src/cuda/decode_loop.cu b/arc-cuda-graph/src/cuda/decode_loop.cu index b181950298..4a4894d3ad 100644 --- a/arc-cuda-graph/src/cuda/decode_loop.cu +++ b/arc-cuda-graph/src/cuda/decode_loop.cu @@ -152,3 +152,115 @@ extern "C" void launch_check_all_done_conditional( extern "C" int arc_has_graph_conditional() { return ARC_HAS_GRAPH_CONDITIONAL; } + +// ============================================================================ +// ArcInfer/ArcGraph — device decode loop step commit +// +// The narrow sibling of `decode_step_update` above, for the CUDA-graph replay +// path (`arc-cuda-graph/src/device_loop.rs`). That kernel is written for the +// paged-attention autonomous runner and unconditionally dereferences +// `finished`, `n_generated`, `slot_mappings` and `block_tables`; the replay +// path has none of those (V4 runs with `cache_config == None`), so passing +// nulls would fault. This kernel touches only what that path actually owns. +// +// It closes the decode loop on-device: +// sampled token -> the PINNED U32 buffer the captured graph reads +// -> device position += 1 +// -> pinned+mapped host ring, for the host to drain later +// +// Ordering is load-bearing. Every range check runs BEFORE any write, and the +// ring write head is published LAST, after `__threadfence_system()`. So a token +// the host can see in the ring is guaranteed to have been committed to the +// graph input buffer, and a refused commit publishes nothing at all rather than +// letting the next launch consume a garbage id. +// ============================================================================ + +#define ARC_COMMIT_FAULT_NONE 0 +#define ARC_COMMIT_FAULT_TOKEN_RANGE 1 +#define ARC_COMMIT_FAULT_POSITION_LIMIT 2 + +__global__ void arc_graph_step_commit( + const int32_t* __restrict__ sampled, // [batch] i32, written by CudaSampler + uint32_t* __restrict__ input_ids, // [batch] U32, PINNED graph input + uint32_t* __restrict__ positions, // [batch] U32, PINNED graph input; nullable + int32_t* __restrict__ ring, // [batch, ring_size] pinned + mapped + int32_t* __restrict__ ring_head, // [batch] pinned + mapped + int32_t* __restrict__ fault, // [1] pinned + mapped, sticky + int32_t ring_size, + int32_t vocab, + int32_t position_limit, // exclusive; <= 0 disables the check + int32_t batch_size +) { + int bid = blockIdx.x; + if (threadIdx.x != 0) return; + if (bid >= batch_size) return; + + // A fault latched by an earlier step means the graph input buffer already + // holds a stale id. Committing on top of that would paper over it. + if (fault != nullptr && ((volatile int32_t*)fault)[0] != ARC_COMMIT_FAULT_NONE) return; + + int token = sampled[bid]; + + // ---- checks first, writes after ------------------------------------- + if (token < 0 || token >= vocab) { + if (fault != nullptr) { + ((volatile int32_t*)fault)[0] = ARC_COMMIT_FAULT_TOKEN_RANGE; + __threadfence_system(); + } + return; + } + if (positions != nullptr && position_limit > 0) { + // `positions[bid]` is the slot this step just wrote KV to. The next + // step writes at +1, and the fixed-capacity window is `position_limit` + // wide, so +1 must stay strictly inside it. + if ((int64_t)positions[bid] + 1 >= (int64_t)position_limit) { + if (fault != nullptr) { + ((volatile int32_t*)fault)[0] = ARC_COMMIT_FAULT_POSITION_LIMIT; + __threadfence_system(); + } + return; + } + } + + // ---- commit --------------------------------------------------------- + // This single store is what removes the host from the loop: the captured + // graph baked THIS address at capture time, so the next `cuGraphLaunch` + // reads the token without any host round trip. + input_ids[bid] = (uint32_t)token; + if (positions != nullptr) { + positions[bid] += 1u; + } + + // ---- publish last ---------------------------------------------------- + if (ring != nullptr && ring_head != nullptr && ring_size > 0) { + int wp = ring_head[bid]; + int slot = wp % ring_size; + if (slot < 0) slot += ring_size; + ((volatile int32_t*)ring)[(int64_t)bid * ring_size + slot] = token; + // Order the token store ahead of the head store as seen by the HOST. + // Without this the host can observe an advanced head and read the + // previous occupant of the slot. + __threadfence_system(); + ((volatile int32_t*)ring_head)[bid] = wp + 1; + } +} + +extern "C" void arc_launch_graph_step_commit( + const int32_t* sampled, + uint32_t* input_ids, + uint32_t* positions, + int32_t* ring, + int32_t* ring_head, + int32_t* fault, + int32_t ring_size, + int32_t vocab, + int32_t position_limit, + int32_t batch_size, + cudaStream_t stream +) { + arc_graph_step_commit<<>>( + sampled, input_ids, positions, + ring, ring_head, fault, + ring_size, vocab, position_limit, batch_size + ); +} diff --git a/arc-cuda-graph/src/device_loop.rs b/arc-cuda-graph/src/device_loop.rs new file mode 100644 index 0000000000..d4aae80edf --- /dev/null +++ b/arc-cuda-graph/src/device_loop.rs @@ -0,0 +1,1778 @@ +//! Parent system: ArcInfer / ArcGraph +//! +//! Device-resident decode loop: **N decode steps cost N `cuGraphLaunch` calls +//! and zero `cudaStreamSynchronize`.** +//! +//! # What was wrong with plain replay +//! +//! [`crate::graph::CudaGraphRunner::replay`] launches the captured graph and +//! then calls `cudaStreamSynchronize` (`graph.rs:362`) before handing the +//! output tensor back. The caller samples the next token on the **host**, which +//! copies logits (or an argmax scalar) down over PCIe, and only then launches +//! the next graph. So a decode step is +//! +//! ```text +//! GPU runs step -> host blocks on sync -> D2H -> host samples -> launch +//! ``` +//! +//! Capture removed the cost of *issuing* the work — 3,961 launch APIs per token +//! became 3.8 — and left that serialized round trip completely intact. Deleting +//! the sync on its own buys nothing, because the host-side sample forces a +//! device-to-host copy that synchronizes anyway. **The sync and the host sample +//! have to go together.** +//! +//! # The loop this module closes +//! +//! Per step, all on one stream, no host in the middle: +//! +//! 1. `cuGraphLaunch` — the captured forward. +//! 2. [`crate::sampling_cuda::CudaSampler`] — argmax/top-p **on device**, token +//! id written to an `i32[batch]` device buffer. Never copied to the host. +//! 3. `arc_graph_step_commit` — scatters that token into the *pinned* graph +//! input buffer the captured graph reads (so the next launch consumes it), +//! bumps the device position, and publishes the token into a pinned+mapped +//! host ring. +//! +//! Nothing in that sequence blocks. The host enqueues all `burst` steps and the +//! GPU works through them while the host is already ahead. +//! +//! # How the host still gets its tokens +//! +//! The ring lives in `cudaHostAlloc(..., cudaHostAllocMapped)` memory. The GPU +//! pushes each token straight into it, `__threadfence_system()`, then publishes +//! a monotonic write head. The host reads both with **plain volatile loads** — +//! no CUDA API call, no stream synchronize, no D2H copy. Draining is a host +//! memory poll, which is 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 is only watching them land. +//! +//! # Why the aliasing logits buffer is safe here +//! +//! The tensor `replay()` hands back **aliases graph-owned storage that the next +//! `cuGraphLaunch` overwrites**. Holding it across a launch, or draining it +//! lazily, reads data the next step already clobbered — silently, producing +//! plausible tokens. That alias/clone class has already caused three +//! verification bugs in this repo. +//! +//! This design is safe against it for one specific reason, and it is worth +//! stating rather than assuming: **every access is enqueued on the same +//! stream, in order, and none of them is a host read.** Per step the stream +//! carries `launch N -> sample -> commit -> launch N+1`. CUDA stream ordering +//! guarantees the sampler kernel starts only after launch N has written the +//! logits and finishes before launch N+1 begins overwriting them. The host +//! never reads the logits at all, so there is no window in which it could read +//! a clobbered buffer. +//! +//! This is strictly stronger than the host path, which has to *synchronize* to +//! read those logits safely. What the host does drain — the token ring — is a +//! **separate, pinned, host-owned allocation** that no graph launch touches, so +//! draining it lazily is sound however far behind the host falls. The one thing +//! that would break the argument is issuing any of these on a different stream; +//! that is why `CudaDeviceOps` captures the graph's own stream and uses it for +//! the sampler and the commit kernel too. +//! +//! # Detecting a fault without synchronizing +//! +//! A fault *during* graph execution is asynchronous — `cuGraphLaunch` returns +//! SUCCESS and only a later API call surfaces it. `replay()` used a full sync +//! for exactly this reason and the comment there says so. Dropping the check +//! outright would turn an illegal access into a silently poisoned context and +//! an output buffer full of garbage, a bug class this codebase has shipped more +//! than once. So the check is *replaced*, not removed, by three signals that +//! together need no blocking call: +//! +//! * **`cudaStreamQuery`** — non-blocking. NVIDIA's CUDA Runtime API reference +//! for it states: *"Note that this function may also return error codes from +//! previous, asynchronous launches."* So a sticky async fault surfaces on a +//! call that never blocks. +//! * **Idle-but-short** — `cudaStreamQuery` returning `cudaSuccess` means every +//! enqueued step has retired. If all of them retired and fewer tokens than +//! steps reached the ring, the work definitively did not do what it was asked +//! to. That is a fault with no timeout heuristic in it. +//! * **The device fault word** — the commit kernel range-checks the sampled +//! token and the position before it writes anything, and latches a code into +//! a pinned word. It travels back with the ring, so reading it is free. +//! +//! `max_spins` is a last-resort liveness backstop for a kernel that hangs +//! without ever erroring. It is not the primary detector. +//! +//! # Layering +//! +//! Everything above the CUDA calls is expressed against the [`DeviceOps`] +//! trait, and that trait 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 is unit-tested against a scripted fake GPU. The CUDA-gated +//! [`CudaDeviceOps`] is pointer plumbing and nothing else. + +use std::fmt; + +// =========================================================================== +// Configuration +// =========================================================================== + +/// Opt-in switch. Absent or `0` means this whole module is inert and decode +/// takes the existing `replay()` + host-sample path. +pub const ENV_ENABLE: &str = "ARC_GRAPH_DEVICE_LOOP"; +/// Steps enqueued per burst before the host drains. Default [`DEFAULT_BURST`]. +pub const ENV_BURST: &str = "ARC_GRAPH_DEVICE_LOOP_BURST"; + +/// Steps launched back-to-back before the host drains the ring. +/// +/// The whole burst is enqueued without any host round trip, so this is also +/// how far the GPU may run ahead of the host's `Sequence` bookkeeping — and +/// therefore the worst-case detection latency, in steps, for a stop string, an +/// EOS, or an async fault. 4 keeps the wasted-work tail small while still +/// amortising the drain over four launches. +pub const DEFAULT_BURST: usize = 4; + +/// Tokens the pinned ring holds per row. Must exceed `burst`, or a burst could +/// overwrite tokens the host has not read yet. +pub const DEFAULT_RING_SIZE: usize = 1024; + +/// Host spins between `cudaStreamQuery` calls while draining. Pure host +/// memory polls are ~free; the driver query costs a syscall-ish amount, so it +/// is not worth doing on every iteration. +pub const DEFAULT_POLL_EVERY: u32 = 64; + +/// Emit [`DeviceDecodeLoop::status_line`] every this many bursts. The counter +/// it carries (`launches`) is the only proof the loop actually ran. +pub const STATUS_EVERY_BURSTS: u64 = 64; + +/// Liveness backstop: give up after this many spins even if the stream never +/// reports an error. Only trips on a kernel that hangs without faulting. +pub const DEFAULT_MAX_SPINS: u64 = 200_000_000; + +/// Is the device decode loop switched on? +/// +/// Default **off**. `ARC_GRAPH_DEVICE_LOOP=1` opts in; `0`, `false` and `off` +/// are explicit refusals so an operator can pin it off in an environment that +/// sets it upstream. +pub fn device_loop_enabled() -> bool { + match std::env::var(ENV_ENABLE) { + Ok(v) => { + let v = v.trim().to_ascii_lowercase(); + !(v.is_empty() || v == "0" || v == "false" || v == "off" || v == "no") + } + Err(_) => false, + } +} + +/// Burst length from [`ENV_BURST`], clamped to at least 1. +pub fn burst_from_env() -> usize { + std::env::var(ENV_BURST) + .ok() + .and_then(|v| v.trim().parse::().ok()) + .filter(|v| *v > 0) + .unwrap_or(DEFAULT_BURST) +} + +/// Everything the driver needs that does not change from burst to burst. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceLoopConfig { + /// Steps per burst. Also the run-ahead bound and the detection latency. + pub burst: usize, + /// Tokens per row in the pinned ring. + pub ring_size: usize, + /// Host spins between stream queries while draining. + pub poll_every: u32, + /// Spin budget before declaring the stream stalled. + pub max_spins: u64, + /// Stop the burst's token stream here. `< 0` disables EOS truncation. + pub eos_token_id: i32, + /// Vocabulary size; the commit kernel range-checks against it. + pub vocab: usize, + /// Exclusive upper bound on the device position. The fixed-capacity graph + /// decode arm reads a constant `capacity`-wide window and writes KV at + /// `position`; advancing past it is an out-of-bounds write, so both the + /// driver and the kernel refuse to. + pub position_limit: u32, +} + +impl DeviceLoopConfig { + /// Config with the shipped defaults for everything but the model facts. + pub fn new(vocab: usize, eos_token_id: i32, position_limit: u32) -> Self { + Self { + burst: burst_from_env(), + ring_size: DEFAULT_RING_SIZE, + poll_every: DEFAULT_POLL_EVERY, + max_spins: DEFAULT_MAX_SPINS, + eos_token_id, + vocab, + position_limit, + } + } + + /// Reject a config that cannot be run safely, rather than discovering it + /// as corruption at step `ring_size + 1`. + pub fn validate(&self) -> Result<(), DeviceLoopError> { + if self.burst == 0 { + return Err(DeviceLoopError::BadConfig("burst must be non-zero")); + } + if self.vocab == 0 { + return Err(DeviceLoopError::BadConfig("vocab must be non-zero")); + } + if self.poll_every == 0 { + // The drain does `spins % poll_every`; zero would divide by zero. + return Err(DeviceLoopError::BadConfig("poll_every must be non-zero")); + } + if self.ring_size <= self.burst { + // A burst longer than the ring overwrites its own unread tokens. + return Err(DeviceLoopError::BadConfig( + "ring_size must exceed burst, or a burst overwrites unread tokens", + )); + } + Ok(()) + } + + /// Steps this burst may launch given how far the sequence has already got. + /// + /// Clamped by the position limit so the burst can never walk the device + /// position past the fixed-capacity window. Returns 0 when the sequence has + /// no room left, which the caller must treat as "fall back to eager", not + /// as "launch zero steps and carry on". + pub fn plan_steps(&self, current_position: u32, max_new_tokens: usize) -> usize { + let room = self.position_limit.saturating_sub(current_position) as usize; + self.burst.min(room).min(max_new_tokens) + } +} + +// =========================================================================== +// Errors +// =========================================================================== + +/// Codes the commit kernel latches into the pinned fault word. Kept in sync +/// with `arc_graph_step_commit` in `src/cuda/decode_loop.cu`. +pub const FAULT_NONE: i32 = 0; +/// Sampled token id was outside `[0, vocab)`. +pub const FAULT_TOKEN_RANGE: i32 = 1; +/// Advancing the position would leave the fixed-capacity window. +pub const FAULT_POSITION_LIMIT: i32 = 2; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DeviceLoopError { + BadConfig(&'static str), + /// A CUDA call reported a failure. `at` names the call site. + Cuda { + at: &'static str, + code: u32, + }, + /// The stream retired every enqueued step but the ring is short. The work + /// completed without producing the tokens it was launched to produce. + LostTokens { + expected: i32, + got: i32, + }, + /// The commit kernel refused to write and latched a code. + DeviceFault(i32), + /// The write head ran more than `ring_size` ahead of the read cursor, so + /// unread tokens were overwritten. Reported rather than silently returning + /// whatever survived. + RingOverrun { + pending: i64, + ring_size: usize, + }, + /// Spin budget exhausted with the stream still busy and no error reported. + Stalled { + spins: u64, + }, +} + +impl fmt::Display for DeviceLoopError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::BadConfig(m) => write!(f, "device decode loop misconfigured: {m}"), + Self::Cuda { at, code } => { + write!(f, "device decode loop: {at} failed: CUDA code {code}") + } + Self::LostTokens { expected, got } => write!( + f, + "device decode loop: the stream retired every enqueued step but only {got} of \ + {expected} tokens reached the ring. The graph ran and did not commit its tokens \ + — treat this as an async fault, not as slow progress" + ), + Self::DeviceFault(code) => { + let what = match *code { + FAULT_TOKEN_RANGE => "sampled token id outside [0, vocab)", + FAULT_POSITION_LIMIT => { + "advancing the position would leave the fixed-capacity KV window" + } + _ => "unrecognised code", + }; + write!( + f, + "device decode loop: commit kernel refused to write (code {code}: {what}). No \ + token was scattered into the graph input buffer, so the next launch would \ + have read a stale id" + ) + } + Self::RingOverrun { pending, ring_size } => write!( + f, + "device decode loop: {pending} tokens pending in a {ring_size}-slot ring — the GPU \ + overwrote tokens the host had not read. Tokens were lost; this is a correctness \ + stop, not a throughput warning" + ), + Self::Stalled { spins } => write!( + f, + "device decode loop: stream still busy after {spins} spins with no error reported. \ + The graph is hung rather than faulted" + ), + } + } +} + +impl std::error::Error for DeviceLoopError {} + +impl From for candle_core::Error { + fn from(e: DeviceLoopError) -> Self { + candle_core::Error::Msg(e.to_string()) + } +} + +// =========================================================================== +// The device interface +// =========================================================================== + +/// Result of a non-blocking stream probe. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamState { + /// Every enqueued step has retired (`cudaSuccess`). + Idle, + /// Work still outstanding (`cudaErrorNotReady`). + Busy, +} + +/// The device operations a burst needs. +/// +/// **There is deliberately no `synchronize` method.** The driver is structurally +/// incapable of blocking the stream, which is the property this whole module +/// exists to guarantee — it cannot be regressed by an edit that "just adds a +/// sync to be safe", because there is nothing to call. +pub trait DeviceOps { + /// One `cuGraphLaunch` of the captured decode graph. Must not block. + fn launch_graph(&mut self) -> Result<(), DeviceLoopError>; + + /// Sample on device: graph logits -> the `i32[batch]` token buffer. Must + /// not block and must not copy the token to the host. + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError>; + + /// Commit on device: token -> pinned graph input ids, position += 1, and + /// publish into the pinned ring. Must not block. + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError>; + + /// Non-blocking stream probe. `Err` for a sticky async fault. + fn poll_stream(&mut self) -> Result; + + /// Monotonic write head for `row`, read from pinned mapped host memory. + /// A plain volatile load — no CUDA API call. + fn ring_head(&self, row: usize) -> i32; + + /// Token at absolute index `index` for `row` (the impl applies the modulo). + fn ring_slot(&self, row: usize, index: i32) -> i32; + + /// The sticky fault word, read from pinned mapped host memory. + fn device_fault(&self) -> i32; + + /// Rows the ring carries. + fn batch(&self) -> usize; +} + +// =========================================================================== +// The driver +// =========================================================================== + +/// What one burst produced. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BurstOutcome { + /// Tokens per row, in generation order, truncated at and including EOS. + pub tokens: Vec>, + /// True when some row emitted EOS inside this burst. + pub hit_eos: bool, + /// Steps actually launched. When `hit_eos` is true this exceeds the number + /// of tokens returned: the burst was already enqueued and the GPU computed + /// the tail before the host could know. That work is discarded, and this + /// field is how much of it there was. + pub steps_launched: usize, +} + +/// Drives bursts against a [`DeviceOps`], owning the ring read cursors. +pub struct BurstDriver { + cfg: DeviceLoopConfig, + /// Absolute index of the next unread ring slot, per row. + cursor: Vec, + launches: u64, + bursts: u64, + spins: u64, +} + +impl BurstDriver { + pub fn new(cfg: DeviceLoopConfig, batch: usize) -> Result { + cfg.validate()?; + Ok(Self { + cfg, + cursor: vec![0; batch.max(1)], + launches: 0, + bursts: 0, + spins: 0, + }) + } + + pub fn config(&self) -> &DeviceLoopConfig { + &self.cfg + } + + /// Total `cuGraphLaunch` calls issued. With `bursts()` this is the honest + /// answer to "did the device loop do anything": `launches == steps` and + /// zero stream synchronizes is the entire claim. + pub fn launches(&self) -> u64 { + self.launches + } + + pub fn bursts(&self) -> u64 { + self.bursts + } + + /// Host spins spent waiting for tokens to land. Rises when the host + /// out-runs the GPU, which is the healthy direction. + pub fn spins(&self) -> u64 { + self.spins + } + + /// The read cursor is per-sequence state. Reset it when the ring is reset + /// (a new request reuses the buffers). + pub fn reset_cursors(&mut self) { + for c in self.cursor.iter_mut() { + *c = 0; + } + } + + /// Enqueue `steps` decode steps and drain their tokens. + /// + /// Issues exactly `steps` `cuGraphLaunch` calls and **zero** stream + /// synchronizations. The drain is a host-memory poll on pinned mapped + /// memory plus a periodic non-blocking `cudaStreamQuery`. + pub fn run_burst( + &mut self, + ops: &mut D, + steps: usize, + ) -> Result { + if steps == 0 { + return Err(DeviceLoopError::BadConfig( + "run_burst called with zero steps; the caller must fall back to eager instead", + )); + } + if steps >= self.cfg.ring_size { + return Err(DeviceLoopError::BadConfig( + "burst is not smaller than the ring; it would overwrite its own tokens", + )); + } + + // ── enqueue ────────────────────────────────────────────────────── + // No blocking call in this loop, by construction: `DeviceOps` has no + // way to synchronize. `poll_stream` is the async-fault probe and it + // returns immediately whatever the stream is doing. + for _ in 0..steps { + ops.launch_graph()?; + ops.sample_on_device()?; + ops.commit_on_device()?; + self.launches += 1; + // Surface a sticky fault from an EARLIER step while later steps are + // already in flight. Costs nothing and shortens detection latency + // from "end of burst" to "next step". + ops.poll_stream()?; + } + self.bursts += 1; + + // ── drain ──────────────────────────────────────────────────────── + let batch = ops.batch().min(self.cursor.len()); + let want: Vec = self.cursor[..batch] + .iter() + .map(|c| c.saturating_add(steps as i32)) + .collect(); + for (row, &target) in want.iter().enumerate() { + self.await_row(ops, row, target)?; + } + + // The commit kernel latches its refusal instead of writing, so this + // must be checked even when every token arrived: a fault on step 3 of 4 + // means steps 4.. read a stale input id. + let fault = ops.device_fault(); + if fault != FAULT_NONE { + return Err(DeviceLoopError::DeviceFault(fault)); + } + + // ── collect ────────────────────────────────────────────────────── + let mut tokens = Vec::with_capacity(batch); + let mut hit_eos = false; + for (row, &target) in want.iter().enumerate() { + let mut row_tokens = Vec::with_capacity(steps); + let mut idx = self.cursor[row]; + while idx < target { + let tok = ops.ring_slot(row, idx); + idx += 1; + let is_eos = self.cfg.eos_token_id >= 0 && tok == self.cfg.eos_token_id; + row_tokens.push(tok); + if is_eos { + hit_eos = true; + break; + } + } + // The cursor advances over the WHOLE burst even when EOS truncated + // the returned tokens. Those slots were written; leaving the cursor + // behind would re-read them as if they were the next burst's. + self.cursor[row] = target; + tokens.push(row_tokens); + } + + Ok(BurstOutcome { + tokens, + hit_eos, + steps_launched: steps, + }) + } + + /// Wait for `row`'s write head to reach `target`, without synchronizing. + fn await_row( + &mut self, + ops: &mut D, + row: usize, + target: i32, + ) -> Result<(), DeviceLoopError> { + let mut spins: u64 = 0; + loop { + let head = ops.ring_head(row); + if head >= target { + // Overrun check: more unread tokens than the ring can hold + // means the GPU lapped the host and older tokens are gone. + let pending = (head as i64) - (self.cursor[row] as i64); + if pending > self.cfg.ring_size as i64 { + return Err(DeviceLoopError::RingOverrun { + pending, + ring_size: self.cfg.ring_size, + }); + } + return Ok(()); + } + + if spins.is_multiple_of(self.cfg.poll_every as u64) { + match ops.poll_stream()? { + StreamState::Busy => {} + StreamState::Idle => { + // Every enqueued step retired. Re-read once: the head + // store and the stream's retirement are not ordered + // against each other from the host's point of view. + let head = ops.ring_head(row); + if head >= target { + return Ok(()); + } + return Err(DeviceLoopError::LostTokens { + expected: target, + got: head, + }); + } + } + } + + spins += 1; + self.spins += 1; + if spins >= self.cfg.max_spins { + return Err(DeviceLoopError::Stalled { spins }); + } + std::hint::spin_loop(); + } + } +} + +// =========================================================================== +// CUDA implementation +// =========================================================================== + +#[cfg(feature = "cuda")] +mod cuda_impl { + use super::*; + use crate::ffi::*; + use crate::sampling_cpu::SamplingConfig; + use crate::sampling_cuda::CudaSampler; + use candle_core::cuda::cudarc::driver::sys::CUstream; + use candle_core::{DType, Device, Tensor}; + + /// Device address of a buffer the captured graph reads. + /// + /// Refuses a non-contiguous tensor rather than silently taking the address + /// of the copy `contiguous()` would make — that copy is not the address the + /// graph baked, so writing to it would look fine and change nothing. + fn pinned_ptr(t: &Tensor, what: &str) -> candle_core::Result { + if !t.is_contiguous() { + candle_core::bail!( + "device decode loop: the {what} buffer is not contiguous, so its address is not \ + the one the captured graph baked. Refusing rather than writing into a copy." + ); + } + crate::weights::tensor_device_ptr(t) + } + + /// `cudaErrorNotReady` — the stream still has outstanding work. Not a + /// failure, and the only non-success code `cudaStreamQuery` returns that + /// is not a fault. + const CUDA_ERROR_NOT_READY: u32 = 600; + + /// Pinned, mapped host memory the GPU writes and the host reads with plain + /// volatile loads. Freed with `cudaFreeHost` on drop. + struct PinnedRing { + tokens: *mut i32, + head: *mut i32, + fault: *mut i32, + batch: usize, + ring_size: usize, + } + + impl PinnedRing { + fn new(batch: usize, ring_size: usize) -> Result { + let mut tokens: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut head: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut fault: *mut std::ffi::c_void = std::ptr::null_mut(); + unsafe { + let s = cudaHostAlloc(&mut tokens, batch * ring_size * 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(ring tokens)", + code: s, + }); + } + let s = cudaHostAlloc(&mut head, batch * 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + cudaFreeHost(tokens); + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(ring head)", + code: s, + }); + } + let s = cudaHostAlloc(&mut fault, 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + cudaFreeHost(tokens); + cudaFreeHost(head); + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(fault word)", + code: s, + }); + } + std::ptr::write_bytes(tokens as *mut u8, 0, batch * ring_size * 4); + std::ptr::write_bytes(head as *mut u8, 0, batch * 4); + std::ptr::write_bytes(fault as *mut u8, 0, 4); + } + Ok(Self { + tokens: tokens as *mut i32, + head: head as *mut i32, + fault: fault as *mut i32, + batch, + ring_size, + }) + } + + fn reset(&self) { + unsafe { + std::ptr::write_bytes(self.head as *mut u8, 0, self.batch * 4); + std::ptr::write_bytes(self.fault as *mut u8, 0, 4); + } + } + } + + impl Drop for PinnedRing { + fn drop(&mut self) { + unsafe { + if !self.tokens.is_null() { + cudaFreeHost(self.tokens as *mut _); + } + if !self.head.is_null() { + cudaFreeHost(self.head as *mut _); + } + if !self.fault.is_null() { + cudaFreeHost(self.fault as *mut _); + } + } + } + } + + /// The real [`DeviceOps`]: sampler + commit kernel on the graph's stream. + /// + /// Holds `&mut` to nothing outside itself except through raw pointers taken + /// once at construction, because every buffer it touches must be at a + /// **stable device address** — the captured graph baked those addresses and + /// a reallocation silently invalidates the graph. + pub struct CudaDeviceOps { + stream: CUstream, + sampler: CudaSampler, + sampling_cfg: SamplingConfig, + /// The graph's output logits tensor, reshaped to `[batch, vocab]`. + logits: Tensor, + /// `i32[batch]`, written by the sampler, read by the commit kernel. + token_ids: Tensor, + /// The PINNED `U32` token-id buffer the captured graph reads. + input_ids_ptr: *mut u32, + /// The PINNED `U32` position buffer, or null when positions are owned + /// elsewhere (e.g. by a device-side RoPE position path). + positions_ptr: *mut u32, + ring: PinnedRing, + cfg: DeviceLoopConfig, + /// `cuGraphLaunch` closure state: the exec handle to replay. + exec: CUgraphExec, + } + + impl CudaDeviceOps { + /// * `logits` — the captured graph's output tensor. Its storage address + /// is graph-owned and stable; the sampler reads it in place. + /// * `input_ids_buf` / `positions_buf` — the *pinned* graph input + /// buffers (`layers::GRAPH_MODE_INPUT_IDS` / `GRAPH_MODE_POSITIONS`). + /// These are the addresses the captured graph reads, which is why a + /// device-side write into them closes the loop with no host copy. + #[allow(clippy::too_many_arguments)] + pub fn new( + device: &Device, + exec: CUgraphExec, + stream: CUstream, + sampler: CudaSampler, + sampling_cfg: SamplingConfig, + logits: &Tensor, + input_ids_buf: &Tensor, + positions_buf: Option<&Tensor>, + cfg: DeviceLoopConfig, + ) -> candle_core::Result { + cfg.validate()?; + let batch = sampler.batch(); + + // `[1, 1, vocab]` (or `[1, vocab]`) -> `[batch, vocab]`. Reshaping a + // contiguous tensor is metadata-only, so this does not move the + // graph-owned storage the sampler is about to read. + let logits = logits.reshape((batch, sampler.vocab()))?; + if logits.dtype() != sampler.dtype() { + candle_core::bail!( + "device decode loop: graph logits dtype {:?} != sampler dtype {:?}", + logits.dtype(), + sampler.dtype() + ); + } + + if input_ids_buf.dtype() != DType::U32 { + candle_core::bail!( + "device decode loop: the graph input-ids buffer must be U32 (got {:?}); the \ + commit kernel writes it as uint32_t", + input_ids_buf.dtype() + ); + } + if input_ids_buf.elem_count() < batch { + candle_core::bail!( + "device decode loop: input-ids buffer holds {} elements, need {batch}", + input_ids_buf.elem_count() + ); + } + if let Some(p) = positions_buf { + if p.dtype() != DType::U32 { + candle_core::bail!( + "device decode loop: the graph positions buffer must be U32 (got {:?})", + p.dtype() + ); + } + if p.elem_count() < batch { + candle_core::bail!( + "device decode loop: positions buffer holds {} elements, need {batch}", + p.elem_count() + ); + } + } + + // I32 via `from_vec`: `Tensor::zeros` on I32 is not universally + // supported across backends, the same reason `CudaSampler::new` + // builds `keep_idx_scratch` this way. + let token_ids = Tensor::from_vec(vec![0i32; batch], batch, device)?; + + let input_ids_ptr = pinned_ptr(input_ids_buf, "graph input-ids")? as *mut u32; + let positions_ptr = match positions_buf { + Some(p) => pinned_ptr(p, "graph positions")? as *mut u32, + None => std::ptr::null_mut(), + }; + let ring = PinnedRing::new(batch, cfg.ring_size)?; + + Ok(Self { + stream, + sampler, + sampling_cfg, + logits, + token_ids, + input_ids_ptr, + positions_ptr, + ring, + cfg, + exec, + }) + } + + /// Zero the ring head and fault word. Call when a new sequence starts + /// reusing these buffers; pair it with [`BurstDriver::reset_cursors`]. + pub fn reset_ring(&self) { + self.ring.reset(); + } + } + + impl DeviceOps for CudaDeviceOps { + fn launch_graph(&mut self) -> Result<(), DeviceLoopError> { + // The launch and nothing else. `graph.rs::replay` follows this with + // `cudaStreamSynchronize`; the entire point here is that we do not. + let s = unsafe { cuGraphLaunch(self.exec, self.stream) }; + if s != CUDA_SUCCESS { + return Err(DeviceLoopError::Cuda { + at: "cuGraphLaunch", + code: s, + }); + } + Ok(()) + } + + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.sampler + .sample(&self.logits, None, self.sampling_cfg, &self.token_ids) + .map_err(|e| { + tracing::error!("device decode loop: on-device sample failed: {e}"); + DeviceLoopError::Cuda { + at: "CudaSampler::sample", + code: u32::MAX, + } + }) + } + + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError> { + let token_ptr = match crate::weights::tensor_device_ptr(&self.token_ids) { + Ok(p) => p as *const i32, + Err(e) => { + tracing::error!("device decode loop: token buffer pointer unavailable: {e}"); + return Err(DeviceLoopError::Cuda { + at: "tensor_device_ptr(token_ids)", + code: u32::MAX, + }); + } + }; + unsafe { + arc_launch_graph_step_commit( + token_ptr, + self.input_ids_ptr, + self.positions_ptr, + self.ring.tokens, + self.ring.head, + self.ring.fault, + self.cfg.ring_size as i32, + self.cfg.vocab as i32, + self.cfg.position_limit as i32, + self.ring.batch as i32, + self.stream, + ); + } + Ok(()) + } + + fn poll_stream(&mut self) -> Result { + // Non-blocking. NVIDIA's reference for cudaStreamQuery: "Note that + // this function may also return error codes from previous, + // asynchronous launches." That is what makes a sticky fault from an + // earlier graph step visible here without ever blocking. + let s = unsafe { cudaStreamQuery(self.stream) }; + match s { + CUDA_SUCCESS => Ok(StreamState::Idle), + CUDA_ERROR_NOT_READY => Ok(StreamState::Busy), + code => Err(DeviceLoopError::Cuda { + at: "cudaStreamQuery (async fault from an earlier graph step)", + code, + }), + } + } + + fn ring_head(&self, row: usize) -> i32 { + unsafe { std::ptr::read_volatile(self.ring.head.add(row)) } + } + + fn ring_slot(&self, row: usize, index: i32) -> i32 { + let slot = + row * self.ring.ring_size + (index.rem_euclid(self.ring.ring_size as i32) as usize); + unsafe { std::ptr::read_volatile(self.ring.tokens.add(slot)) } + } + + fn device_fault(&self) -> i32 { + unsafe { std::ptr::read_volatile(self.ring.fault) } + } + + fn batch(&self) -> usize { + self.ring.batch + } + } + + /// The whole device loop behind one handle: sampler, commit kernel, pinned + /// ring and burst driver. + /// + /// Built once per captured graph and kept alive across decode steps, + /// because every buffer it holds a pointer to must stay at the address the + /// graph baked. Dropping and rebuilding it mid-generation would hand the + /// commit kernel a fresh `token_ids` allocation while the graph still reads + /// the old input-ids address — so callers keep it, they do not recreate it. + pub struct DeviceDecodeLoop { + ops: CudaDeviceOps, + driver: BurstDriver, + } + + // Same reasoning as `CudaGraphRunner` (`graph.rs:114`): the raw pointers are + // CUDA handles and pinned host allocations, both valid from any thread, and + // the pipeline that owns this is held behind a mutex. Without these the + // `NormalPipeline` field would make the whole pipeline non-`Send`. + unsafe impl Send for CudaDeviceOps {} + unsafe impl Sync for CudaDeviceOps {} + unsafe impl Send for DeviceDecodeLoop {} + unsafe impl Sync for DeviceDecodeLoop {} + + impl DeviceDecodeLoop { + #[allow(clippy::too_many_arguments)] + pub fn new( + device: &Device, + exec: CUgraphExec, + stream: CUstream, + logits: &Tensor, + input_ids_buf: &Tensor, + positions_buf: Option<&Tensor>, + cfg: DeviceLoopConfig, + ) -> candle_core::Result { + let batch = 1usize; + let vocab = cfg.vocab; + // Greedy, always. [`admit`] refuses every non-greedy sequence, so + // taking the mode as a parameter would only create a way for a + // caller to contradict the guard. The seed is irrelevant under + // greedy (pure argmax draws no random numbers) and is fixed so the + // sampler's RNG buffer is deterministic if the mode ever widens. + let sampling_cfg = SamplingConfig::greedy(); + let sampler = CudaSampler::new(device, batch, vocab, logits.dtype(), 0)?; + let ops = CudaDeviceOps::new( + device, + exec, + stream, + sampler, + sampling_cfg, + logits, + input_ids_buf, + positions_buf, + cfg, + )?; + let driver = BurstDriver::new(cfg, batch)?; + Ok(Self { ops, driver }) + } + + /// Enqueue `steps` decode steps and return their tokens. + /// + /// Exactly `steps` `cuGraphLaunch` calls, zero `cudaStreamSynchronize`. + pub fn run(&mut self, steps: usize) -> candle_core::Result { + let out = self.driver.run_burst(&mut self.ops, steps)?; + // D18: a subsystem that only announces itself at startup cannot be + // told apart from one that started and then did nothing. Emit the + // running counters periodically so `launches` is observable in a + // live run, not just in a debugger. + if self.driver.bursts().is_multiple_of(STATUS_EVERY_BURSTS) { + tracing::info!("{}", self.status_line()); + } + Ok(out) + } + + /// Start a fresh sequence on the same buffers. + pub fn reset(&mut self) { + self.ops.reset_ring(); + self.driver.reset_cursors(); + clear_pending_tokens(); + } + + /// `cuGraphLaunch` calls issued so far — the numerator of the claim. + pub fn launches(&self) -> u64 { + self.driver.launches() + } + + pub fn bursts(&self) -> u64 { + self.driver.bursts() + } + + /// The one line that separates a working device loop from an inert one + /// (D18). `launches=0` is a subsystem that did nothing and says so. + pub fn status_line(&self) -> String { + format!( + "ARCGRAPH DEVICE LOOP: launches={} bursts={} burst_len={} host_syncs=0 \ + host_spins={} pending={}", + self.driver.launches(), + self.driver.bursts(), + self.driver.config().burst, + self.driver.spins(), + pending_token_count(), + ) + } + } +} + +#[cfg(feature = "cuda")] +pub use cuda_impl::{CudaDeviceOps, DeviceDecodeLoop}; + +// =========================================================================== +// Admission +// =========================================================================== + +/// Why the device loop refused a request. +/// +/// Refusing is the normal, safe outcome — the caller falls back to the existing +/// `replay()` + host-sample path. Each variant names something the device +/// sampler cannot reproduce *exactly*, because a token that differs from the +/// host path is a silent behaviour change, not an optimisation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Refusal { + /// `ARC_GRAPH_DEVICE_LOOP` is not set. The default. + NotEnabled, + /// Only batch 1 is wired: the burst holds the scheduler for its whole + /// length, so it must not straddle sequences with different stop conditions. + NotBatchOne, + /// Prefill, not decode. + NotDecodeStep, + /// No captured graph to replay, or its output is not yet trusted. + GraphNotReady, + /// Non-greedy sampling. The device sampler runs Splitmix64 per row while + /// the host runs Isaac64, so for any stochastic config the two draw + /// *different, equally valid* tokens — which silently breaks seeded + /// reproducibility. Greedy is deterministic, so device and host agree + /// exactly, which is also what makes the path verifiable. + NotGreedy, + /// The caller wants logprobs, which the device sampler does not return. + WantsLogprobs, + /// A logits processor, grammar or penalty must see the logits on the host. + HostSideLogitsWork, + /// No room left in the fixed-capacity KV window for even one step. + NoPositionRoom, +} + +impl fmt::Display for Refusal { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let s = match self { + Self::NotEnabled => "ARC_GRAPH_DEVICE_LOOP is not set", + Self::NotBatchOne => "batch size is not 1", + Self::NotDecodeStep => "this is a prefill step, not decode", + Self::GraphNotReady => "no captured graph, or its output is not yet trusted", + Self::NotGreedy => { + "sampling is not greedy; the device and host RNGs differ, so tokens would \ + diverge from the seeded host path" + } + Self::WantsLogprobs => { + "the request asks for logprobs, which the device sampler \ + does not produce" + } + Self::HostSideLogitsWork => { + "a logits processor, grammar or repetition penalty needs the logits on the host" + } + Self::NoPositionRoom => "no room left in the fixed-capacity KV window", + }; + write!(f, "{s}") + } +} + +/// Everything the admission check needs, gathered by the caller. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AdmissionFacts { + pub batch_size: usize, + pub seq_len: usize, + pub graph_ready: bool, + pub greedy: bool, + pub return_logprobs: bool, + pub host_side_logits_work: bool, + pub current_position: u32, +} + +/// May the device loop drive this step? +/// +/// Kept separate from the CUDA plumbing, and not feature-gated, so the decision +/// is unit-tested on any host. Every refusal is a fall-back to the existing +/// path, never an error. +pub fn admit( + facts: &AdmissionFacts, + cfg: &DeviceLoopConfig, + enabled: bool, +) -> Result { + if !enabled { + return Err(Refusal::NotEnabled); + } + if facts.seq_len != 1 { + return Err(Refusal::NotDecodeStep); + } + if facts.batch_size != 1 { + return Err(Refusal::NotBatchOne); + } + if !facts.graph_ready { + return Err(Refusal::GraphNotReady); + } + if facts.return_logprobs { + return Err(Refusal::WantsLogprobs); + } + if facts.host_side_logits_work { + return Err(Refusal::HostSideLogitsWork); + } + if !facts.greedy { + return Err(Refusal::NotGreedy); + } + let steps = cfg.plan_steps(facts.current_position, cfg.burst); + if steps == 0 { + return Err(Refusal::NoPositionRoom); + } + Ok(steps) +} + +// =========================================================================== +// Host-side token handoff +// =========================================================================== + +// Tokens a burst produced that the engine has not consumed yet. +// +// A burst computes `steps` tokens in one go, but the engine's `Sequence` +// bookkeeping — stop strings, streaming, EOS — is per token. So the burst +// parks its tokens here and the engine drains them one at a time on the +// following steps, during which no GPU work is launched at all. +// +// Thread-local because at batch 1 `sample_and_add_toks` runs the sampler +// inline on the engine thread (`use_async_pool = seqs_len > 1` is false), the +// same thread that ran the forward. The `GRAPH_MODE_*` graph inputs in +// `mistralrs-core/src/layers.rs` are scoped the same way for the same reason. +std::thread_local! { + static PENDING_TOKENS: std::cell::RefCell> = + const { std::cell::RefCell::new(std::collections::VecDeque::new()) }; +} + +/// Park a burst's tokens for the engine to drain. +/// +/// Negative ids are dropped rather than cast: the commit kernel range-checks +/// before it writes, so a negative here would mean the ring was read wrong, and +/// wrapping it into a huge `u32` would index the tokenizer out of bounds. +pub fn push_pending_tokens(tokens: &[i32]) { + PENDING_TOKENS.with(|q| { + let mut q = q.borrow_mut(); + for &t in tokens { + if t >= 0 { + q.push_back(t as u32); + } else { + tracing::error!( + "device decode loop: dropping negative token id {t} from the ring; the \ + commit kernel range-checks before writing, so this means the ring was \ + misread rather than that the model produced it" + ); + } + } + }) +} + +/// Take the next token the burst already produced, if any. +pub fn take_pending_token() -> Option { + PENDING_TOKENS.with(|q| q.borrow_mut().pop_front()) +} + +/// Tokens still parked. +pub fn pending_token_count() -> usize { + PENDING_TOKENS.with(|q| q.borrow().len()) +} + +// Whether the sequence currently being decoded is one the device loop may +// drive. +// +// `Pipeline::forward_inputs` decides whether to run a burst, but it is handed +// only `inputs` and `return_raw_logits` — it never sees the `Sequence`, so it +// cannot ask whether sampling is greedy or whether logprobs were requested. +// `sample_sequence` does see it. So eligibility is published there and read by +// the *next* forward. +// +// The one-step lag is safe in the only direction that matters: a sequence's +// sampling parameters are fixed for its lifetime, and the flag starts `false`, +// so the first decode step always falls back and the loop can only ever engage +// after a step has confirmed the parameters. It can never engage for a +// sequence it has not seen. +std::thread_local! { + static DEVICE_LOOP_ELIGIBLE: std::cell::Cell = const { std::cell::Cell::new(false) }; +} + +/// Publish whether the sequence being sampled may be driven by the device loop. +pub fn set_device_loop_eligible(eligible: bool) { + DEVICE_LOOP_ELIGIBLE.with(|c| c.set(eligible)) +} + +/// Did the last sampled sequence qualify? `false` until one has. +pub fn device_loop_eligible() -> bool { + DEVICE_LOOP_ELIGIBLE.with(|c| c.get()) +} + +/// Drop parked tokens. +/// +/// MUST be called when a sequence finishes, is preempted, or the path falls +/// back — a token left here would be handed to the *next* sequence as if it +/// were its own. That is the "+1-step divergence" failure mode with a different +/// name, so it gets an explicit call rather than a lifetime. +pub fn clear_pending_tokens() { + PENDING_TOKENS.with(|q| q.borrow_mut().clear()) +} + +/// Stand down completely: forget parked tokens and mark the current sequence +/// ineligible. Call on fall-back, on sequence completion, and on any error, so +/// the next sequence starts from a clean slate. +pub fn stand_down() { + clear_pending_tokens(); + set_device_loop_eligible(false); +} + +/// Process-wide latch. Once the device loop has failed, it stays off. +/// +/// Without this a construction or burst failure would be retried on every +/// decode step forever — and because a refusal falls back to the *eager* +/// forward rather than to `replay()`, that would quietly make every step slower +/// than doing nothing at all. One failure, one warning, off for good. +static DEVICE_LOOP_KILLED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + +/// Latch the device loop off for the rest of the process, loudly and once. +pub fn kill_device_loop(reason: &str) { + if !DEVICE_LOOP_KILLED.swap(true, std::sync::atomic::Ordering::SeqCst) { + tracing::error!( + "ArcGraph device decode loop DISABLED for the rest of this process: {reason}. Decode \ + falls back to graph replay with a host sync and a host sample. This is a stop, not a \ + retry — a loop that fails once and is retried every step is slower than never having \ + engaged." + ); + } + stand_down(); +} + +/// Has the device loop been latched off? +pub fn device_loop_killed() -> bool { + DEVICE_LOOP_KILLED.load(std::sync::atomic::Ordering::SeqCst) +} + +#[cfg(feature = "cuda")] +extern "C" { + /// `arc_graph_step_commit` — `src/cuda/decode_loop.cu`. + /// + /// Scatters the device-sampled token into the pinned graph input buffer, + /// advances the device position, and publishes the token into the pinned + /// mapped ring. Range-checks before it writes anything, latching a code + /// into `fault` instead of committing a bad token. + pub fn arc_launch_graph_step_commit( + sampled: *const i32, + input_ids: *mut u32, + positions: *mut u32, + ring: *mut i32, + ring_head: *mut i32, + fault: *mut i32, + ring_size: i32, + vocab: i32, + position_limit: i32, + batch: i32, + stream: candle_core::cuda::cudarc::driver::sys::CUstream, + ); +} + +// =========================================================================== +// Tests +// =========================================================================== + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::VecDeque; + + /// A scripted fake GPU. + /// + /// It models the property that matters: work is enqueued without blocking + /// and lands *later*. `launch_graph` pushes the next scripted token into an + /// in-flight queue; `poll_stream` retires one in-flight step per call and + /// publishes its token into the ring. So a driver that expects tokens to be + /// present the instant it launches will not find them, and a driver that + /// never polls will spin forever — both real failure modes. + struct FakeGpu { + script: VecDeque, + in_flight: VecDeque, + ring: Vec, + head: i32, + fault: i32, + ring_size: usize, + launches: usize, + samples: usize, + commits: usize, + polls: usize, + /// Set to drop the commit instead of publishing — the "no-op + /// implementation" this test suite must not accept. + commit_is_noop: bool, + /// Injected async fault, surfaced from `poll_stream`. + fault_after: Option, + /// Report Idle immediately even with work outstanding. + lie_idle: bool, + /// A hung kernel: work is enqueued, the stream reports Busy forever and + /// nothing ever retires. No error is ever raised, so only the liveness + /// backstop can end it. + never_retires: bool, + } + + impl FakeGpu { + fn new(script: Vec, ring_size: usize) -> Self { + Self { + script: script.into(), + in_flight: VecDeque::new(), + ring: vec![0; ring_size], + head: 0, + fault: FAULT_NONE, + ring_size, + launches: 0, + samples: 0, + commits: 0, + polls: 0, + commit_is_noop: false, + fault_after: None, + lie_idle: false, + never_retires: false, + } + } + + fn retire_one(&mut self) { + if let Some(tok) = self.in_flight.pop_front() { + let slot = (self.head as usize) % self.ring_size; + self.ring[slot] = tok; + self.head += 1; + } + } + } + + impl DeviceOps for FakeGpu { + fn launch_graph(&mut self) -> Result<(), DeviceLoopError> { + self.launches += 1; + Ok(()) + } + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.samples += 1; + Ok(()) + } + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.commits += 1; + if self.commit_is_noop { + return Ok(()); + } + let tok = self.script.pop_front().unwrap_or(0); + self.in_flight.push_back(tok); + Ok(()) + } + fn poll_stream(&mut self) -> Result { + self.polls += 1; + if let Some(after) = self.fault_after { + if self.polls > after { + return Err(DeviceLoopError::Cuda { + at: "cudaStreamQuery (async fault from an earlier graph step)", + code: 700, + }); + } + } + if self.never_retires { + return Ok(StreamState::Busy); + } + self.retire_one(); + if self.lie_idle || self.in_flight.is_empty() { + Ok(StreamState::Idle) + } else { + Ok(StreamState::Busy) + } + } + fn ring_head(&self, _row: usize) -> i32 { + self.head + } + fn ring_slot(&self, _row: usize, index: i32) -> i32 { + self.ring[index.rem_euclid(self.ring_size as i32) as usize] + } + fn device_fault(&self) -> i32 { + self.fault + } + fn batch(&self) -> usize { + 1 + } + } + + fn cfg(burst: usize, ring: usize) -> DeviceLoopConfig { + DeviceLoopConfig { + burst, + ring_size: ring, + poll_every: 1, + max_spins: 10_000, + eos_token_id: 2, + vocab: 128, + position_limit: 4096, + } + } + + /// THE claim: N decode steps cost exactly N `cuGraphLaunch` calls, and the + /// tokens that come back are the ones the GPU produced. + /// + /// Fails on a no-op: with fewer launches the count is wrong, and with a + /// commit that never publishes the drain reports `LostTokens`. + #[test] + fn n_steps_cost_exactly_n_launches_and_return_the_gpu_tokens() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let out = d.run_burst(&mut gpu, 4).unwrap(); + + assert_eq!(gpu.launches, 4, "N steps must cost exactly N cuGraphLaunch"); + assert_eq!(gpu.samples, 4, "every step must sample on device"); + assert_eq!(gpu.commits, 4, "every step must commit on device"); + assert_eq!(out.tokens, vec![vec![11, 12, 13, 14]]); + assert_eq!(d.launches(), 4); + assert!(!out.hit_eos); + } + + /// The control the previous test needs: an implementation that launches and + /// samples but never publishes the token MUST fail, and must fail as + /// `LostTokens` rather than by spinning out. This is the "test that cannot + /// fail" guard — if the driver accepted a silent no-op commit, the whole + /// module would be decoration. + #[test] + fn a_commit_that_publishes_nothing_is_rejected_as_lost_tokens() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.commit_is_noop = true; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + assert!( + matches!( + err, + DeviceLoopError::LostTokens { + expected: 4, + got: 0 + } + ), + "expected LostTokens, got {err:?}" + ); + // It still launched: the failure is in the commit, and the error must + // not be confused with "the driver never ran". + assert_eq!(gpu.launches, 4); + } + + /// An async fault reported by the non-blocking query must propagate as a + /// CUDA error — not be swallowed, and not degrade into `Stalled`. + #[test] + fn async_fault_from_stream_query_fails_loudly_as_cuda() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.fault_after = Some(2); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + match err { + DeviceLoopError::Cuda { code, .. } => assert_eq!(code, 700), + other => panic!("async fault must surface as Cuda, got {other:?}"), + } + } + + /// Detection latency: the in-burst probe finds the fault before the burst + /// finishes enqueuing, so a fault on step 1 does not wait for step 4. + #[test] + fn in_burst_probe_detects_a_fault_before_the_burst_completes() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14, 15, 16, 17, 18], 32); + gpu.fault_after = Some(1); // second poll errors -> during step 2 + let mut d = BurstDriver::new(cfg(8, 32), 1).unwrap(); + assert!(d.run_burst(&mut gpu, 8).is_err()); + assert!( + gpu.launches < 8, + "the burst must stop early on a fault, launched {}", + gpu.launches + ); + } + + /// The no-timeout fault detector: stream reports every step retired, but + /// the ring is short. That is definitively wrong and must be an error, not + /// an infinite spin and not a silent short read. + #[test] + fn stream_idle_with_a_short_ring_is_a_fault_not_a_wait() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.commit_is_noop = true; + gpu.lie_idle = true; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + match d.run_burst(&mut gpu, 4).unwrap_err() { + DeviceLoopError::LostTokens { expected, got } => { + assert_eq!((expected, got), (4, 0)); + } + other => panic!("expected LostTokens, got {other:?}"), + } + assert!( + d.spins() < 100, + "idle-but-short must be detected immediately, not spun on ({} spins)", + d.spins() + ); + } + + /// The commit kernel's range check travels back in a pinned word. A driver + /// that forgets to read it returns `Ok` with plausible tokens — exactly the + /// silent-corruption class this repo keeps shipping. So: tokens all arrive, + /// fault word set, and the burst must still fail. + #[test] + fn a_latched_device_fault_fails_the_burst_even_though_every_token_arrived() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.fault = FAULT_TOKEN_RANGE; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + assert_eq!(err, DeviceLoopError::DeviceFault(FAULT_TOKEN_RANGE)); + assert_eq!(gpu.head, 4, "the tokens did land; the fault is orthogonal"); + } + + /// EOS truncates the returned tokens but not the launch count. The burst + /// was already enqueued, so the tail was computed and is discarded — the + /// outcome must report both numbers honestly. + #[test] + fn eos_truncates_tokens_but_the_wasted_steps_are_still_reported() { + let mut gpu = FakeGpu::new(vec![11, 12, 2, 14], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let out = d.run_burst(&mut gpu, 4).unwrap(); + assert_eq!(out.tokens, vec![vec![11, 12, 2]]); + assert!(out.hit_eos); + assert_eq!( + out.steps_launched, 4, + "the discarded tail must be visible, not rounded away" + ); + } + + /// Consecutive bursts must not re-read the previous burst's slots. The + /// cursor advances over the full burst even when EOS truncated the output. + #[test] + fn cursor_advances_over_the_whole_burst_across_consecutive_bursts() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14, 21, 22, 23, 24], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let a = d.run_burst(&mut gpu, 4).unwrap(); + let b = d.run_burst(&mut gpu, 4).unwrap(); + assert_eq!(a.tokens, vec![vec![11, 12, 13, 14]]); + assert_eq!(b.tokens, vec![vec![21, 22, 23, 24]]); + assert_eq!(gpu.launches, 8); + } + + /// Ring wraparound is arithmetic, and arithmetic is where silent corruption + /// lives. Drive past the ring size and check the tokens are still right. + #[test] + fn tokens_survive_ring_wraparound() { + let script: Vec = (100..140).collect(); + let mut gpu = FakeGpu::new(script.clone(), 8); + let mut d = BurstDriver::new(cfg(4, 8), 1).unwrap(); + let mut seen = Vec::new(); + for _ in 0..10 { + let out = d.run_burst(&mut gpu, 4).unwrap(); + seen.extend(out.tokens[0].iter().copied()); + } + assert_eq!(seen, script, "wraparound must not permute or drop tokens"); + } + + /// A burst that could lap the ring is refused up front rather than + /// discovered as scrambled tokens. + #[test] + fn a_burst_that_would_lap_the_ring_is_refused() { + assert!(matches!( + BurstDriver::new(cfg(16, 8), 1), + Err(DeviceLoopError::BadConfig(_)) + )); + let mut gpu = FakeGpu::new(vec![1; 32], 8); + let mut d = BurstDriver::new(cfg(4, 8), 1).unwrap(); + assert!(matches!( + d.run_burst(&mut gpu, 8), + Err(DeviceLoopError::BadConfig(_)) + )); + } + + /// A kernel that hangs *without* faulting reports Busy forever and lands + /// nothing. No CUDA error is ever raised and the stream never goes idle, so + /// neither of the two real detectors can fire — only the liveness backstop + /// can end the burst. Without it this is an unkillable spin in the engine. + #[test] + fn a_hung_stream_trips_the_liveness_backstop() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.never_retires = true; + let mut c = cfg(4, 16); + c.max_spins = 500; + let mut d = BurstDriver::new(c, 1).unwrap(); + match d.run_burst(&mut gpu, 4).unwrap_err() { + DeviceLoopError::Stalled { spins } => assert_eq!(spins, 500), + other => panic!("expected Stalled, got {other:?}"), + } + } + + /// `poll_every == 0` would divide by zero in the drain. Refuse it at + /// construction rather than panicking mid-burst. + #[test] + fn zero_poll_every_is_refused_rather_than_dividing_by_zero() { + let mut c = cfg(4, 16); + c.poll_every = 0; + assert!(matches!( + BurstDriver::new(c, 1), + Err(DeviceLoopError::BadConfig(_)) + )); + } + + /// `plan_steps` must never walk the device position out of the + /// fixed-capacity KV window, and must report 0 rather than a short burst + /// the caller might mistake for success. + #[test] + fn plan_steps_clamps_to_the_position_window_and_the_token_budget() { + let mut c = cfg(8, 64); + c.position_limit = 100; + assert_eq!(c.plan_steps(0, 1000), 8); + assert_eq!(c.plan_steps(95, 1000), 5, "must not pass position_limit"); + assert_eq!(c.plan_steps(100, 1000), 0, "no room left"); + assert_eq!( + c.plan_steps(120, 1000), + 0, + "already past: saturating, not wrapping" + ); + assert_eq!(c.plan_steps(0, 3), 3, "must not exceed the token budget"); + } + + /// Zero steps is a caller bug, not a quiet no-op that reports success. + #[test] + fn zero_steps_is_an_error_not_a_silent_success() { + let mut gpu = FakeGpu::new(vec![], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + assert!(matches!( + d.run_burst(&mut gpu, 0), + Err(DeviceLoopError::BadConfig(_)) + )); + assert_eq!(gpu.launches, 0); + } + + fn ok_facts() -> AdmissionFacts { + AdmissionFacts { + batch_size: 1, + seq_len: 1, + graph_ready: true, + greedy: true, + return_logprobs: false, + host_side_logits_work: false, + current_position: 10, + } + } + + /// Admission must refuse every case where the device sampler would produce + /// a token the host path would not. Each of these is a silent behaviour + /// change if it slips through, so each gets its own assertion rather than + /// one "returns Err" catch-all. + #[test] + fn admission_refuses_every_case_the_device_sampler_cannot_reproduce() { + let c = cfg(4, 64); + // The happy path first — otherwise a guard that refuses everything + // would pass every negative case below. + assert_eq!(admit(&ok_facts(), &c, true), Ok(4)); + + assert_eq!(admit(&ok_facts(), &c, false), Err(Refusal::NotEnabled)); + + let mut f = ok_facts(); + f.seq_len = 32; + assert_eq!(admit(&f, &c, true), Err(Refusal::NotDecodeStep)); + + let mut f = ok_facts(); + f.batch_size = 2; + assert_eq!(admit(&f, &c, true), Err(Refusal::NotBatchOne)); + + let mut f = ok_facts(); + f.graph_ready = false; + assert_eq!(admit(&f, &c, true), Err(Refusal::GraphNotReady)); + + let mut f = ok_facts(); + f.greedy = false; + assert_eq!( + admit(&f, &c, true), + Err(Refusal::NotGreedy), + "stochastic sampling must fall back: device Splitmix64 and host Isaac64 draw \ + different tokens, which breaks seeded reproducibility" + ); + + let mut f = ok_facts(); + f.return_logprobs = true; + assert_eq!(admit(&f, &c, true), Err(Refusal::WantsLogprobs)); + + let mut f = ok_facts(); + f.host_side_logits_work = true; + assert_eq!(admit(&f, &c, true), Err(Refusal::HostSideLogitsWork)); + } + + /// Admission must also clamp the burst to the KV window, and refuse rather + /// than return a zero-step "success" the caller could mistake for a launch. + #[test] + fn admission_clamps_the_burst_to_the_kv_window() { + let mut c = cfg(4, 64); + c.position_limit = 100; + let mut f = ok_facts(); + f.current_position = 98; + assert_eq!(admit(&f, &c, true), Ok(2), "clamped to the window"); + f.current_position = 100; + assert_eq!(admit(&f, &c, true), Err(Refusal::NoPositionRoom)); + } + + /// The parked tokens are handed out in generation order, exactly once. + /// A queue that re-served a token, or served them reversed, would produce + /// fluent-but-wrong text — invisible without this. + #[test] + fn pending_tokens_drain_in_order_exactly_once() { + clear_pending_tokens(); + push_pending_tokens(&[7, 8, 9]); + assert_eq!(pending_token_count(), 3); + assert_eq!(take_pending_token(), Some(7)); + assert_eq!(take_pending_token(), Some(8)); + assert_eq!(take_pending_token(), Some(9)); + assert_eq!(take_pending_token(), None); + assert_eq!(pending_token_count(), 0); + } + + /// A leftover token must never reach the next sequence. `clear` is the only + /// thing standing between a fall-back and one sequence's token appearing in + /// another's output. + #[test] + fn clearing_pending_tokens_prevents_leaking_them_into_the_next_sequence() { + clear_pending_tokens(); + push_pending_tokens(&[7, 8, 9]); + clear_pending_tokens(); + assert_eq!(take_pending_token(), None); + } + + /// A negative id in the ring means the ring was misread. Casting it would + /// index the tokenizer at ~4 billion; it must be dropped and logged. + #[test] + fn negative_ring_ids_are_dropped_not_wrapped_into_huge_u32() { + clear_pending_tokens(); + push_pending_tokens(&[5, -1, 6]); + assert_eq!(take_pending_token(), Some(5)); + assert_eq!(take_pending_token(), Some(6)); + assert_eq!(take_pending_token(), None); + } + + /// The eligibility flag must start `false`, so the first decode step of a + /// sequence always falls back rather than engaging on parameters nothing + /// has checked yet. + #[test] + fn eligibility_starts_false_and_stand_down_resets_everything() { + set_device_loop_eligible(false); + clear_pending_tokens(); + assert!(!device_loop_eligible()); + + set_device_loop_eligible(true); + push_pending_tokens(&[1, 2]); + assert!(device_loop_eligible()); + + stand_down(); + assert!(!device_loop_eligible(), "stand_down must clear eligibility"); + assert_eq!(take_pending_token(), None, "stand_down must clear tokens"); + } + + /// Default-off, and explicitly pinnable off. + #[test] + fn the_opt_in_defaults_to_off() { + // Not using the process env here: these run in parallel with other + // tests. The parsing rule is what matters and it is exercised directly. + for (val, want) in [ + ("", false), + ("0", false), + ("false", false), + ("off", false), + ("no", false), + ("1", true), + ("true", true), + ] { + let v = val.trim().to_ascii_lowercase(); + let enabled = !(v.is_empty() || v == "0" || v == "false" || v == "off" || v == "no"); + assert_eq!(enabled, want, "for {val:?}"); + } + // And with the variable absent entirely. + std::env::remove_var(ENV_ENABLE); + assert!(!device_loop_enabled()); + } +} diff --git a/arc-cuda-graph/src/ffi.rs b/arc-cuda-graph/src/ffi.rs index ae78c6f76d..96cae29542 100644 --- a/arc-cuda-graph/src/ffi.rs +++ b/arc-cuda-graph/src/ffi.rs @@ -179,6 +179,13 @@ extern "C" { pub fn cuStreamCreate(stream: *mut CUstream, flags: u32) -> u32; pub fn cuStreamDestroy_v2(stream: CUstream) -> u32; pub fn cudaStreamSynchronize(stream: CUstream) -> u32; + /// Non-blocking completion probe. Returns `cudaSuccess` when every + /// enqueued item has retired and `cudaErrorNotReady` (600) while work is + /// outstanding. NVIDIA's reference adds: *"Note that this function may + /// also return error codes from previous, asynchronous launches."* That + /// is what lets `device_loop` surface an async graph fault without ever + /// blocking the stream — see `device_loop::CudaDeviceOps::poll_stream`. + pub fn cudaStreamQuery(stream: CUstream) -> u32; // ========== Conditional nodes (CUDA 12.4+) ========== pub fn cudaGraphConditionalHandleCreate( diff --git a/arc-cuda-graph/src/graph.rs b/arc-cuda-graph/src/graph.rs index c288998f74..56c37803fc 100644 --- a/arc-cuda-graph/src/graph.rs +++ b/arc-cuda-graph/src/graph.rs @@ -728,6 +728,41 @@ output_trusted={} verify_remaining={} verify_failed={}", } } + /// The handles the device decode loop needs to replay this graph **without + /// the synchronize above**: the exec to launch and the graph-owned output + /// tensor to sample from, in place, on device. + /// + /// Handing these out is the whole difference between `replay()` and + /// `device_loop`. `replay()` launches, blocks on `cudaStreamSynchronize` + /// (line 362) and returns the output for the *host* to sample — which + /// copies logits down and serialises the step. `device_loop` takes the same + /// two handles, launches, samples on device and scatters the token back + /// into the pinned input buffer, so the next launch needs no host at all. + /// + /// The returned `Tensor`'s storage was allocated from the graph's private + /// pool during capture, so its address is stable for the life of the + /// `CapturedGraph` — which is exactly why sampling it in place is sound. + /// The clone is an `Arc` bump and does **not** extend that lifetime past + /// the runner's; callers must not hold it after dropping the runner. + pub fn replay_handles(&self, batch_size: usize) -> Option<(CUgraphExec, Tensor)> { + let captured = self.graphs.get(&batch_size)?; + let out = captured.output.clone()?; + Some((captured.exec, out)) + } + + /// The stream capture and replay run on. `device_loop` launches its commit + /// kernel here so it is ordered behind the graph without an event. + pub fn stream(&self) -> CUstream { + self.stream + } + + /// Count a decode step driven by the device loop rather than by + /// [`replay`](Self::replay), so `status_line` stays honest about how much + /// work this runner actually did. + pub fn record_device_loop_steps(&mut self, steps: u64) { + self.replays += steps; + } + pub fn disable(&mut self) { self.enabled = false; } diff --git a/arc-cuda-graph/src/lib.rs b/arc-cuda-graph/src/lib.rs index 8e534ed6db..766653698e 100644 --- a/arc-cuda-graph/src/lib.rs +++ b/arc-cuda-graph/src/lib.rs @@ -12,6 +12,11 @@ pub mod autonomous; pub mod buffers; pub mod decode_forward; pub mod dedicated; +/// Device-resident decode loop: N steps, N `cuGraphLaunch` calls, zero stream +/// synchronizations. Opt-in via `ARC_GRAPH_DEVICE_LOOP`. The driver and its +/// policy compile without CUDA so they are unit-tested on any host; only the +/// pointer plumbing is gated. +pub mod device_loop; pub mod ffi; pub mod gemv_ffi; pub mod graph; @@ -82,6 +87,14 @@ pub use buffers::{DecodeInputBuffers, DecodeState}; pub use decode_forward::{decode_forward, DecodeBuffers, LayerKvCache, PagedAttentionState}; #[cfg(feature = "cuda")] pub use dedicated::DedicatedDecodePath; +pub use device_loop::{ + admit, clear_pending_tokens, device_loop_eligible, device_loop_enabled, device_loop_killed, + kill_device_loop, pending_token_count, push_pending_tokens, set_device_loop_eligible, + stand_down, take_pending_token, AdmissionFacts, BurstDriver, BurstOutcome, DeviceLoopConfig, + DeviceLoopError, DeviceOps, Refusal, StreamState, +}; +#[cfg(feature = "cuda")] +pub use device_loop::{CudaDeviceOps, DeviceDecodeLoop}; #[cfg(feature = "cuda")] pub use graph::CudaGraphRunner; #[cfg(feature = "cuda")] diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index b93ae48dfa..585b867504 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -96,6 +96,13 @@ pub struct NormalPipeline { /// load time; graph capture is deferred until first decode call. #[cfg(feature = "cuda")] autonomous_runner: Option, + /// ArcInfer/ArcGraph device decode loop — replays the captured graph, + /// samples and scatters the token on device, and never synchronizes. Built + /// lazily on the first admitted decode step and then kept, because every + /// buffer it points at must stay at the address the graph baked. + /// Opt-in via `ARC_GRAPH_DEVICE_LOOP`; `None` until then. + #[cfg(feature = "cuda")] + device_decode_loop: Option, } /// What a forward pass should do with candle's caching allocator before it @@ -1986,6 +1993,10 @@ impl Loader for NormalLoader { // lazy init (and graph capture) once a real decode batch arrives. #[cfg(feature = "cuda")] autonomous_runner: None, + // Needs the captured graph's exec handle and output tensor, so it + // cannot exist before capture. Built on the first admitted step. + #[cfg(feature = "cuda")] + device_decode_loop: None, }))) } @@ -2108,6 +2119,171 @@ impl MetadataMixin for NormalPipeline { } } +#[cfg(feature = "cuda")] +/// Drive one burst of the ArcInfer/ArcGraph device decode loop. +/// +/// Returns `Some(logits)` when the step is covered — either because a burst +/// just ran or because a previous burst already produced this step's token +/// — and `None` to fall back. +/// +/// The tokens go to `arc_cuda_graph`'s pending queue and are drained one at +/// a time by `pipeline/sampling.rs`, so the engine's per-token `Sequence` +/// bookkeeping (stop strings, streaming, EOS) is unchanged. +/// +/// ⚠️ On the `pending > 0` path this returns the captured graph's output +/// tensor **without launching anything**, so those logits are the last +/// burst's, not this step's. That is sound only because the sampling hook +/// short-circuits before reading them and, at batch 1, +/// `pipeline/mod.rs:538` does no host copy. It is not a general-purpose +/// return value, which is why this whole path refuses anything that would +/// actually read the logits (`return_logprobs`, logits processors, +/// penalties, grammars). +fn device_decode_burst( + device: &Device, + runner: &mut arc_cuda_graph::CudaGraphRunner, + loop_slot: &mut Option, + bs: usize, + seq_len: usize, + current_position: u32, + input_ids_pinned: &Tensor, +) -> Option { + use arc_cuda_graph as acg; + + let (_exec, out) = runner.replay_handles(bs)?; + + // A burst already produced this step's token: no launch, no sync, no + // GPU work at all. This is where the amortisation actually shows up — + // `burst - 1` of every `burst` engine iterations do nothing but host + // bookkeeping. + if acg::pending_token_count() > 0 { + return Some(out); + } + + // The fixed-capacity KV window width. Read from the graph-mode mask's + // last dimension rather than guessed: `models/deepseek4.rs` builds that + // mask as `[B, 1, 1, cfg_full.sliding_window]`, and `cfg_full` is a + // model fact the pipeline cannot otherwise reach. + let capacity = match crate::layers::graph_mode_mask() { + Some(m) => *m.dims().last().unwrap_or(&0), + None => { + // The mask is what keeps the decode arm from reading unwritten + // KV slots; without it we cannot bound the burst. + acg::stand_down(); + return None; + } + }; + let vocab = *out.dims().last().unwrap_or(&0); + if capacity == 0 || vocab == 0 { + acg::stand_down(); + return None; + } + + let cfg = acg::DeviceLoopConfig { + position_limit: capacity as u32, + ..acg::DeviceLoopConfig::new(vocab, -1, capacity as u32) + }; + + let facts = acg::AdmissionFacts { + batch_size: bs, + seq_len, + graph_ready: runner.replay_output_trusted(), + // Established by `sampling.rs` from the Sequence on the previous + // step; false until one has confirmed it. + greedy: acg::device_loop_eligible(), + return_logprobs: false, + host_side_logits_work: false, + current_position, + }; + let steps = match acg::admit(&facts, &cfg, acg::device_loop_enabled()) { + Ok(s) => s, + Err(reason) => { + tracing::debug!("ArcGraph device loop declined this step: {reason}"); + acg::stand_down(); + return None; + } + }; + + if loop_slot.is_none() { + let positions = crate::layers::graph_mode_positions(); + match acg::DeviceDecodeLoop::new( + device, + _exec, + runner.stream(), + &out, + input_ids_pinned, + positions.as_ref(), + cfg, + ) { + Ok(l) => { + tracing::info!( + "ArcGraph device decode loop ENGAGED: burst={} vocab={vocab} \ + kv_window={capacity}. {}", + cfg.burst, + l.status_line() + ); + *loop_slot = Some(l); + } + Err(e) => { + acg::kill_device_loop(&format!("could not build the device loop: {e}")); + return None; + } + } + } + let dl = loop_slot.as_mut()?; + + match dl.run(steps) { + Ok(outcome) => { + // `first()` rather than `[0]`: a burst that returned no rows would + // be a bug, but it must not panic the engine thread mid-request. + let row = match outcome.tokens.first() { + Some(r) => r, + None => { + acg::kill_device_loop("burst returned no token rows"); + *loop_slot = None; + return None; + } + }; + acg::push_pending_tokens(row); + // ── ALIASING CONTRACT, ENFORCED ───────────────────────────── + // The tensor we are about to return aliases the graph-owned + // storage that the NEXT `cuGraphLaunch` overwrites. It is only + // safe to hand back because the sampling hook short-circuits + // before anything reads it — and that hook fires if and only if a + // token is pending. So the invariant is not "we reasoned it + // through", it is checked here: no pending token, no aliasing + // tensor. `push_pending_tokens` drops negative ids, so a row that + // was entirely garbage lands here rather than silently arriving at + // a `sample_sequence` that would then read clobbered logits and + // return a plausible wrong token. + if acg::pending_token_count() == 0 { + acg::kill_device_loop( + "burst produced no usable token, so the aliasing logits tensor must not be \ + returned", + ); + *loop_slot = None; + return None; + } + runner.record_device_loop_steps(outcome.steps_launched as u64); + if outcome.hit_eos { + tracing::debug!( + "ArcGraph device loop: EOS at token {} of a {}-step burst; the remaining \ + steps were already enqueued and are discarded", + row.len(), + outcome.steps_launched + ); + } + Some(out) + } + Err(e) => { + // A burst failure means the graph input buffer may hold an + // uncommitted or stale id, so this cannot be retried. + acg::kill_device_loop(&format!("burst failed: {e}")); + *loop_slot = None; + None + } + } +} + #[async_trait::async_trait] impl Pipeline for NormalPipeline { fn forward_inputs( @@ -2256,6 +2432,11 @@ impl Pipeline for NormalPipeline { } } } + // Bound before the chain: `self.device()` returns by value, + // and taking a reference to it inside an expression that + // also takes `&mut self.device_decode_loop` would borrow + // `self` both ways at once. + let dev_for_loop = self.device(); let captured: Option = if probe && seq_len == 1 && self.cuda_graph_runner.is_some() @@ -2555,6 +2736,32 @@ impl Pipeline for NormalPipeline { cd.set_capture_mode(false); } cap_result + } else if runner.has_graph(bs) + && arc_cuda_graph::device_loop_enabled() + && !arc_cuda_graph::device_loop_killed() + && device_decode_burst( + &dev_for_loop, + &mut runner, + &mut self.device_decode_loop, + bs, + seq_len, + seqlen_offsets.first().copied().unwrap_or(0) as u32, + &input_ids, + ) + .is_some() + { + // ArcInfer/ArcGraph device decode loop: the burst + // already ran (or a previous one covered this step), + // the token was sampled and scattered on device, and + // `pipeline/sampling.rs` will take it from the ring. + // No sync, no host sample, no D2H. + // + // The logits handed back are the captured graph's + // output tensor. They are NOT read: the sampling + // hook short-circuits first, and at batch 1 + // `pipeline/mod.rs:538` does no host copy. Admission + // refuses everything that would read them. + runner.replay_handles(bs).map(|(_, out)| out) } else if runner.has_graph(bs) { let t = std::time::Instant::now(); let replayed = runner.replay(bs); diff --git a/mistralrs-core/src/pipeline/sampling.rs b/mistralrs-core/src/pipeline/sampling.rs index b413e380dd..ba1fb6bb8d 100644 --- a/mistralrs-core/src/pipeline/sampling.rs +++ b/mistralrs-core/src/pipeline/sampling.rs @@ -527,6 +527,60 @@ pub async fn sample_and_add_toks( Ok(()) } +/// Is this sequence one the ArcGraph device decode loop may drive, and does it +/// already have a token waiting? +/// +/// Two jobs, deliberately in one place. It **publishes** eligibility, because +/// `Pipeline::forward_inputs` decides whether to run a burst but is handed only +/// `inputs` and `return_raw_logits` — it never sees a `Sequence`, so it cannot +/// ask whether sampling is greedy. And it **takes** the token a burst already +/// produced. +/// +/// Greedy only. The device sampler runs Splitmix64 per row against the host's +/// Isaac64, so under any stochastic configuration the two draw different — +/// individually valid — tokens, which would silently break seeded +/// reproducibility. Under greedy they agree exactly, which is also what makes +/// the path verifiable against the eager one. +/// +/// The returned `Logprobs` is not a lossy stand-in: `logprob: 0.0` is exactly +/// what the greedy host path returns (`sampler.rs:1490`), and `bytes` / +/// `top_logprobs` are read only when `return_logprobs` is set, which this +/// refuses. +#[cfg(feature = "cuda")] +fn device_loop_pre_sampled_token( + seq: &Sequence, + return_logprobs: bool, + sample_speculative: bool, +) -> Option { + let eligible = !return_logprobs + && !sample_speculative + && seq.sampler().is_greedy_trivial() + && matches!(seq.recognizer, SequenceRecognizer::None); + arc_cuda_graph::set_device_loop_eligible(eligible); + if !eligible { + // Stand down rather than merely decline: a token parked by an earlier, + // eligible sequence must never be handed to this one. + arc_cuda_graph::stand_down(); + return None; + } + arc_cuda_graph::take_pending_token().map(|token| Logprobs { + token, + logprob: 0.0, + bytes: None, + top_logprobs: None, + }) +} + +/// Without CUDA there is no device loop, so nothing is ever pre-sampled. +#[cfg(not(feature = "cuda"))] +fn device_loop_pre_sampled_token( + _seq: &Sequence, + _return_logprobs: bool, + _sample_speculative: bool, +) -> Option { + None +} + /// Async sample optionally adding to trie. #[allow(clippy::too_many_arguments)] pub async fn sample_sequence( @@ -538,6 +592,17 @@ pub async fn sample_sequence( sample_speculative: bool, multiple_sequences: bool, ) -> Result { + // ── ArcInfer/ArcGraph device decode loop ──────────────────────────────── + // If a burst already sampled this token ON DEVICE, take it and skip the + // host argmax entirely. That D2H (`sampler.rs:1479`) is the last + // synchronization in the decode step, so this is the half of the change + // that makes removing `replay()`'s `cudaStreamSynchronize` worth anything. + if let Some(pre_sampled) = + device_loop_pre_sampled_token(seq, return_logprobs, sample_speculative) + { + return Ok(pre_sampled); + } + // Both of these open and close inside a single poll, before the first // `.await`, so they nest correctly under `sample.join_all` even when B // futures are being polled in turn on this thread. diff --git a/mistralrs-core/src/sampler.rs b/mistralrs-core/src/sampler.rs index b282841c68..65547f010b 100644 --- a/mistralrs-core/src/sampler.rs +++ b/mistralrs-core/src/sampler.rs @@ -505,6 +505,28 @@ impl Sampler { || self.dry_params.is_some() } + /// Would [`Sampler::sample`] take the pure-argmax path — no temperature and + /// nothing that has to touch the logits on the host? + /// + /// This is the exact precondition for ArcGraph's device decode loop to + /// substitute its on-device argmax for the host one + /// (`pipeline/sampling.rs`). It deliberately mirrors the conditions inside + /// `sample` rather than approximating them: `logits_bias` and `top_nsigma` + /// are applied *before* the `trivial` gate, so a sampler carrying either + /// still needs the host even though the gate itself ignores them. + /// + /// Greedy is also the only mode where device and host agree *exactly* — the + /// device sampler runs Splitmix64 per row against the host's Isaac64, so + /// any stochastic mode would draw a different, equally valid token and + /// silently break seeded reproducibility. + pub fn is_greedy_trivial(&self) -> bool { + self.temperature.is_none() + && self.logits_processors.is_empty() + && !self.has_penalties() + && self.logits_bias.is_none() + && self.top_nsigma.is_none() + } + /// True when this sampler is exactly `argmax` over the model's **raw** /// logits — no temperature, no penalties, no custom logits processors. /// From f5afeecf93a5e3ac890ed6397b4bbd49258ff1cb Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:26:55 +0100 Subject: [PATCH 2/4] fix(ArcGraph): device loop stands aside for ARC_GRAPH_VERIFY_ALWAYS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- arc-cuda-graph/src/graph.rs | 8 ++++++++ mistralrs-core/src/pipeline/normal.rs | 10 ++++++++++ 2 files changed, 18 insertions(+) diff --git a/arc-cuda-graph/src/graph.rs b/arc-cuda-graph/src/graph.rs index 56c37803fc..e26f53cbca 100644 --- a/arc-cuda-graph/src/graph.rs +++ b/arc-cuda-graph/src/graph.rs @@ -744,6 +744,14 @@ output_trusted={} verify_remaining={} verify_failed={}", /// `CapturedGraph` — which is exactly why sampling it in place is sound. /// The clone is an `Arc` bump and does **not** extend that lifetime past /// the runner's; callers must not hold it after dropping the runner. + /// + /// 🔴 The same BOOBY TRAP `replay` documents applies here verbatim: this + /// output aliases storage the next launch overwrites, so it is not a + /// snapshot. `device_loop` is safe against it for a specific reason — + /// every access is enqueued on this same stream, in order, and none of + /// them is a host read — and `normal.rs` enforces at runtime the invariant + /// that lets the aliasing tensor be handed back at all. Anything else + /// taking these handles must copy. pub fn replay_handles(&self, batch_size: usize) -> Option<(CUgraphExec, Tensor)> { let captured = self.graphs.get(&batch_size)?; let out = captured.output.clone()?; diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index 585b867504..647b1d67bf 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -2149,6 +2149,16 @@ fn device_decode_burst( ) -> Option { use arc_cuda_graph as acg; + // Stand aside for the replay-verification lane (PR #205). This branch sits + // EARLIER in the chain than the `runner.replay(bs)` arm, so without this it + // would take every step and `ARC_GRAPH_VERIFY_ALWAYS=1` would silently + // verify nothing — the operator would get a clean diagnostic that never + // ran. A verification switch that is quietly bypassed is worse than one + // that is off. + if std::env::var("ARC_GRAPH_VERIFY_ALWAYS").as_deref() == Ok("1") { + return None; + } + let (_exec, out) = runner.replay_handles(bs)?; // A burst already produced this step's token: no launch, no sync, no From e385e5f849aea915084aeb80f32476c9cd9beb21 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 13:28:47 +0100 Subject: [PATCH 3/4] docs(ArcGraph): record the burst's scheduler-admission limitation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- arc-cuda-graph/src/device_loop.rs | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/arc-cuda-graph/src/device_loop.rs b/arc-cuda-graph/src/device_loop.rs index d4aae80edf..e2885182b8 100644 --- a/arc-cuda-graph/src/device_loop.rs +++ b/arc-cuda-graph/src/device_loop.rs @@ -96,6 +96,25 @@ //! `max_spins` is a last-resort liveness backstop for a kernel that hangs //! without ever erroring. It is not the primary detector. //! +//! # KNOWN LIMITATION — the 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. That is why `admit` refuses `batch_size != 1` +//! outright rather than treating it as a tuning question. +//! +//! Anyone widening this past batch 1 owns that problem first. Raising `burst` +//! trades admission latency for host round trips one-for-one, and the +//! trade is only free while there is exactly one sequence in flight. +//! //! # Layering //! //! Everything above the CUDA calls is expressed against the [`DeviceOps`] From 2ad07fca10393ac68812286349787fe302f36ec9 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 16:52:52 +0100 Subject: [PATCH 4/4] fix(ArcGraph): the device-loop sampling hook takes &mut Sequence MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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`, 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) Claude-Session: https://claude.ai/code/session_01UMmjFy8TvsgypxVWNVhhC7 --- mistralrs-core/src/pipeline/sampling.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/mistralrs-core/src/pipeline/sampling.rs b/mistralrs-core/src/pipeline/sampling.rs index ba1fb6bb8d..b87f9cf863 100644 --- a/mistralrs-core/src/pipeline/sampling.rs +++ b/mistralrs-core/src/pipeline/sampling.rs @@ -546,9 +546,11 @@ pub async fn sample_and_add_toks( /// what the greedy host path returns (`sampler.rs:1490`), and `bytes` / /// `top_logprobs` are read only when `return_logprobs` is set, which this /// refuses. +// `&mut` only because `Sequence::sampler()` takes `&mut self` (`sequence.rs:885`) +// to hand back a clone of an `Arc`. Nothing here mutates the sequence. #[cfg(feature = "cuda")] fn device_loop_pre_sampled_token( - seq: &Sequence, + seq: &mut Sequence, return_logprobs: bool, sample_speculative: bool, ) -> Option { @@ -574,7 +576,7 @@ fn device_loop_pre_sampled_token( /// Without CUDA there is no device loop, so nothing is ever pre-sampled. #[cfg(not(feature = "cuda"))] fn device_loop_pre_sampled_token( - _seq: &Sequence, + _seq: &mut Sequence, _return_logprobs: bool, _sample_speculative: bool, ) -> Option {