From be7f971db87dc17d429aeb2c461d7ce2b8e46f06 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:08:07 +0100 Subject: [PATCH 1/2] =?UTF-8?q?fix(ArcInfer/ArcGraph):=20parked=20burst=20?= =?UTF-8?q?tokens=20no=20longer=20outlive=20their=20sequence=20=E2=80=94?= =?UTF-8?q?=20one=20funnel,=20one=20leak=20test?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A device-loop burst runs with eos_token_id = -1, so it always parks its full burst; a sequence stopping mid-burst (stop string, max_tokens, cancel, error, panic, eviction) left its undrained tokens parked, and the NEXT sequence's first sample_sequence — including the one after its prefill — consumed them as its own. One user's tokens in another user's response. clear_pending_tokens documented the contract but its only production call sites were fallback paths. The fix is one call at the one funnel every completion path passes through: Sequence::set_state stands the device loop down on every transition out of the running set (Done, Error, FinishedAborted, FinishedIgnored, Waiting, Swapped). Clearing on preemption is deliberate fail-closed: the sequence re-prefills on re-admission, so its parked tokens are stale. To make the funnel testable on machines without CUDA — the cfg(cuda) darkness is exactly what let this ship — arc-cuda-graph is now a non-optional dependency of mistralrs-core (its host surface is plain Rust; CUDA plumbing stays behind its own cuda feature, still enabled by mistralrs-core/cuda) and device_loop_pre_sampled_token loses its cfg split. The new leak test drives the real sample_sequence path both ways: wired, B's first token is its own argmax and the queue is empty; with the funnel removed, the queue holds 3 of A's tokens and B's first sampled token is 22 — A's parked id. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --- arc-cuda-graph/src/device_loop.rs | 13 +- mistralrs-core/Cargo.toml | 10 +- mistralrs-core/src/pipeline/sampling.rs | 237 ++++++++++++++++++++++-- mistralrs-core/src/sequence.rs | 35 ++++ 4 files changed, 281 insertions(+), 14 deletions(-) diff --git a/arc-cuda-graph/src/device_loop.rs b/arc-cuda-graph/src/device_loop.rs index e2885182b..e89dcc2a2 100644 --- a/arc-cuda-graph/src/device_loop.rs +++ b/arc-cuda-graph/src/device_loop.rs @@ -1224,13 +1224,24 @@ pub fn device_loop_eligible() -> bool { /// 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. +/// +/// The sequence-completion half of that contract is enforced at ONE funnel: +/// `mistralrs-core`'s `Sequence::set_state` calls [`stand_down`] on every +/// transition out of the running set (`Done`, `Error`, `FinishedAborted`, +/// `FinishedIgnored`, `Waiting`, `Swapped`), which is why this crate's +/// dependency in `mistralrs-core` is deliberately not optional — the funnel +/// and its leak test compile and run on hosts without CUDA. The tests here +/// cover only this function's own behaviour; the call-site wiring is covered +/// by `pipeline/sampling.rs`'s `device_loop_cross_sequence_leak_tests`. 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. +/// the next sequence starts from a clean slate. Sequence completion — every +/// transition out of the running set — calls this from the +/// `Sequence::set_state` funnel; see [`clear_pending_tokens`]. pub fn stand_down() { clear_pending_tokens(); set_device_loop_eligible(false); diff --git a/mistralrs-core/Cargo.toml b/mistralrs-core/Cargo.toml index a58d0a9ff..dd7728002 100644 --- a/mistralrs-core/Cargo.toml +++ b/mistralrs-core/Cargo.toml @@ -71,7 +71,14 @@ mistralrs-paged-attn = { workspace = true, optional = true } mistralrs-quant.workspace = true arc-turbo = { path = "../arc-turbo" } arc-profiler = { path = "../arc-profiler" } -arc-cuda-graph = { path = "../arc-cuda-graph", optional = true, features = [] } +# NOT optional, deliberately: the sequence-completion funnel +# (`Sequence::set_state`) must stand the device decode loop down on every +# finished transition, and a gate that only compiles under +# `cfg(feature = "cuda")` cannot be tested on any machine that reviews it — +# the exact class that let the cross-request token leak ship. The crate's +# host-side surface (pending-token queue, burst driver, admission) is plain +# Rust; the CUDA plumbing inside it stays behind its own `cuda` feature. +arc-cuda-graph = { path = "../arc-cuda-graph", features = [] } uuid = { workspace = true, features = ["v4"] } utoipa = { workspace = true, optional = true } schemars.workspace = true @@ -124,7 +131,6 @@ cuda = [ "mistralrs-quant/cuda", "dep:mistralrs-paged-attn", "mistralrs-paged-attn/cuda", - "dep:arc-cuda-graph", "arc-cuda-graph/cuda", "arc-profiler/cuda", ] diff --git a/mistralrs-core/src/pipeline/sampling.rs b/mistralrs-core/src/pipeline/sampling.rs index b87f9cf86..b7cc8045f 100644 --- a/mistralrs-core/src/pipeline/sampling.rs +++ b/mistralrs-core/src/pipeline/sampling.rs @@ -548,7 +548,12 @@ pub async fn sample_and_add_toks( /// 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")] +// +// Deliberately NOT `cfg(feature = "cuda")`: only the CUDA plumbing can *park* +// tokens, but the take/stand-down half is plain host code, and gating it made +// the cross-request leak untestable on every machine that reviewed it. On a +// non-CUDA build nothing ever parks, `take_pending_token` is always `None`, +// and this is one thread-local read per sample. fn device_loop_pre_sampled_token( seq: &mut Sequence, return_logprobs: bool, @@ -573,16 +578,6 @@ 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: &mut 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( @@ -993,3 +988,223 @@ mod speculative_verification_tests { ); } } + +/// ArcInfer/ArcGraph: a device-loop burst that outlives its sequence must +/// never reach the next sequence. +/// +/// The burst runs with `eos_token_id = -1` (`device_decode_burst` passes -1, +/// so no on-device EOS truncation ever fires) and parks its full length; the +/// engine drains one token per step. A sequence that stops mid-burst — stop +/// string, `max_tokens`, cancel, error — leaves the tail parked, and +/// `device_loop_pre_sampled_token` hands parked tokens to whichever eligible +/// sequence samples next, **including the next request's first sample after +/// its prefill**. That is one user's tokens inside another user's response. +/// +/// The fix under test is the funnel: `Sequence::set_state` stands the device +/// loop down on every transition out of the running set. These tests drive the +/// REAL `sample_sequence` path, so they fail loudly if that call is removed — +/// unlike the arc-cuda-graph unit tests, which only ever exercised +/// `clear_pending_tokens` itself, not that anything calls it. +#[cfg(test)] +mod device_loop_cross_sequence_leak_tests { + use super::*; + use crate::sampler::Sampler; + use crate::sequence::{SeqStepType, SequenceGroup, SequenceState, StopReason}; + use candle_core::Device; + use rand::SeedableRng; + use std::cell::RefCell; + + const VOCAB: usize = 32; + /// Argmax of every logits tensor handed to sequence A. + const PEAK_A: u32 = 7; + /// Argmax of every logits tensor handed to sequence B. Distinct from + /// `PEAK_A` and from the burst, so the origin of B's token is unambiguous. + const PEAK_B: u32 = 9; + /// The burst sequence A's device loop parked. A drains the first token, + /// then stops mid-burst, leaving `[22, 23, 24]`. + const BURST: [i32; 4] = [21, 22, 23, 24]; + + thread_local! { + /// Keeps each fixture sequence's `Receiver` alive; a dropped receiver + /// models a disconnected client and would silently change behaviour. + static LIVE_CLIENTS: RefCell>> = + const { RefCell::new(Vec::new()) }; + } + + /// A greedy-trivial sequence: no temperature, penalties, bias, processors + /// or recognizer — exactly the shape `device_loop_pre_sampled_token` + /// accepts, so the pending queue is live for it. + fn greedy_seq() -> Sequence { + let (dummy_sender, rx) = tokio::sync::mpsc::channel(1); + LIVE_CLIENTS.with(|k| k.borrow_mut().push(rx)); + let group = Arc::new(std::sync::Mutex::new(SequenceGroup::new( + 1, false, false, None, + ))); + let sampler = Sampler::new( + None, + 0, + None, + None, + None, + None, + None, + -1, + 1.0, + 0.0, + None, + vec![], + None, + ) + .unwrap(); + Sequence::new_waiting( + vec![1u32; 4], + String::new(), + 0, + 0, + 1, + dummy_sender, + sampler, + vec![], + vec![], + None, + false, + false, + group, + 0, + 0, + SequenceRecognizer::None, + None, + None, + None, + None, + None, + None, + None, + SeqStepType::PromptAndDecode, + None, + None, + None, + false, + vec![], + ) + } + + /// `(1, 1, VOCAB)` logits peaking hard at `peak` — the shape + /// `sample_sequence` receives for one decode step of one sequence. + fn peaked_logits(peak: u32) -> Tensor { + let mut raw = vec![0f32; VOCAB]; + raw[peak as usize] = 20.0; + Tensor::from_vec(raw, (1, 1, VOCAB), &Device::Cpu).unwrap() + } + + async fn sample(seq: &mut Sequence, peak: u32, rng: &Arc>) -> u32 { + sample_sequence( + peaked_logits(peak), + seq, + false, + rng.clone(), + false, + false, + false, + ) + .await + .unwrap() + .token + } + + /// THE leak test: A parks a burst, stops mid-burst on a non-fallback path + /// (`max_tokens` → `StopReason::Length`, the transition + /// `finish_or_add_toks_to_seq` performs at `sampling.rs:267`/`:277`), and + /// B's first sample must consume nothing of A's. + /// + /// Mutation check (run both ways before shipping): with the + /// `stand_down` arm removed from `Sequence::set_state`, this fails at the + /// queue assertion with 3 parked tokens, and B's first token comes back + /// `22` — A's first undrained parked id — instead of `PEAK_B`. + #[tokio::test] + async fn a_sequence_stopping_mid_burst_leaks_nothing_into_the_next_sequence() { + // This thread's queue, not shared with other tests (thread-local). + arc_cuda_graph::stand_down(); + let rng = Arc::new(std::sync::Mutex::new(Isaac64Rng::seed_from_u64(0))); + + // ── Sequence A, decode step N: queue empty, host argmax, and the + // pre-hook publishes A's eligibility — the same order the engine + // runs (sample N publishes, forward N+1 bursts). + let mut seq_a = greedy_seq(); + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); + + // ── Forward N+1: the burst parks its FULL length (eos -1). + arc_cuda_graph::push_pending_tokens(&BURST); + + // ── Step N+1's sample drains one parked token. This is the control + // that proves the pending-token path is LIVE in this build: if the + // pre-hook were compiled out or inert, this would return `PEAK_A` + // and the leak below could never manifest — a test that cannot fail. + assert_eq!( + sample(&mut seq_a, PEAK_A, &rng).await, + BURST[0] as u32, + "the device-loop take path is not live in this build; the leak assertions below \ + would be vacuous" + ); + + // ── A stops MID-BURST via a non-fallback path: max_tokens. + seq_a.set_state(SequenceState::Done(StopReason::Length(2))); + + // ── The claim, part 1: nothing of A's survives its completion. + assert_eq!( + arc_cuda_graph::pending_token_count(), + 0, + "sequence A finished mid-burst but its parked tokens survived; the next \ + sequence's first sample will consume them" + ); + + // ── Sequence B's FIRST sample (the one right after its prefill). + let mut seq_b = greedy_seq(); + let b_first = sample(&mut seq_b, PEAK_B, &rng).await; + + // ── The claim, part 2: B's token comes from B's own logits and is + // none of A's parked ids. + assert!( + !BURST.contains(&(b_first as i32)), + "cross-request token leak: sequence B's first sampled token is {b_first}, one of \ + sequence A's parked burst tokens" + ); + assert_eq!( + b_first, PEAK_B, + "sequence B's first token must be the argmax of B's own logits" + ); + } + + /// Every transition out of the running set clears the queue — not just + /// `Done(Length)`. This pins the whole funnel map: client cancel, + /// forward-error/panic recovery (`Error`), PagedAttention eviction + /// (`FinishedAborted`/`FinishedIgnored`), and preemption + /// (`Waiting`/`Swapped`), where consuming the stale burst after + /// re-prefill would duplicate output. + #[tokio::test] + async fn every_transition_out_of_the_running_set_clears_parked_tokens() { + for state in [ + SequenceState::Done(StopReason::Canceled), + SequenceState::Done(StopReason::StopString { + stop_string_idx: 0, + completion_bytes_pos: 0, + }), + SequenceState::Error, + SequenceState::FinishedAborted, + SequenceState::FinishedIgnored, + SequenceState::Waiting, + SequenceState::Swapped, + ] { + arc_cuda_graph::stand_down(); + arc_cuda_graph::push_pending_tokens(&BURST); + assert_eq!(arc_cuda_graph::pending_token_count(), BURST.len()); + let seq = greedy_seq(); + seq.set_state(state); + assert_eq!( + arc_cuda_graph::pending_token_count(), + 0, + "transition to {state:?} left parked tokens for the next sequence" + ); + } + } +} diff --git a/mistralrs-core/src/sequence.rs b/mistralrs-core/src/sequence.rs index 46c8cdd41..bdb5b1638 100644 --- a/mistralrs-core/src/sequence.rs +++ b/mistralrs-core/src/sequence.rs @@ -1085,6 +1085,41 @@ impl Sequence { let mut group = get_mut_group!(self); group.n_choices = group.n_choices.saturating_sub(1); } + // ArcInfer/ArcGraph: this is the one funnel every path that takes a + // sequence out of the running set passes through — EOS / stop-token / + // stop-string / length / tool-calls (`pipeline/sampling.rs` + // `finish_or_add_toks_to_seq`), the schedulers' disconnected-client + // sweeps and `TERMINATE_ALL_NEXT_STEP` (`Done(Canceled)`), the + // engine's forward-error and panic recovery (`SequenceState::Error` + // via `handle_pipeline_forward_error!` / `step_catching_panics`), and + // PagedAttention eviction (`FinishedAborted` / `FinishedIgnored` / + // `Waiting` / `Swapped`). + // + // A device decode loop burst runs with `eos_token_id = -1`, so it + // always parks its full burst; a sequence that stops mid-burst leaves + // tokens parked, and a parked token is handed to the NEXT sequence's + // first `sample_sequence` — including the one after its prefill — as + // if it were its own. That is one user's tokens in another user's + // response, which is why the stand-down happens HERE, on the + // transition itself, rather than on N scattered completion paths. + // + // The pending queue is thread-local to the engine thread that parked + // it (the same thread that samples at batch 1); on any other thread + // this clears an empty queue, which is a no-op. Clearing on + // preemption (`Waiting` / `Swapped`) is deliberate fail-closed: the + // sequence re-prefills on re-admission, so its parked tokens are + // stale and consuming them would duplicate output. + match state { + SequenceState::Done(_) + | SequenceState::Error + | SequenceState::Waiting + | SequenceState::FinishedAborted + | SequenceState::FinishedIgnored + | SequenceState::Swapped => arc_cuda_graph::stand_down(), + SequenceState::RunningPrompt + | SequenceState::RunningCompletion + | SequenceState::RunningPrefillPrompt => {} + } *self.state.write().unwrap() = state; } From a30bec297009c52afd5e836f69752488e5a3155f Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:32:39 +0100 Subject: [PATCH 2/2] =?UTF-8?q?fix(ArcInfer/ArcGraph):=20owner-tag=20the?= =?UTF-8?q?=20pending=20queue=20=E2=80=94=20the=20leak=20becomes=20structu?= =?UTF-8?q?rally=20impossible,=20not=20funnel-dependent?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The completion funnel cannot see a sequence the DefaultScheduler moves aside WITHOUT a state change (the bucketing waitlist says so in its own doc), and batch growth has the same shape. So the queue itself now refuses: tokens are parked tagged with the owning sequence's id (published by the sampling pre-hook with eligibility), take requires a matching id, and a foreign queue is dropped whole with the grep-able instrument line "ArcGraph: dropped N foreign parked tokens". The forward's pending>0 short-circuit gains the two guards its aliasing contract always needed: seq_len == 1 (a parked queue at a prefill step must not replace the prefill) and pending_owned_by_current(). If the short-circuit fired and the sample still cannot take (scheduler swapped sequences in between), the step fails loudly via the aliased-logits marker instead of host-sampling another sequence's stale graph output. DeviceDecodeLoop finally resets: every stand_down (including the funnel's) bumps a device-state generation; run() compares and zeroes ring, cursors and fault word before a newer-generation burst. Reset at engagement, not completion, because the funnel cannot reach the pipeline-owned loop and between bursts is the only point with no in-flight ring writes. ARC_GRAPH_DEVICE_LOOP is now registered in capability_reachability.rs (Symbol: device_loop_enabled), closing the refuted premise from the previous commit. Interleave test (A waitlisted, no set_state): wired, B samples its own argmax 9 and the foreign tail is dropped; with the owner comparison removed, B's first sampled token is 22 — A's parked id. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --- arc-cuda-graph/src/device_loop.rs | 281 ++++++++++++++++-- arc-cuda-graph/src/lib.rs | 9 +- mistralrs-core/src/pipeline/normal.rs | 44 ++- mistralrs-core/src/pipeline/sampling.rs | 181 ++++++++++- .../tests/capability_reachability.rs | 14 + 5 files changed, 464 insertions(+), 65 deletions(-) diff --git a/arc-cuda-graph/src/device_loop.rs b/arc-cuda-graph/src/device_loop.rs index e89dcc2a2..757c31a7f 100644 --- a/arc-cuda-graph/src/device_loop.rs +++ b/arc-cuda-graph/src/device_loop.rs @@ -933,6 +933,15 @@ mod cuda_impl { pub struct DeviceDecodeLoop { ops: CudaDeviceOps, driver: BurstDriver, + /// The [`super::device_loop_generation`] this loop's ring, cursors and + /// fault word are valid for. Every `stand_down` — including the one + /// the `Sequence::set_state` completion funnel issues — bumps the + /// generation, and [`Self::run`] resets the per-sequence device state + /// before running a burst from a newer generation. Reset happens HERE, + /// at engagement, not at completion: the funnel cannot reach this + /// object (it lives in the pipeline), and between bursts is the only + /// point where no device write to the ring can still be in flight. + generation: u64, } // Same reasoning as `CudaGraphRunner` (`graph.rs:114`): the raw pointers are @@ -976,13 +985,26 @@ mod cuda_impl { cfg, )?; let driver = BurstDriver::new(cfg, batch)?; - Ok(Self { ops, driver }) + Ok(Self { + ops, + driver, + generation: super::device_loop_generation(), + }) } /// Enqueue `steps` decode steps and return their tokens. /// /// Exactly `steps` `cuGraphLaunch` calls, zero `cudaStreamSynchronize`. pub fn run(&mut self, steps: usize) -> candle_core::Result { + // A stand-down happened since the last burst — sequence completed, + // errored, or was preempted — so the ring head, read cursors and + // fault word describe a sequence that no longer exists. Zero them + // before this burst rather than carrying them across sequences. + let current = super::device_loop_generation(); + if self.generation != current { + self.reset(); + self.generation = current; + } 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 @@ -1156,14 +1178,37 @@ pub fn admit( std::thread_local! { static PENDING_TOKENS: std::cell::RefCell> = const { std::cell::RefCell::new(std::collections::VecDeque::new()) }; + /// The sequence the parked tokens belong to, stamped at park time from + /// [`CURRENT_SEQ`]. The completion funnel (`Sequence::set_state`) is the + /// primary defence against a parked token outliving its sequence; this tag + /// is the structural one — it covers the paths NO funnel can see, such as + /// the DefaultScheduler bucketing waitlist, which moves a running sequence + /// aside *without a state modification* (`default_scheduler.rs`, the + /// `bucket_and_waitlist_seqs_waiting` doc says so explicitly). + static PENDING_OWNER: std::cell::Cell> = const { std::cell::Cell::new(None) }; + /// The sequence id the last `sample_sequence` published alongside + /// eligibility. At batch 1 — the only shape the loop admits — this is + /// exactly the sequence the next forward's burst decodes for. + static CURRENT_SEQ: std::cell::Cell> = const { std::cell::Cell::new(None) }; + /// Set by the forward when it served this step from the pending queue and + /// handed back the ALIASING graph logits tensor without launching anything + /// (`normal.rs::device_decode_burst`, the `pending > 0` arm). Consumed by + /// the very next sample. If that sample cannot take a parked token, the + /// logits it holds are a stale alias of graph-owned storage and sampling + /// them would return a plausible token from the WRONG sequence's + /// distribution — so the sample must fail loudly instead. + static ALIASED_LOGITS_SERVED: std::cell::Cell = const { std::cell::Cell::new(false) }; } -/// Park a burst's tokens for the engine to drain. +/// Park a burst's tokens for the engine to drain. The queue is stamped with +/// the current decoding sequence ([`set_device_loop_eligible`]) as its owner; +/// [`take_pending_token`] refuses to hand them to anyone else. /// /// 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_OWNER.with(|o| o.set(CURRENT_SEQ.with(|c| c.get()))); PENDING_TOKENS.with(|q| { let mut q = q.borrow_mut(); for &t in tokens { @@ -1180,9 +1225,53 @@ pub fn push_pending_tokens(tokens: &[i32]) { }) } -/// 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()) +/// What [`take_pending_token`] found. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PendingTake { + /// A token parked by a burst for THIS sequence. + Taken(u32), + /// Nothing parked. + Empty, + /// The queue held another sequence's tokens. They were dropped — handing + /// even one across would put one user's tokens in another user's response + /// — and `dropped` says how many. Logged as + /// `"ArcGraph: dropped N foreign parked tokens"`; on a box that line is + /// the violation instrument, so keep it grep-able. + Foreign { dropped: usize }, +} + +/// Take the next parked token, if it belongs to `seq_id`. +/// +/// A queue owned by a different sequence is dropped whole, fail-closed: the +/// owner has left the running set (or been waitlisted with no state change), +/// its device-side burst state is stale, and replaying its tokens into any +/// sequence — including the owner after a re-prefill — would corrupt output. +pub fn take_pending_token(seq_id: usize) -> PendingTake { + let owner = PENDING_OWNER.with(|o| o.get()); + PENDING_TOKENS.with(|q| { + let mut q = q.borrow_mut(); + if q.is_empty() { + return PendingTake::Empty; + } + if owner == Some(seq_id) { + match q.pop_front() { + Some(t) => PendingTake::Taken(t), + None => PendingTake::Empty, + } + } else { + let dropped = q.len(); + q.clear(); + PENDING_OWNER.with(|o| o.set(None)); + tracing::warn!( + "ArcGraph: dropped {dropped} foreign parked tokens (parked for sequence \ + {owner:?}, sampler is sequence {seq_id}) — cross-sequence leak prevented. \ + Expected once when the scheduler moves a sequence aside mid-burst; frequent \ + occurrences mean the device loop is engaging under scheduler churn it cannot \ + amortise" + ); + PendingTake::Foreign { dropped } + } + }) } /// Tokens still parked. @@ -1190,6 +1279,29 @@ pub fn pending_token_count() -> usize { PENDING_TOKENS.with(|q| q.borrow().len()) } +/// Does the pending queue belong to the sequence whose sample last published +/// eligibility? The forward's `pending > 0` short-circuit must not fire for +/// anyone else: it returns stale aliasing logits and launches nothing, which +/// is only sound when the very next sample will take a parked token. +pub fn pending_owned_by_current() -> bool { + let owner = PENDING_OWNER.with(|o| o.get()); + owner.is_some() && owner == CURRENT_SEQ.with(|c| c.get()) +} + +/// The forward served this step from the pending queue: no launch, and the +/// returned logits tensor aliases graph-owned storage. The next sample MUST +/// take a parked token; [`take_aliased_logits_marker`] is how it checks. +pub fn note_aliased_logits_served() { + ALIASED_LOGITS_SERVED.with(|c| c.set(true)); +} + +/// Consume the aliased-logits marker for this step. `true` means the logits +/// the current sample holds were NOT produced by a launch for this step and +/// must not be host-sampled. +pub fn take_aliased_logits_marker() -> bool { + ALIASED_LOGITS_SERVED.with(|c| c.replace(false)) +} + // Whether the sequence currently being decoded is one the device loop may // drive. // @@ -1208,8 +1320,12 @@ 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) { +/// Publish whether the sequence being sampled may be driven by the device +/// loop, and WHICH sequence that is. The id becomes the owner stamped onto +/// anything the next burst parks, which is what makes a parked token +/// unconsumable by any other sequence. +pub fn set_device_loop_eligible(seq_id: usize, eligible: bool) { + CURRENT_SEQ.with(|c| c.set(Some(seq_id))); DEVICE_LOOP_ELIGIBLE.with(|c| c.set(eligible)) } @@ -1234,17 +1350,43 @@ pub fn device_loop_eligible() -> bool { /// cover only this function's own behaviour; the call-site wiring is covered /// by `pipeline/sampling.rs`'s `device_loop_cross_sequence_leak_tests`. pub fn clear_pending_tokens() { + PENDING_OWNER.with(|o| o.set(None)); 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. Sequence completion — every -/// transition out of the running set — calls this from the -/// `Sequence::set_state` funnel; see [`clear_pending_tokens`]. +/// Stand down completely: forget parked tokens (and their owner), drop the +/// aliased-logits marker, mark the current sequence ineligible, and bump the +/// device-state generation so a kept-alive [`DeviceDecodeLoop`] resets its +/// ring and cursors before it next runs. Call on fall-back, on sequence +/// completion, and on any error, so the next sequence starts from a clean +/// slate. Sequence completion — every transition out of the running set — +/// calls this from the `Sequence::set_state` funnel; see +/// [`clear_pending_tokens`]. pub fn stand_down() { clear_pending_tokens(); - set_device_loop_eligible(false); + ALIASED_LOGITS_SERVED.with(|c| c.set(false)); + DEVICE_LOOP_ELIGIBLE.with(|c| c.set(false)); + DEVICE_LOOP_GENERATION.with(|g| g.set(g.get().wrapping_add(1))); +} + +// The device-state generation. `DeviceDecodeLoop` holds ring cursors, the +// pinned ring head and the device-side position — all per-sequence state — +// but it is deliberately kept alive across sequences because every buffer it +// points at must stay at the address the captured graph baked +// (`DeviceDecodeLoop` docs). So instead of the funnel destroying it (it +// cannot even reach it: the loop lives in the pipeline, the funnel in +// `Sequence`), every `stand_down` bumps this generation and the loop compares +// on its next `run`, resetting itself first if the world moved on. Reset at +// engagement time is also the only SAFE time: between bursts nothing is in +// flight on the stream, whereas a reset at completion time could race a +// still-draining burst's device writes. +std::thread_local! { + static DEVICE_LOOP_GENERATION: std::cell::Cell = const { std::cell::Cell::new(0) }; +} + +/// The current device-state generation. Bumped by every [`stand_down`]. +pub fn device_loop_generation() -> u64 { + DEVICE_LOOP_GENERATION.with(|g| g.get()) } /// Process-wide latch. Once the device loop has failed, it stays off. @@ -1728,18 +1870,20 @@ mod tests { 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. + /// The parked tokens are handed out in generation order, exactly once, to + /// the sequence that owns them. 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(); + stand_down(); + set_device_loop_eligible(11, true); 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!(take_pending_token(11), PendingTake::Taken(7)); + assert_eq!(take_pending_token(11), PendingTake::Taken(8)); + assert_eq!(take_pending_token(11), PendingTake::Taken(9)); + assert_eq!(take_pending_token(11), PendingTake::Empty); assert_eq!(pending_token_count(), 0); } @@ -1748,21 +1892,93 @@ mod tests { /// another's output. #[test] fn clearing_pending_tokens_prevents_leaking_them_into_the_next_sequence() { - clear_pending_tokens(); + stand_down(); + set_device_loop_eligible(11, true); push_pending_tokens(&[7, 8, 9]); clear_pending_tokens(); - assert_eq!(take_pending_token(), None); + assert_eq!(take_pending_token(11), PendingTake::Empty); + } + + /// THE structural guarantee: tokens parked while sequence 11 was current + /// cannot be taken by sequence 12 — they are dropped whole, and the drop + /// is reported. This is the defence the completion funnel cannot provide: + /// the DefaultScheduler bucketing waitlist moves a running sequence aside + /// *without a state modification*, so no funnel ever fires. + #[test] + fn foreign_parked_tokens_are_dropped_not_served() { + stand_down(); + set_device_loop_eligible(11, true); + push_pending_tokens(&[7, 8, 9]); + + // Sequence 12 samples next (11 was waitlisted, no state change). + set_device_loop_eligible(12, true); + assert_eq!( + take_pending_token(12), + PendingTake::Foreign { dropped: 3 }, + "another sequence's parked tokens must be dropped, not served" + ); + assert_eq!(pending_token_count(), 0, "the foreign queue is gone whole"); + // And they are gone for the original owner too: its device-side burst + // state is stale, so replaying them after a re-prefill would corrupt. + assert_eq!(take_pending_token(11), PendingTake::Empty); + } + + /// The forward-side guard: the `pending > 0` short-circuit may only fire + /// for the sequence that owns the queue. + #[test] + fn pending_ownership_tracks_the_published_sequence() { + stand_down(); + assert!( + !pending_owned_by_current(), + "empty queue is owned by nobody" + ); + set_device_loop_eligible(11, true); + push_pending_tokens(&[7]); + assert!(pending_owned_by_current()); + set_device_loop_eligible(12, true); + assert!( + !pending_owned_by_current(), + "the moment another sequence publishes, the queue is foreign" + ); + } + + /// The aliased-logits marker is consumed exactly once, and stand_down + /// drops it: a marker surviving a stand-down would poison the next + /// sequence's first REAL sample. + #[test] + fn aliased_logits_marker_is_one_shot_and_cleared_by_stand_down() { + stand_down(); + assert!(!take_aliased_logits_marker()); + note_aliased_logits_served(); + assert!(take_aliased_logits_marker()); + assert!(!take_aliased_logits_marker(), "consumed exactly once"); + note_aliased_logits_served(); + stand_down(); + assert!(!take_aliased_logits_marker(), "stand_down must drop it"); + } + + /// Every stand_down bumps the device-state generation; that is what makes + /// a kept-alive `DeviceDecodeLoop` reset its ring and cursors before the + /// next sequence's first burst instead of carrying them across. + #[test] + fn stand_down_bumps_the_device_state_generation() { + let g0 = device_loop_generation(); + stand_down(); + assert_eq!(device_loop_generation(), g0.wrapping_add(1)); + stand_down(); + assert_eq!(device_loop_generation(), g0.wrapping_add(2)); } /// 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(); + stand_down(); + set_device_loop_eligible(11, true); 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); + assert_eq!(take_pending_token(11), PendingTake::Taken(5)); + assert_eq!(take_pending_token(11), PendingTake::Taken(6)); + assert_eq!(take_pending_token(11), PendingTake::Empty); } /// The eligibility flag must start `false`, so the first decode step of a @@ -1770,17 +1986,20 @@ mod tests { /// has checked yet. #[test] fn eligibility_starts_false_and_stand_down_resets_everything() { - set_device_loop_eligible(false); - clear_pending_tokens(); + stand_down(); assert!(!device_loop_eligible()); - set_device_loop_eligible(true); + set_device_loop_eligible(11, 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"); + assert_eq!( + take_pending_token(11), + PendingTake::Empty, + "stand_down must clear tokens" + ); } /// Default-off, and explicitly pinnable off. diff --git a/arc-cuda-graph/src/lib.rs b/arc-cuda-graph/src/lib.rs index 766653698..1760ad210 100644 --- a/arc-cuda-graph/src/lib.rs +++ b/arc-cuda-graph/src/lib.rs @@ -88,10 +88,11 @@ pub use decode_forward::{decode_forward, DecodeBuffers, LayerKvCache, PagedAtten #[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, + admit, clear_pending_tokens, device_loop_eligible, device_loop_enabled, device_loop_generation, + device_loop_killed, kill_device_loop, note_aliased_logits_served, pending_owned_by_current, + pending_token_count, push_pending_tokens, set_device_loop_eligible, stand_down, + take_aliased_logits_marker, take_pending_token, AdmissionFacts, BurstDriver, BurstOutcome, + DeviceLoopConfig, DeviceLoopError, DeviceOps, PendingTake, Refusal, StreamState, }; #[cfg(feature = "cuda")] pub use device_loop::{CudaDeviceOps, DeviceDecodeLoop}; diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index 0048e34c9..eaaaddcb9 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -167,11 +167,7 @@ pub(crate) enum AllocCacheAction { /// Pure, so the policy is testable without a GPU — which matters, because the /// allocator itself has no tests at all in either repo. #[cfg_attr(not(feature = "cuda"), allow(dead_code))] -pub(crate) fn alloc_cache_action( - seq_len: usize, - enabled: bool, - killed: bool, -) -> AllocCacheAction { +pub(crate) fn alloc_cache_action(seq_len: usize, enabled: bool, killed: bool) -> AllocCacheAction { let want = !killed && seq_len == 1; match (want, enabled) { (true, false) => AllocCacheAction::Enable, @@ -2165,8 +2161,32 @@ fn device_decode_burst( // GPU work at all. This is where the amortisation actually shows up — // `burst - 1` of every `burst` engine iterations do nothing but host // bookkeeping. + // + // Guarded twice, because returning `out` here hands back an ALIASING + // tensor from an earlier launch and skips this step's forward entirely — + // sound only when the very next sample takes a parked token: + // * `seq_len == 1` — parked tokens present at a PREFILL step belong to a + // sequence that left the running set without a state change (the + // bucketing waitlist); serving the new request's prefill from them + // would skip the prefill and leak the old sequence's tokens. + // * `pending_owned_by_current()` — the queue must belong to the sequence + // whose sample published eligibility last, i.e. the one this decode + // step is for at batch 1. + // Either violation drops the parked tokens (loudly) and falls through to + // the real forward. if acg::pending_token_count() > 0 { - return Some(out); + if seq_len == 1 && acg::pending_owned_by_current() { + acg::note_aliased_logits_served(); + return Some(out); + } + tracing::warn!( + "ArcGraph: dropped {} foreign parked tokens — parked queue found at a step it does \ + not belong to (seq_len={}, owned_by_current={}); running the real forward instead", + acg::pending_token_count(), + seq_len, + acg::pending_owned_by_current(), + ); + acg::stand_down(); } // The fixed-capacity KV window width. Read from the graph-mode mask's @@ -2388,9 +2408,7 @@ impl Pipeline for NormalPipeline { cd.set_alloc_cache_enabled(true); arm_alloc_cache_capacity(cd, &alloc_cache_device); } - AllocCacheAction::DrainAndDisable => { - cd.set_alloc_cache_enabled(false) - } + AllocCacheAction::DrainAndDisable => cd.set_alloc_cache_enabled(false), // Steady-state decode. The cap has to be re-armed // here and not only on the Enable edge: free VRAM // shrinks token by token as the KV cache grows, so @@ -2743,7 +2761,8 @@ impl Pipeline for NormalPipeline { ); if arc_layer_bisect() { let bufs = - crate::layers::arc_layer_trace_take(); + crate::layers::arc_layer_trace_take( + ); tracing::error!( "ARC BISECT: holding {} captured layer buffers", bufs.as_ref() @@ -3779,10 +3798,7 @@ mod alloc_cache_policy_tests { /// enabling the cache. #[test] fn degenerate_zero_length_forward_does_not_enable() { - assert_eq!( - alloc_cache_action(0, false, false), - AllocCacheAction::Leave - ); + assert_eq!(alloc_cache_action(0, false, false), AllocCacheAction::Leave); assert_eq!( alloc_cache_action(0, true, false), AllocCacheAction::DrainAndDisable diff --git a/mistralrs-core/src/pipeline/sampling.rs b/mistralrs-core/src/pipeline/sampling.rs index b7cc8045f..3e99d7212 100644 --- a/mistralrs-core/src/pipeline/sampling.rs +++ b/mistralrs-core/src/pipeline/sampling.rs @@ -558,24 +558,73 @@ fn device_loop_pre_sampled_token( seq: &mut Sequence, return_logprobs: bool, sample_speculative: bool, -) -> Option { +) -> Result> { + let seq_id = *seq.id(); let eligible = !return_logprobs && !sample_speculative && seq.sampler().is_greedy_trivial() && matches!(seq.recognizer, SequenceRecognizer::None); - arc_cuda_graph::set_device_loop_eligible(eligible); + arc_cuda_graph::set_device_loop_eligible(seq_id, eligible); + // Did the forward serve this step from the pending queue WITHOUT + // launching? Then the logits this sample holds alias graph-owned storage + // from an earlier step — sound only if a parked token is taken below. + // Consumed here, once, whatever the outcome. + let aliased_logits = arc_cuda_graph::take_aliased_logits_marker(); 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; + if aliased_logits { + candle_core::bail!( + "ArcGraph: this step's logits were served from the device loop's pending queue \ + for another sequence and alias graph-owned storage; this sequence (id {seq_id}) \ + is not device-loop eligible and cannot consume them. Failing the step rather \ + than sampling stale logits." + ); + } + return Ok(None); + } + match arc_cuda_graph::take_pending_token(seq_id) { + arc_cuda_graph::PendingTake::Taken(token) => Ok(Some(Logprobs { + token, + logprob: 0.0, + bytes: None, + top_logprobs: None, + })), + // Nothing parked: the forward really ran for this step (the pending + // short-circuit only fires with a non-empty queue), so the logits are + // real and the host samples them. `aliased_logits` cannot be set here + // in a consistent engine: it implies the queue was non-empty when the + // forward ran and empty now, on the same thread with no take between. + arc_cuda_graph::PendingTake::Empty => { + if aliased_logits { + candle_core::bail!( + "ArcGraph: the forward served this step from the pending queue but the queue \ + is empty at sampling time — the aliased logits cannot be sampled. Failing \ + the step." + ); + } + Ok(None) + } + // The queue belonged to another sequence and was dropped whole (and + // logged: "ArcGraph: dropped N foreign parked tokens"). If the + // forward ALSO short-circuited on that foreign queue, these logits + // are a stale alias and sampling them would hand this sequence a + // plausible token from the other sequence's distribution — fail the + // step loudly instead. Without the marker the forward genuinely ran + // (e.g. a batched eager step), the logits are this sequence's own, + // and sampling proceeds normally. + arc_cuda_graph::PendingTake::Foreign { dropped } => { + if aliased_logits { + candle_core::bail!( + "ArcGraph: this step's logits alias graph storage served against {dropped} \ + parked token(s) belonging to another sequence; sampling them would leak \ + another user's distribution into sequence {seq_id}. Failing the step." + ); + } + Ok(None) + } } - arc_cuda_graph::take_pending_token().map(|token| Logprobs { - token, - logprob: 0.0, - bytes: None, - top_logprobs: None, - }) } /// Async sample optionally adding to trie. @@ -595,7 +644,7 @@ pub async fn sample_sequence( // 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) + device_loop_pre_sampled_token(seq, return_logprobs, sample_speculative)? { return Ok(pre_sampled); } @@ -1033,8 +1082,10 @@ mod device_loop_cross_sequence_leak_tests { /// A greedy-trivial sequence: no temperature, penalties, bias, processors /// or recognizer — exactly the shape `device_loop_pre_sampled_token` - /// accepts, so the pending queue is live for it. - fn greedy_seq() -> Sequence { + /// accepts, so the pending queue is live for it. `id` matters: the + /// pending queue is owner-tagged by sequence id, so every test sequence + /// must carry a distinct one, as real requests do. + fn greedy_seq(id: usize) -> Sequence { let (dummy_sender, rx) = tokio::sync::mpsc::channel(1); LIVE_CLIENTS.with(|k| k.borrow_mut().push(rx)); let group = Arc::new(std::sync::Mutex::new(SequenceGroup::new( @@ -1059,7 +1110,7 @@ mod device_loop_cross_sequence_leak_tests { Sequence::new_waiting( vec![1u32; 4], String::new(), - 0, + id, 0, 1, dummy_sender, @@ -1130,7 +1181,7 @@ mod device_loop_cross_sequence_leak_tests { // ── Sequence A, decode step N: queue empty, host argmax, and the // pre-hook publishes A's eligibility — the same order the engine // runs (sample N publishes, forward N+1 bursts). - let mut seq_a = greedy_seq(); + let mut seq_a = greedy_seq(1); assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); // ── Forward N+1: the burst parks its FULL length (eos -1). @@ -1159,7 +1210,7 @@ mod device_loop_cross_sequence_leak_tests { ); // ── Sequence B's FIRST sample (the one right after its prefill). - let mut seq_b = greedy_seq(); + let mut seq_b = greedy_seq(2); let b_first = sample(&mut seq_b, PEAK_B, &rng).await; // ── The claim, part 2: B's token comes from B's own logits and is @@ -1198,7 +1249,7 @@ mod device_loop_cross_sequence_leak_tests { arc_cuda_graph::stand_down(); arc_cuda_graph::push_pending_tokens(&BURST); assert_eq!(arc_cuda_graph::pending_token_count(), BURST.len()); - let seq = greedy_seq(); + let seq = greedy_seq(3); seq.set_state(state); assert_eq!( arc_cuda_graph::pending_token_count(), @@ -1207,4 +1258,102 @@ mod device_loop_cross_sequence_leak_tests { ); } } + + /// The interleave the completion funnel CANNOT see: the DefaultScheduler + /// bucketing waitlist moves a running sequence aside **without a state + /// modification** (`default_scheduler.rs`, `bucket_and_waitlist_seqs_waiting` + /// says so in its doc), so no `set_state` ever fires for A. The owner tag + /// is the structural defence: B's first sample must consume nothing of + /// A's — the foreign queue is dropped whole (logged as + /// "ArcGraph: dropped N foreign parked tokens") rather than served, to A + /// included, whose device-side burst state is stale by then. + /// + /// Mutation check (run both ways before shipping): with the owner + /// comparison in `take_pending_token` forced to `true`, B's first token + /// comes back `22` — A's parked id — instead of `PEAK_B`. + #[tokio::test] + async fn a_waitlisted_sequence_with_no_state_change_cannot_leak_into_the_next() { + arc_cuda_graph::stand_down(); + let rng = Arc::new(std::sync::Mutex::new(Isaac64Rng::seed_from_u64(0))); + + // ── A (id 31) decodes at batch 1; its burst parks; it drains one + // (the control proving the take path is live in this build). + let mut seq_a = greedy_seq(31); + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); + arc_cuda_graph::push_pending_tokens(&BURST); + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, BURST[0] as u32); + + // ── A is WAITLISTED: no set_state, no funnel, nothing runs. Its + // undrained tail [22, 23, 24] is still parked when B is scheduled. + + // ── B (id 32) runs; its forward really ran, so its logits are real. + let mut seq_b = greedy_seq(32); + let b_first = sample(&mut seq_b, PEAK_B, &rng).await; + assert!( + !BURST.contains(&(b_first as i32)), + "cross-request token leak through the waitlist interleave: B's first sampled token \ + is {b_first}, one of A's parked burst tokens" + ); + assert_eq!(b_first, PEAK_B, "B must sample its own logits"); + + // ── A's foreign-tagged tail was dropped whole, not left to resurface. + assert_eq!( + arc_cuda_graph::pending_token_count(), + 0, + "the foreign queue must be dropped at B's take, not survive it" + ); + + // ── A resumes: nothing stale to consume, so it host-samples its own + // logits. (Its dropped tokens were computed against device state that + // no longer matches; replaying them would corrupt A's output.) + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); + } + + /// The double fault: the forward served a step from the pending queue + /// (returning the ALIASED stale logits tensor, launching nothing) and the + /// scheduler then swapped sequences before the sample ran. The taker is + /// foreign, so there is no parked token to return — and the logits in + /// hand are another sequence's stale graph output. Sampling them would + /// leak A's distribution into B; the step must fail loudly instead. + #[tokio::test] + async fn an_aliased_step_served_for_another_sequence_fails_instead_of_sampling_stale_logits() { + arc_cuda_graph::stand_down(); + let rng = Arc::new(std::sync::Mutex::new(Isaac64Rng::seed_from_u64(0))); + + // A (id 41) parks a burst; the forward short-circuits on it. + let mut seq_a = greedy_seq(41); + assert_eq!(sample(&mut seq_a, PEAK_A, &rng).await, PEAK_A); + arc_cuda_graph::push_pending_tokens(&BURST); + arc_cuda_graph::note_aliased_logits_served(); + + // B (id 42) is what actually samples that step. + let mut seq_b = greedy_seq(42); + let res = sample_sequence( + peaked_logits(PEAK_B), + &mut seq_b, + false, + rng.clone(), + false, + false, + false, + ) + .await; + assert!( + res.is_err(), + "a stale aliased-logits step served against a foreign queue must fail the step, \ + not return a plausible token" + ); + + // Control: the marker is harmless when the owner itself samples — the + // parked token is taken and the aliased tensor is never read. + let mut seq_c = greedy_seq(43); + assert_eq!(sample(&mut seq_c, PEAK_A, &rng).await, PEAK_A); + arc_cuda_graph::push_pending_tokens(&[25]); + arc_cuda_graph::note_aliased_logits_served(); + assert_eq!( + sample(&mut seq_c, PEAK_A, &rng).await, + 25, + "the owner's own aliased-served step must take its parked token normally" + ); + } } diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index c203fd761..27bb56d56 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -257,6 +257,20 @@ static REGISTRY: &[Capability] = &[ blocking condition as the top_k entry above.", }, }, + Capability { + name: "ArcGraph device decode loop opt-in gate (ARC_GRAPH_DEVICE_LOOP)", + parent: "ArcInfer / ArcGraph", + check: Check::Symbol { + symbol: "device_loop_enabled", + defined_in: "arc-cuda-graph/src/device_loop.rs", + }, + // The gate is opt-in default-off, but the CALL PATH must be live: + // `normal.rs::forward_inputs` consults `device_loop_enabled()` before + // `device_decode_burst`. If this reference disappears, the env var + // parses forever and engages nothing — the switched-off class this + // file exists to catch. + status: Status::Live, + }, ]; /// Symbols that are defined in Arc-owned source and have **no production