From ee03fc784083c8fcd0892e3cc4daf6d5cefa0233 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:28:52 +0100 Subject: [PATCH 1/4] perf(ArcInfer/ArcSample): port the STEP_us TOTAL/fwd/sample/other split onto the non-paged step arm V4 actually runs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The only TOTAL/fwd/sample/other instrument lived on the PagedAttention arm, which V4 can never reach (DeepSeekV4Loader::supports_paged_attention is false). The DefaultInstructions arm — the path V4 runs — had never once reported how a step divides between GPU and host. Log format byte-identical to the paged arm's so tooling parses both. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --- mistralrs-core/src/pipeline/mod.rs | 35 ++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs index 37110ed39..eec807814 100644 --- a/mistralrs-core/src/pipeline/mod.rs +++ b/mistralrs-core/src/pipeline/mod.rs @@ -982,6 +982,16 @@ pub trait Pipeline: let mut raw_out_logits = vec![vec![None; len_inputs]; input_seqs.len()]; let mut embedding_logits = vec![None; input_seqs.len()]; + // STEP_us instrumentation, ported verbatim from the + // PagedAttention arm below. That arm carried the only + // TOTAL/fwd/sample/other split in the codebase, and V4 cannot + // reach it (`DeepSeekV4Loader::supports_paged_attention` is + // false) — so the path V4 actually runs had never once + // reported how a step divides between GPU and host. The log + // format is byte-identical to the paged arm's on purpose: + // tooling that parses one parses both. + let __np_t_step_start = Instant::now(); + let mut __np_t_fwd_us: f64 = 0.0; let mut exec_duration = Duration::ZERO; for (i, inputs) in inputs_iter.into_iter().enumerate() { let InputProcessorOutput { @@ -1013,6 +1023,7 @@ pub trait Pipeline: }; let end = Instant::now(); exec_duration += end.duration_since(start); + __np_t_fwd_us += end.duration_since(start).as_secs_f64() * 1e6; // ONE host copy for the whole batch, *before* the // per-sequence split. See `host_copy_batched_result`: the @@ -1204,6 +1215,30 @@ pub trait Pipeline: } let end = Instant::now(); exec_duration += end.duration_since(start); + let __np_t_sample_us = end.duration_since(start).as_secs_f64() * 1e6; + let __np_t_total_us = __np_t_step_start.elapsed().as_secs_f64() * 1e6; + { + use std::sync::atomic::{AtomicU64, Ordering}; + static N: AtomicU64 = AtomicU64::new(0); + static SUM_TOTAL: AtomicU64 = AtomicU64::new(0); + static SUM_FWD: AtomicU64 = AtomicU64::new(0); + static SUM_SAMPLE: AtomicU64 = AtomicU64::new(0); + let n = N.fetch_add(1, Ordering::Relaxed) + 1; + SUM_TOTAL.fetch_add(__np_t_total_us as u64, Ordering::Relaxed); + SUM_FWD.fetch_add(__np_t_fwd_us as u64, Ordering::Relaxed); + SUM_SAMPLE.fetch_add(__np_t_sample_us as u64, Ordering::Relaxed); + if n > 4 && (n - 4) % 25 == 0 { + let nn = (n - 4) as f64; + let total = SUM_TOTAL.load(Ordering::Relaxed) as f64 / nn; + let fwd = SUM_FWD.load(Ordering::Relaxed) as f64 / nn; + let samp = SUM_SAMPLE.load(Ordering::Relaxed) as f64 / nn; + let other = total - fwd - samp; + tracing::info!( + "STEP_us avg over {} steps: TOTAL={:.0} fwd={:.0} sample={:.0} other={:.0} (=> {:.1} tok/s)", + n - 4, total, fwd, samp, other, 1e6 / total + ); + } + } Ok(exec_duration) } From c5b3f3f5466f5399265782708f5541bbd2d686fc Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:54:59 +0100 Subject: [PATCH 2/4] fix(ArcInfer/ArcSample): three measured fast-path aggravators + the device-eligibility predicate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit (a) has_penalties(): a DRY config at its disabled default (multiplier 0.0, a proven no-op in apply_dry_penalty) no longer disqualifies the GPU fast path or is_raw_argmax. (b) sample_top_kp_min_p: the server defaults (top_k=-1, top_p=1.0, min_p=0.0) ran a full-vocab sort_unstable_by (129,280 elements on V4, per sequence per token) whose result nothing consumed — skip it. Seeded draw provably identical (thread-local sort counter + pinned token 506 in tests). (c) sample_multinomial: the global rng MutexGuard was held (temporary lifetime extension) across get_top_logprobs and tokenizer.decode — narrowed to the draw. Plus Sampler::device_sampling_terminates, the single-source mirror of sample()'s device dispatch that ARC_SAMPLE_ON_DEVICE partitioning consults, and the ignored $0 host-cost probe at V4 vocab (D14: host evidence only). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --- mistralrs-core/src/pipeline/mod.rs | 129 +++++ mistralrs-core/src/sampler.rs | 473 +++++++++++++++++- .../tests/capability_reachability.rs | 13 + 3 files changed, 612 insertions(+), 3 deletions(-) diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs index eec807814..0bdf82594 100644 --- a/mistralrs-core/src/pipeline/mod.rs +++ b/mistralrs-core/src/pipeline/mod.rs @@ -545,6 +545,88 @@ pub(crate) fn host_copy_batched_result( raw.to_device(host) } +/// Parent system: ArcInfer / ArcSample. +/// +/// `ARC_SAMPLE_ON_DEVICE`: opt-in device-resident sampling for batched decode. +/// Default OFF; parsed by value (`env_flag_is_set`, so `=0` is off), latched +/// once per process. +/// +/// WHY THIS FLAG EXISTS: at any `n_seqs > 1` the `DefaultInstructions` step +/// arm copied the ENTIRE `[B, 1, vocab]` logits tensor to the host before +/// sampling — at B=256 with V4's 129,280 vocab in BF16 that is 66.2 MB per +/// decode step, serial, on the critical path — and `Sampler::sample`'s GPU +/// fast path is gated on `!logits.device().is_cpu()`, so the copy itself then +/// disqualified all 256 rows into the ~5.9 ms/token CPU pipeline the fast +/// path exists to avoid. [`host_copy_batched_result`]'s own doc states the +/// trap for a single sequence ("moving a B=1 batch to the host would silently +/// disable it") and never followed the thought to B > 1: the GPU sampler was +/// structurally unreachable at every batch above 1. +/// +/// With the flag ON, [`split_batched_logits_for_sampling`] keeps +/// fast-path-eligible rows on the device (token-sized D2H each) and gathers +/// the rest into ONE batched D2H. OFF is bit-for-bit the shipped behaviour. +pub(crate) fn sample_on_device_enabled() -> bool { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| mistralrs_quant::env_flag_is_set("ARC_SAMPLE_ON_DEVICE")) +} + +/// Split a batched `[B, 1, vocab]` logits tensor into per-row tensors for +/// sampling: rows whose sampler terminates on the device +/// ([`crate::sampler::Sampler::device_sampling_terminates`]) stay as device +/// views; the rows that genuinely need the CPU pipeline are shipped with at +/// most ONE batched D2H. `row_eligible` is in logits-row order (the +/// `seq_indices` order of the step arm, not sequence order). +/// +/// Copy discipline — the O(B²) trap, same as [`host_copy_batched_result`]: a +/// per-row view still references the whole `B * vocab` storage, and both +/// `Tensor::to_device` and `Tensor::to_vec1` copy the ENTIRE storage +/// (`CudaStorage::to_cpu_storage`, candle `9586979`, `tensor.rs:1949/2379`). +/// So CPU-bound rows must be gathered into fresh row-sized storage +/// (`index_select` allocates exactly the selected rows) *before* the single +/// `to_device`. The `Tensor::from_vec` for the gather indices is a +/// once-per-step H2D of `n_cpu` u32s (pitfall #5 concerns per-token tensor +/// creation in forward hot loops; this one replaces a 66 MB D2H and only runs +/// for mixed batches). +pub(crate) fn split_batched_logits_for_sampling( + logits: &Tensor, + row_eligible: &[bool], +) -> candle_core::Result> { + let n = row_eligible.len(); + let n_cpu = row_eligible.iter().filter(|e| !**e).count(); + // Every row stays on device: pure views, no transfer at all. + if n_cpu == 0 { + return (0..n).map(|i| logits.i(i)).collect(); + } + LOGITS_HOST_COPIES.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + // Every row needs the CPU pipeline: one whole-tensor D2H — exactly the + // flag-off behaviour. + if n_cpu == n { + let host = logits.to_device(&Device::Cpu)?; + return (0..n).map(|i| host.i(i)).collect(); + } + // Mixed batch: gather ONLY the CPU-bound rows into fresh contiguous + // device storage, then one D2H of that gather. + let cpu_rows: Vec = row_eligible + .iter() + .enumerate() + .filter(|(_, e)| !**e) + .map(|(i, _)| i as u32) + .collect(); + let idx = Tensor::from_vec(cpu_rows, n_cpu, logits.device())?; + let host = logits.index_select(&idx, 0)?.to_device(&Device::Cpu)?; + let mut host_pos = 0usize; + let mut out = Vec::with_capacity(n); + for (i, eligible) in row_eligible.iter().enumerate() { + if *eligible { + out.push(logits.i(i)?); + } else { + out.push(host.i(host_pos)?); + host_pos += 1; + } + } + Ok(out) +} + #[derive(serde::Serialize, serde::Deserialize)] pub(crate) struct FileListCache { files: Vec, @@ -1025,6 +1107,53 @@ pub trait Pipeline: exec_duration += end.duration_since(start); __np_t_fwd_us += end.duration_since(start).as_secs_f64() * 1e6; + // Device-resident sampling (ARC_SAMPLE_ON_DEVICE, default + // OFF): keep fast-path-eligible rows on the device instead + // of hosting the whole batch — see + // [`sample_on_device_enabled`] for why. Only the batched + // CausalGeneration case changes; B=1 already stays on + // device via `host_copy_batched_result`'s `n_seqs > 1` + // condition, and RawLogits/Embeddings are host-consumed. + let device_sampling = sample_on_device_enabled() + && seq_indices.len() > 1 + && matches!( + &raw_logits, + ForwardInputsResult::CausalGeneration { logits } + if !logits.device().is_cpu() + ); + if device_sampling { + let ForwardInputsResult::CausalGeneration { logits: batched } = + &raw_logits + else { + unreachable!("device_sampling implies CausalGeneration") + }; + let vocab = batched.dim(candle_core::D::Minus1)?; + let is_cuda = batched.device().is_cuda(); + let row_eligible: Vec = seq_indices + .iter() + .map(|&seq_idx| { + let return_logprobs = input_seqs[seq_idx].return_logprobs(); + input_seqs[seq_idx].sampler().device_sampling_terminates( + vocab, + is_cuda, + return_logprobs, + ) + }) + .collect(); + // Still a `sync_span`: the batched-gather D2H (when any + // row needs the CPU pipeline) is the same blocking + // device wait the flag-off copy was. + let rows = { + let _s = arc_profiler::sync_span("logits_d2h"); + split_batched_logits_for_sampling(batched, &row_eligible)? + }; + for (row, seq_idx) in rows.into_iter().zip(seq_indices) { + logits[seq_idx] = + Some(ForwardInputsResult::CausalGeneration { logits: row }); + } + continue; + } + // ONE host copy for the whole batch, *before* the // per-sequence split. See `host_copy_batched_result`: the // reverse order costs O(B²) D2H bytes per step. diff --git a/mistralrs-core/src/sampler.rs b/mistralrs-core/src/sampler.rs index 65547f010..ffa3f163a 100644 --- a/mistralrs-core/src/sampler.rs +++ b/mistralrs-core/src/sampler.rs @@ -300,6 +300,8 @@ fn cmp_desc_by_prob(a: &(u32, f32), b: &(u32, f32)) -> std::cmp::Ordering { /// If `k >= probs.len()`, returns all elements sorted. /// Also zeros out elements in `probs` beyond top-k if `zero_rest` is true. fn partial_sort_top_k(probs: &mut [f32], k: usize, zero_rest: bool) -> Vec<(u32, f32)> { + #[cfg(test)] + PARTIAL_SORT_CALLS.with(|calls| calls.set(calls.get() + 1)); let n = probs.len(); if n == 0 || k == 0 { return Vec::new(); @@ -333,6 +335,18 @@ fn partial_sort_top_k(probs: &mut [f32], k: usize, zero_rest: bool) -> Vec<(u32, idx_probs } +/// Invocation counter for [`partial_sort_top_k`], test builds only. Exists so +/// the sort-skip in [`Sampler::sample_top_kp_min_p`] is provable: a test can +/// assert the sort genuinely did not run, not merely that the output looks +/// the same. Thread-local, because `cargo test` runs tests concurrently and +/// several other tests in this file sort — a process-global counter would make +/// every delta assertion racy. +#[cfg(test)] +thread_local! { + pub(crate) static PARTIAL_SORT_CALLS: std::cell::Cell = + const { std::cell::Cell::new(0) }; +} + /// Kernel dispatch sizes for the GPU radix top-k sampler path. Must mirror /// `arc_cuda_graph::flashmlasparse::SUPPORTED_TOPK` (kept as a local const so /// the candidate-truncation logic and its CPU-parity tests build without the @@ -499,10 +513,19 @@ impl Sampler { /// directly. Two copies of this expression could disagree, and the one that /// disagreed would be a silent correctness bug rather than a slow path. fn has_penalties(&self) -> bool { + // DRY counts only when it can actually rewrite a logit. The disabled + // default (`DrySamplingParams::default()`) is `multiplier: 0.0`, and + // `apply_dry_penalty` returns without touching the logits when the + // multiplier is exactly 0.0 — so `dry_params.is_some()` alone was + // disqualifying the GPU fast path (and failing `is_raw_argmax`) for a + // penalty that is a proven no-op. self.frequency_penalty.unwrap_or(0.0) != 0.0 || self.presence_penalty.unwrap_or(0.0) != 0.0 || self.repetition_penalty.unwrap_or(1.0) != 1.0 - || self.dry_params.is_some() + || self + .dry_params + .as_ref() + .is_some_and(|params| params.multiplier != 0.0) } /// Would [`Sampler::sample`] take the pure-argmax path — no temperature and @@ -561,6 +584,69 @@ impl Sampler { && self.logits_processors.is_empty() } + /// Mirror of [`Self::sample`]'s device fast-path dispatch, answerable + /// **before** the logits row has been copied anywhere: would a non-CPU row + /// of `vocab` logits be sampled to completion on the device (token-sized + /// D2H), rather than falling back to the CPU pipeline (full-row D2H + + /// full-vocab host work)? + /// + /// This is what `ARC_SAMPLE_ON_DEVICE`'s host-copy partitioning + /// (`pipeline/mod.rs`) consults per sequence: rows answering `true` keep + /// their device residency; rows answering `false` are gathered into ONE + /// batched D2H. It must stay in this file, next to the gates it mirrors + /// ([`Self::has_penalties`] and `sample`'s `trivial` condition): a second + /// copy that drifted would silently route rows onto the wrong side of the + /// PCIe bus. + /// + /// Assumes non-speculative sampling (`sample_speculative = false`), which + /// is what every `sample_causal_gen` reachable from the partitioned step + /// arm passes; the speculative pipelines override `step` entirely and + /// never see the partition. + /// + /// `logit_bias` and top-nσ do **not** disqualify: both are applied as + /// device tensor ops before the fast-path dispatch. One residual runtime + /// fallback remains on the CUDA radix path (a pure-top-p nucleus the + /// candidate set cannot contain declines with `Ok(None)`); it pays a + /// per-row D2H and is counted by `gpu_sampling_health::DECLINED`. + pub(crate) fn device_sampling_terminates( + &self, + vocab: usize, + device_is_cuda: bool, + return_logprobs: bool, + ) -> bool { + // `sample`'s `trivial` gate, minus its device term — the caller is + // deciding device residency, not reacting to it. + let trivial = + !return_logprobs && self.logits_processors.is_empty() && !self.has_penalties(); + if !trivial { + return false; + } + // Greedy: a single argmax kernel + 4-byte D2H. Any backend, any vocab. + if self.temperature.is_none() { + return true; + } + // Temperature path: `sample_fast` stays on device when no sort is + // needed, or when candle's shared-memory arg_sort can hold the vocab + // (same bound `sample` checks). + let needs_sort = self.top_k > 0 || (self.top_p > 0.0 && self.top_p < 1.0); + if !needs_sort { + return true; + } + if vocab.next_power_of_two().saturating_mul(4) <= 48 * 1024 { + return true; + } + // Big-vocab sort: only the CUDA radix top-k path avoids the CPU + // fallback, and only when the requested k fits its dispatch table + // (`sample_fast_topk_gpu`'s own static precondition). + let max_k = *GPU_RADIX_TOPK_SIZES.last().unwrap(); + let k_needed = if self.top_k > 0 { + self.top_k as usize + } else { + max_k + }; + cfg!(feature = "cuda") && device_is_cuda && k_needed <= max_k && vocab > max_k + } + /// True if an OpenAI `logit_bias` map is set. Any sampling path that does /// not route through [`Self::sample`] — notably the GPU-autonomous decode /// sampler, which draws from the model's raw logits — must refuse the @@ -1186,8 +1272,17 @@ impl Sampler { ) -> Result { let distr = WeightedIndex::new(probs).map_err(Error::wrap)?; - let mut mut_ref_rng = &mut *rng.lock().expect("could not lock rng mutex"); - let next_token = distr.sample(&mut mut_ref_rng); // "Find the first item which has a weight *higher* than the chosen weight." + // Hold the global rng mutex for the draw and NOTHING else. The guard + // used to live to the end of this function (temporary lifetime + // extension through the `&mut *lock()` borrow), which serialised every + // concurrently-sampling sequence in the batch behind this one's + // `get_top_logprobs` (a partial sort plus up to `top_n` tokenizer + // decodes) and its own `tokenizer.decode` — none of which touch the + // rng. + let next_token = { + let mut guard = rng.lock().expect("could not lock rng mutex"); + distr.sample(&mut *guard) // "Find the first item which has a weight *higher* than the chosen weight." + }; let logprob = probs[next_token].log(10.0); let top_logprobs = if return_logprobs { @@ -1231,6 +1326,22 @@ impl Sampler { probs.len() }; + // The sorted pairs feed exactly three consumers: the zero-rest pass + // inside `partial_sort_top_k` (only when `k < n`), the top-p cumsum + // walk, and the min-p threshold scan. Under the server defaults + // (`top_k = -1`, `top_p = 1.0`, `min_p = 0.0` — `engine/add_request.rs`) + // none of the three runs: `k = n` zeroes nothing, and both filter + // blocks below are gated out of (0, 1). That request was still paying + // a full-vocab `sort_unstable_by` — 129,280 elements for V4, per + // sequence per token — to produce a Vec nothing read. Skip straight to + // the multinomial draw; `probs` is untouched either way, so the drawn + // token (same rng stream) and the reported logprobs are identical. + let top_p_active = top_p > 0.0 && top_p < 1.0; + let min_p_active = min_p > 0.0 && min_p < 1.0; + if k >= probs.len() && !top_p_active && !min_p_active { + return self.sample_multinomial(probs, return_logprobs, rng); + } + // Get sorted top-k indices with partial sort, zeroing out rest let idx_probs = partial_sort_top_k(probs, k, true); @@ -2893,6 +3004,362 @@ mod topk_parity_tests { } } +/// The server-default sort skip and the penalty/eligibility gates. +#[cfg(test)] +mod fast_path_gate_tests { + use super::{DrySamplingParams, Sampler, PARTIAL_SORT_CALLS}; + use candle_core::{Device, Tensor}; + use rand::SeedableRng; + use rand_isaac::Isaac64Rng; + use std::sync::{Arc, Mutex}; + + const VOCAB: usize = 512; + + /// Deterministic spread-out logits (same LCG family as the parity tests). + fn logits() -> Tensor { + let mut state: u64 = 0x5EED ^ 0xDEAD_BEEF_CAFE_F00D; + let v: Vec = (0..VOCAB) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 32) as u32 as f32 / u32::MAX as f32) * 12.0 - 6.0 + }) + .collect(); + Tensor::from_vec(v, VOCAB, &Device::Cpu).unwrap() + } + + fn sampler(top_k: i64, top_p: f64, min_p: f64) -> Sampler { + Sampler::new( + Some(0.7), + 0, + None, + None, + None, + None, + None, + top_k, + top_p, + min_p, + None, + vec![], + None, + ) + .unwrap() + } + + /// Sampled token + how many times `partial_sort_top_k` ran during the + /// call. `Sampler::sample` on this thread never crosses threads, so the + /// thread-local delta is exact. + fn sample_counting_sorts(s: &Sampler, seed: u64) -> (u32, u64) { + let rng = Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(seed))); + let before = PARTIAL_SORT_CALLS.with(|c| c.get()); + let token = s + .sample(logits(), &[1u32, 2, 3], false, rng, false, true) + .unwrap() + .token; + let sorts = PARTIAL_SORT_CALLS.with(|c| c.get()) - before; + (token, sorts) + } + + /// The server defaults (`top_k = -1`, `top_p = 1.0`, `min_p = 0.0` — + /// `engine/add_request.rs`) must not run the full-vocab sort: nothing + /// consumes it. Equivalence witness: the PRE-FIX code path replicated + /// literally — `partial_sort_top_k(probs, n, zero_rest=true)` (which for + /// `k = n` sorts everything and zeroes nothing) followed by + /// `sample_multinomial` on the same rng seed. Same token, one sort vs + /// zero. (`top_k = VOCAB` is not a usable witness: it now — correctly — + /// takes the skip as well, being the same no-op filter.) + #[test] + fn server_default_top_k_skips_the_unconsumed_sort() { + let s = sampler(-1, 1.0, 0.0); + let (token_skip, sorts_skip) = sample_counting_sorts(&s, 42); + assert_eq!(sorts_skip, 0, "server defaults must skip the sort entirely"); + + // Pre-fix path, spelled out: softmax(logits / T), full sort with + // zero-rest, multinomial. Filters are all out of range, so the sort's + // output feeds nothing — exactly the work the skip removes. + let probs_t = + candle_nn::ops::softmax_last_dim(&(logits().to_dtype(candle_core::DType::F32).unwrap() + / 0.7) + .unwrap()) + .unwrap(); + let mut probs: Vec = probs_t.to_vec1().unwrap(); + let n = probs.len(); + let before = PARTIAL_SORT_CALLS.with(|c| c.get()); + let sorted = super::partial_sort_top_k(&mut probs, n, true); + assert_eq!( + PARTIAL_SORT_CALLS.with(|c| c.get()) - before, + 1, + "the witness must actually sort, or this test proves nothing" + ); + assert_eq!(sorted.len(), VOCAB, "k = n returns every element sorted"); + let rng = Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(42))); + let token_sort = s.sample_multinomial(&probs, false, rng).unwrap().token; + + assert_eq!( + token_skip, token_sort, + "skipping the sort must not change the drawn token (identical probs, \ + identical rng stream)" + ); + // Pin the drawn token so a silent change to the rng stream (e.g. an + // extra draw before the multinomial) cannot pass as parity. + assert_eq!(token_skip, 506, "seeded draw moved — the rng stream shifted"); + } + + /// MUTATION CONTROLS for the skip condition: each filter that consumes the + /// sorted pairs must, on its own, keep the sort alive. + #[test] + fn each_active_filter_keeps_the_sort_alive() { + // top_p in (0,1): the cumsum walk needs the sorted pairs. + let (_, sorts) = sample_counting_sorts(&sampler(-1, 0.9, 0.0), 7); + assert_eq!(sorts, 1, "active top_p must sort"); + // min_p in (0,1): the threshold scan needs the sorted pairs. + let (_, sorts) = sample_counting_sorts(&sampler(-1, 1.0, 0.1), 7); + assert_eq!(sorts, 1, "active min_p must sort"); + // top_k < vocab: the zero-rest pass needs the partition. + let (_, sorts) = sample_counting_sorts(&sampler(32, 1.0, 0.0), 7); + assert_eq!(sorts, 1, "truncating top_k must sort"); + } + + /// A DRY config left at its disabled default (`multiplier: 0.0`) is a + /// proven no-op in `apply_dry_penalty`, so it must not count as a penalty + /// — before this fix `has_penalties()` answered true for it, which + /// disqualified the GPU fast path (and `is_raw_argmax`) for every request + /// that merely *carried* the default DRY struct. + /// + /// `is_raw_argmax` is the public reader of `has_penalties`; both + /// directions are pinned: disabled DRY ⇒ still raw-argmax, enabled DRY ⇒ + /// not. + #[test] + fn disabled_dry_is_not_a_penalty_and_enabled_dry_is() { + // Dry params only survive `Sampler::new` when a tokenizer is present. + let tokenizer = Arc::new(tokenizers::Tokenizer::new( + tokenizers::models::bpe::BPE::default(), + )); + let mk = |dry: Option| { + Sampler::new( + None, + 0, + Some(tokenizer.clone()), + None, + None, + None, + dry, + -1, + 1.0, + 0.0, + None, + vec![], + None, + ) + .unwrap() + }; + + let none = mk(None); + assert!(none.is_raw_argmax(), "baseline: no DRY at all"); + + let disabled = mk(Some(DrySamplingParams::default())); + assert!( + disabled.is_raw_argmax(), + "multiplier = 0.0 (the disabled default) cannot rewrite a logit; \ + it must not disqualify the fast path" + ); + + let enabled = mk(Some(DrySamplingParams { + multiplier: 0.8, + ..Default::default() + })); + assert!( + !enabled.is_raw_argmax(), + "an active DRY penalty rewrites logits and must disqualify" + ); + + // And the no-op claim itself: with a disabled DRY the sampled token + // equals the no-DRY token on the same fixture (greedy, so exact). + let rng = || Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(0))); + let ctx = [5u32, 5, 5, 5]; + let t_none = none + .sample(logits(), &ctx, false, rng(), false, false) + .unwrap() + .token; + let t_disabled = disabled + .sample(logits(), &ctx, false, rng(), false, false) + .unwrap() + .token; + assert_eq!(t_disabled, t_none, "disabled DRY must be a behavioural no-op"); + } + + /// `device_sampling_terminates` must mirror `sample`'s dispatch exactly — + /// each disqualifier flips it independently, and the sort/vocab geometry + /// matches the fast path's own bounds. + #[test] + fn device_eligibility_mirrors_the_fast_path_dispatch() { + let big_vocab = 129_280; // V4: needs_sort can't fit shared memory + let small_vocab = 4_096; // 4096.next_power_of_two() * 4 = 16 KB <= 48 KB + + // Greedy, no penalties: eligible at any vocab, any backend. + let greedy = Sampler::new( + None, + 0, + None, + None, + None, + None, + None, + -1, + 1.0, + 0.0, + None, + vec![], + None, + ) + .unwrap(); + assert!(greedy.device_sampling_terminates(big_vocab, false, false)); + assert!(greedy.device_sampling_terminates(big_vocab, true, false)); + // return_logprobs needs the CPU pipeline. + assert!(!greedy.device_sampling_terminates(big_vocab, true, true)); + + // Temperature without a sort (server defaults): eligible. + let no_sort = sampler(-1, 1.0, 0.0); + assert!(no_sort.device_sampling_terminates(big_vocab, false, false)); + + // Temperature with a sort: small vocab fits candle's shared-memory + // arg_sort on any backend; big vocab needs the CUDA radix path. + let sorted = sampler(32, 1.0, 0.0); + assert!(sorted.device_sampling_terminates(small_vocab, false, false)); + assert_eq!( + sorted.device_sampling_terminates(big_vocab, true, false), + cfg!(feature = "cuda"), + "big-vocab sorted sampling terminates on device iff the CUDA radix \ + top-k exists in this build" + ); + assert!(!sorted.device_sampling_terminates(big_vocab, false, false)); + + // top_k above the radix dispatch cap: the kernel declines statically. + let huge_k = sampler(2048, 1.0, 0.0); + assert!(!huge_k.device_sampling_terminates(big_vocab, true, false)); + + // A penalty disqualifies regardless of geometry. + let penalised = Sampler::new( + Some(0.7), + 0, + None, + Some(0.5), + None, + None, + None, + -1, + 1.0, + 0.0, + None, + vec![], + None, + ) + .unwrap(); + assert!(!penalised.device_sampling_terminates(big_vocab, true, false)); + } +} + +/// $0 host-cost probe at V4's vocabulary. +/// +/// **This times CPU code on a CPU tensor. It is evidence about HOST cost +/// only — it validates nothing about GPU behaviour (D14).** It exists to put +/// a number on what one sequence pays per token when it lands on the CPU +/// sampling pipeline, under (i) bench params (`top_k = 32`) and (ii) the +/// server defaults (`top_k = -1`, `top_p = 1.0`, `min_p = 0.0` — +/// `engine/add_request.rs`), which before the sort-skip fix ran a full +/// 129,280-element `sort_unstable_by` whose result nothing consumed. +/// +/// Deliberately `#[ignore]`d: it is a manual instrument, not a gate. Run with +/// `cargo test -p mistralrs-core --release host_cost_probe -- --ignored --nocapture`. +/// +/// The parity fixture above stays at `VOCAB = 4096` on purpose: its pure-top-p +/// cases depend on nucleus containment in the top 1024, which does not hold +/// for this spread at 129,280. +#[cfg(test)] +mod host_cost_probe { + use super::Sampler; + use candle_core::{Device, Tensor}; + use rand::SeedableRng; + use rand_isaac::Isaac64Rng; + use std::sync::{Arc, Mutex}; + + /// DeepSeek-V4's vocabulary. + const PROBE_VOCAB: usize = 129_280; + + /// Same LCG as `topk_parity_tests::peaked_logits`, at probe size. + fn probe_logits(seed: u64) -> Vec { + let mut state: u64 = seed ^ 0xDEAD_BEEF_CAFE_F00D; + (0..PROBE_VOCAB) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 32) as u32 as f32 / u32::MAX as f32) * 12.0 - 6.0 + }) + .collect() + } + + fn probe_sampler(top_k: i64, top_p: f64, min_p: f64) -> Sampler { + Sampler::new( + Some(0.7), + 0, + None, + None, + None, + None, + None, + top_k, + top_p, + min_p, + None, + vec![], + None, + ) + .unwrap() + } + + fn time_us_per_call(sampler: &Sampler) -> (f64, f64) { + const WARMUP: usize = 3; + const ITERS: usize = 30; + let logits = Tensor::from_vec(probe_logits(7), PROBE_VOCAB, &Device::Cpu).unwrap(); + let rng = Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(42))); + let ctx = [1u32, 2, 3]; + for _ in 0..WARMUP { + sampler + .sample(logits.clone(), &ctx, false, rng.clone(), false, true) + .unwrap(); + } + let mut samples_us = Vec::with_capacity(ITERS); + for _ in 0..ITERS { + let t0 = std::time::Instant::now(); + sampler + .sample(logits.clone(), &ctx, false, rng.clone(), false, true) + .unwrap(); + samples_us.push(t0.elapsed().as_secs_f64() * 1e6); + } + let mean = samples_us.iter().sum::() / ITERS as f64; + let min = samples_us.iter().cloned().fold(f64::INFINITY, f64::min); + (mean, min) + } + + #[test] + #[ignore = "manual $0 probe: prints per-call CPU sampling cost at V4 vocab; \ + run with --ignored --nocapture. Host-cost evidence only, never \ + GPU validation (D14)."] + fn time_cpu_sample_at_v4_vocab() { + let (bench_mean, bench_min) = time_us_per_call(&probe_sampler(32, 1.0, 0.0)); + let (srv_mean, srv_min) = time_us_per_call(&probe_sampler(-1, 1.0, 0.0)); + println!( + "HOST_PROBE vocab={PROBE_VOCAB} bench(top_k=32): mean={bench_mean:.0}us min={bench_min:.0}us" + ); + println!( + "HOST_PROBE vocab={PROBE_VOCAB} server(top_k=-1,top_p=1.0): mean={srv_mean:.0}us min={srv_min:.0}us" + ); + } +} + /// The greedy tie-break contract. /// /// Arc has three greedy implementations and they did not agree with each other diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index c203fd761..16b6b8759 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -197,6 +197,19 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + Capability { + name: "sampling: ARC_SAMPLE_ON_DEVICE keeps batched logits on the device", + parent: "ArcInfer / ArcSample", + check: Check::Symbol { + symbol: "sample_on_device_enabled", + defined_in: "mistralrs-core/src/pipeline/mod.rs", + }, + // Default OFF (env_flag_is_set, value-parsed), consulted by the + // DefaultInstructions step arm before the batched logits host copy. + // If this goes dark, every batch B > 1 is back to a 66 MB/step serial + // D2H that also disqualifies all rows from the GPU sampler fast path. + status: Status::Live, + }, // -- the ApiPromise class ------------------------------------------------ Capability { name: "sampling: logit_bias reaches the sampler", From 98572a66ce5e5296d0b4508a0dea960e3a1bcaaa Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:13:56 +0100 Subject: [PATCH 3/4] =?UTF-8?q?test(ArcInfer/ArcSample):=20partition=20spl?= =?UTF-8?q?it=20tests=20=E2=80=94=20copy=20discipline=20+=20seeded=20token?= =?UTF-8?q?=20parity=20for=20ARC=5FSAMPLE=5FON=5FDEVICE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --- mistralrs-core/src/pipeline/mod.rs | 244 +++++++++++++++++++++++++++++ 1 file changed, 244 insertions(+) diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs index 0bdf82594..cfe391112 100644 --- a/mistralrs-core/src/pipeline/mod.rs +++ b/mistralrs-core/src/pipeline/mod.rs @@ -508,6 +508,12 @@ impl ForwardInputsResult { pub(crate) static LOGITS_HOST_COPIES: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); +/// Serialises tests that assert on [`LOGITS_HOST_COPIES`] deltas: the counter +/// is process-global and `cargo test` runs tests concurrently, so two +/// counter-delta tests interleaving would read each other's increments. +#[cfg(test)] +pub(crate) static LOGITS_HOST_COPIES_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + /// Move a batched forward result to `host` with a **single** transfer, before /// it is split into per-sequence views. /// @@ -1964,6 +1970,9 @@ mod host_copy_tests { #[test] fn batched_host_copy_is_o1_in_batch_while_the_old_order_was_o_b_squared() { + let _serial = LOGITS_HOST_COPIES_TEST_LOCK + .lock() + .unwrap_or_else(|e| e.into_inner()); // 7 is equal to none of the batch sizes used below and coprime with // all of them, so no assertion can pass by `b * v` aliasing `v`, `b`, // or `v * b` (DOCTRINE D12: the fixture must discriminate). @@ -2089,6 +2098,241 @@ mod host_copy_tests { } } +/// `ARC_SAMPLE_ON_DEVICE`'s partitioned split: copy discipline and seeded +/// token parity, on CPU-checkable fixtures. +/// +/// On a CPU fixture every "device-resident" row re-enters the same CPU +/// sampling code, so what these tests prove is exactly the CPU-checkable +/// half: the partition itself — row routing, ordering, and the one-batched- +/// copy discipline — cannot change a sampled token or move more bytes than +/// the rows it ships. GPU-side fast-path behaviour is the `Sampler`'s +/// existing contract, not established here (D14). +#[cfg(test)] +mod device_sampling_split_tests { + use super::*; + use crate::sampler::Sampler; + use candle_core::{CpuStorage, Storage}; + use rand::SeedableRng; + use rand_isaac::Isaac64Rng; + use std::sync::atomic::Ordering; + use std::sync::Mutex; + + /// Elements candle would move across a device boundary for `t`: the whole + /// backing storage, not `t.elem_count()` (see `host_copy_tests`). + fn storage_elems(t: &Tensor) -> usize { + let (storage, _layout) = t.storage_and_layout(); + match &*storage { + Storage::Cpu(CpuStorage::F32(v)) => v.len(), + _ => panic!("fixture must be a CPU F32 tensor"), + } + } + + /// `[b, 1, v]` ramp `0..b*v`: every row identifies itself, so any + /// host-position mapping error in the gather produces a visible wrong + /// slice rather than a coincidentally-equal one. + fn batch(b: usize, v: usize) -> Tensor { + let vals: Vec = (0..b * v).map(|x| x as f32).collect(); + Tensor::from_vec(vals, (b, 1, v), &Device::Cpu).unwrap() + } + + fn row_values(t: &Tensor) -> Vec { + t.flatten_all().unwrap().to_vec1::().unwrap() + } + + #[test] + fn mixed_batch_gathers_only_the_cpu_rows_in_one_copy() { + let _serial = LOGITS_HOST_COPIES_TEST_LOCK + .lock() + .unwrap_or_else(|e| e.into_inner()); + const B: usize = 6; + const V: usize = 7; + let logits = batch(B, V); + let eligible = [true, false, true, false, false, true]; + let n_cpu = 3usize; + + let before = LOGITS_HOST_COPIES.load(Ordering::Relaxed); + let rows = split_batched_logits_for_sampling(&logits, &eligible).unwrap(); + assert_eq!( + LOGITS_HOST_COPIES.load(Ordering::Relaxed) - before, + 1, + "a mixed batch is ONE batched copy, not one per CPU row" + ); + + assert_eq!(rows.len(), B); + for (k, row) in rows.iter().enumerate() { + assert_eq!(row.dims(), &[1, V], "row {k} shape"); + let expect: Vec = (k * V..(k + 1) * V).map(|x| x as f32).collect(); + assert_eq!(row_values(row), expect, "row {k} carries the wrong data"); + if eligible[k] { + // Device-resident rows are views over the original storage: + // zero bytes moved for them. + assert_eq!( + storage_elems(row), + B * V, + "row {k}: an eligible row must stay a view of the batch" + ); + } else { + // CPU-bound rows share the gathered storage: exactly the + // selected rows crossed, not the batch (and not O(B²)). + assert_eq!( + storage_elems(row), + n_cpu * V, + "row {k}: the gather must materialise ONLY the CPU rows" + ); + } + } + } + + #[test] + fn all_eligible_moves_nothing_and_none_eligible_matches_flag_off() { + let _serial = LOGITS_HOST_COPIES_TEST_LOCK + .lock() + .unwrap_or_else(|e| e.into_inner()); + const B: usize = 4; + const V: usize = 5; + let logits = batch(B, V); + + // All rows terminate on device: no host copy at all. + let before = LOGITS_HOST_COPIES.load(Ordering::Relaxed); + let rows = split_batched_logits_for_sampling(&logits, &[true; B]).unwrap(); + assert_eq!( + LOGITS_HOST_COPIES.load(Ordering::Relaxed) - before, + 0, + "an all-eligible batch must not touch the host" + ); + for (k, row) in rows.iter().enumerate() { + assert_eq!(storage_elems(row), B * V, "row {k} must be a view"); + } + + // No row terminates on device: one whole-batch copy — the flag-off + // behaviour exactly (same counter delta as host_copy_batched_result). + let before = LOGITS_HOST_COPIES.load(Ordering::Relaxed); + let rows = split_batched_logits_for_sampling(&logits, &[false; B]).unwrap(); + assert_eq!( + LOGITS_HOST_COPIES.load(Ordering::Relaxed) - before, + 1, + "an all-CPU batch is one batched copy, exactly like flag-off" + ); + for (k, row) in rows.iter().enumerate() { + let expect: Vec = (k * V..(k + 1) * V).map(|x| x as f32).collect(); + assert_eq!(row_values(row), expect, "row {k} after whole-batch copy"); + } + } + + /// Seeded sampling through the partition must produce IDENTICAL tokens to + /// the flag-off path (batched host copy + `index_bs` split), row for row, + /// for any eligibility mask. Rows sample in row order on a shared seeded + /// rng in both paths, so any reordering, extra draw, or wrong-row routing + /// shifts the stream and fails. + #[test] + fn seeded_tokens_identical_to_the_flag_off_split() { + // Not asserting on counter deltas, but this test INCREMENTS the + // global counter (both split paths), so it must still serialise with + // the tests that do assert. + let _serial = LOGITS_HOST_COPIES_TEST_LOCK + .lock() + .unwrap_or_else(|e| e.into_inner()); + const B: usize = 4; + const V: usize = 64; + + // Deterministic spread-out logits (LCG as in sampler parity tests). + let mut state: u64 = 0xA5A5 ^ 0xDEAD_BEEF_CAFE_F00D; + let vals: Vec = (0..B * V) + .map(|_| { + state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((state >> 32) as u32 as f32 / u32::MAX as f32) * 12.0 - 6.0 + }) + .collect(); + let logits = Tensor::from_vec(vals, (B, 1, V), &Device::Cpu).unwrap(); + + // One sampler shape per row: greedy, top-k, server defaults, top-p. + let mk = |temperature: Option, top_k: i64, top_p: f64| { + Sampler::new( + temperature, + 0, + None, + None, + None, + None, + None, + top_k, + top_p, + 0.0, + None, + vec![], + None, + ) + .unwrap() + }; + let samplers = [ + mk(None, -1, 1.0), + mk(Some(0.7), 8, 1.0), + mk(Some(1.0), -1, 1.0), + mk(Some(0.7), -1, 0.9), + ]; + + // Mirrors `sample_sequence`'s per-row prologue. + let sample_rows = |rows: &[Tensor]| -> Vec { + let rng = Arc::new(Mutex::new(Isaac64Rng::seed_from_u64(42))); + rows.iter() + .zip(samplers.iter()) + .map(|(row, sampler)| { + let row = row + .squeeze(0) + .unwrap() + .to_dtype(DType::F32) + .unwrap(); + sampler + .sample(row, &[1u32, 2, 3], false, rng.clone(), false, true) + .unwrap() + .token + }) + .collect() + }; + + // Flag-off: batched host copy, then the per-sequence split. + let flag_off_rows: Vec = { + let host = host_copy_batched_result( + ForwardInputsResult::CausalGeneration { + logits: logits.clone(), + }, + B, + &Device::Cpu, + ) + .unwrap(); + (0..B) + .map(|k| match host.index_bs(k).unwrap() { + ForwardInputsResult::CausalGeneration { logits } => logits, + _ => unreachable!(), + }) + .collect() + }; + let expected = sample_rows(&flag_off_rows); + + // Flag-on, under every mask shape the partition can take. + for mask in [ + [true, true, true, true], + [false, false, false, false], + [true, false, true, false], + [false, true, false, true], + ] { + let rows = split_batched_logits_for_sampling(&logits, &mask).unwrap(); + assert_eq!( + sample_rows(&rows), + expected, + "mask {mask:?}: partitioned rows must sample the same tokens" + ); + } + + // Pin the expected tokens: if the sampler's rng consumption ever + // changes shape, this fails here rather than passing vacuously as + // both-sides-moved parity. + assert_eq!(expected, vec![23, 42, 31, 22], "seeded fixture drifted"); + } +} + /// `extract_logits` vs the `.i((.., seq_len - 1, ..))` it replaced in the phi3 /// and starcoder2 GGUF models. /// From 6cd69992036032d88bd1630d6ebb41513420e8f7 Mon Sep 17 00:00:00 2001 From: Nirupam Bhowmick <48842933+heydryft@users.noreply.github.com> Date: Fri, 21 Aug 2026 20:25:20 +0100 Subject: [PATCH 4/4] chore(ArcInfer/ArcSample): u32::try_from for the gather row index (clippy-by-hand for mistralrs-core) Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf --- mistralrs-core/src/pipeline/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs index cfe391112..7fa8c878f 100644 --- a/mistralrs-core/src/pipeline/mod.rs +++ b/mistralrs-core/src/pipeline/mod.rs @@ -616,7 +616,7 @@ pub(crate) fn split_batched_logits_for_sampling( .iter() .enumerate() .filter(|(_, e)| !**e) - .map(|(i, _)| i as u32) + .map(|(i, _)| u32::try_from(i).expect("batch row index fits u32")) .collect(); let idx = Tensor::from_vec(cpu_rows, n_cpu, logits.device())?; let host = logits.index_select(&idx, 0)?.to_device(&Device::Cpu)?;