fix(ArcLab/ArcQuant): make prefill measurable, then ship the two kernels it proves were switched off

arc-cuda-graph/src/dedicated.rs

@@ -1106,6 +1106,31 @@ impl DedicatedDecodePath {
 #[cfg(feature = "cuda")]
 impl Drop for DedicatedDecodePath {
     fn drop(&mut self) {
+        // 🔴 Bind the context before freeing, exactly as `new()` does.
+        //
+        // `new()` binds first and its comment states the consequence of not
+        // doing so: "a SEGV in libcuda MOVAPS on the first cudaMalloc — no
+        // error code, just a fault", because the runtime's lazy-init walks a
+        // thread-local that has not been seeded. This `Drop` then performed the
+        // *same class* of raw Driver/Runtime calls — `cuGraphExecDestroy` plus
+        // a long list of `cudaFree` — with no bind at all.
+        //
+        // `unsafe impl Send`/`Sync` above make it legal for this value to be
+        // dropped on a thread that never touched CUDA, which is precisely what
+        // happens when the engine thread outlives `main` and teardown lands on
+        // a foreign thread. That asymmetry (bind in `new`, none in `drop`) is
+        // the second half of the teardown SIGSEGV this binary died from.
+        //
+        // Failure here is not fatal: if we cannot bind we are already in a
+        // teardown we cannot improve, and skipping the frees leaks into a
+        // process that is exiting anyway — strictly better than faulting.
+        if Self::bind_context(&self.weights).is_err() {
+            tracing::warn!(
+                "DedicatedDecodePath::drop: could not bind CUDA context; \
+                 skipping device frees to avoid faulting during teardown."
+            );
+            return;
+        }
         unsafe {
             if let Some(exec) = self.graph_exec {
                 cuGraphExecDestroy(exec);

mistralrs-bench/Cargo.toml

@@ -20,6 +20,7 @@ serde.workspace = true
 serde_json.workspace = true
 clap.workspace = true
 mistralrs-core.workspace = true
+mistralrs-quant.workspace = true
 tracing.workspace = true
 tokio.workspace = true
 cli-table.workspace = true

mistralrs-bench/src/main.rs

@@ -10,6 +10,7 @@ use mistralrs_core::{
 };
 use std::fmt::Display;
 use std::sync::Arc;
+use std::time::Instant;
 use tokio::sync::mpsc::channel;
 use tracing::{info, warn};
 
@@ -32,6 +33,13 @@ struct BenchResult {
     usages: Vec<Usage>,
     concurrency: usize,
     test_name: TestName,
+    /// Wall-clock seconds for the whole send->all-Done window, measured by the
+    /// client with no engine cooperation. This is the SECOND instrument: the
+    /// engine's own `Usage` numbers are only believable when they agree with a
+    /// clock that lives outside the engine.
+    wall_secs: f32,
+    /// Total prompt tokens the engine says it processed across all responses.
+    prompt_toks_seen: usize,
 }
 
 struct UncertainTokSec {
@@ -74,29 +82,41 @@ async fn run_bench(
     let sender = mistralrs.get_sender(None).unwrap();
     let (tx, mut rx) = channel(10_000);
 
-    let req = Request::Normal(Box::new(NormalRequest {
-        id: mistralrs.next_request_id(),
-        messages: prompt,
-        sampling_params: sampling_params.clone(),
-        response: tx,
-        return_logprobs: false,
-        is_streaming: false,
-        constraint: Constraint::None,
-        suffix: None,
-        tools: None,
-        tool_choice: None,
-        logits_processors: None,
-        return_raw_logits: false,
-        web_search_options: None,
-        model_id: None,
-        truncate_sequence: false,
-    }));
+    // 🔴 A fresh request id per send.
+    //
+    // This used to call `next_request_id()` exactly once and then `.clone()`
+    // the same `NormalRequest` for every concurrent copy and every repetition,
+    // so `concurrency * repetitions` sequences were in flight sharing ONE id.
+    // The id is the engine's handle for a sequence group; issuing duplicates
+    // means unrelated sequences alias each other in engine-side bookkeeping.
+    // A benchmark must not be the only caller in the system that does this.
+    let make_req = |id: usize| {
+        Request::Normal(Box::new(NormalRequest {
+            id,
+            messages: prompt.clone(),
+            sampling_params: sampling_params.clone(),
+            response: tx.clone(),
+            return_logprobs: false,
+            is_streaming: false,
+            constraint: Constraint::None,
+            suffix: None,
+            tools: None,
+            tool_choice: None,
+            logits_processors: None,
+            return_raw_logits: false,
+            web_search_options: None,
+            model_id: None,
+            truncate_sequence: false,
+        }))
+    };
 
     let mut usages = Vec::new();
+    let wall_start = Instant::now();
 
     for _ in 0..repetitions {
         for _ in 0..concurrency {
-            if sender.send(req.clone()).await.is_err() {
+            let req = make_req(mistralrs.next_request_id());
+            if sender.send(req).await.is_err() {
                 eprintln!("Receiver disconnected");
             }
         }
@@ -131,10 +151,15 @@ async fn run_bench(
         }
     }
 
+    let wall_secs = wall_start.elapsed().as_secs_f32();
+    let prompt_toks_seen = usages.iter().map(|u| u.prompt_tokens).sum();
+
     Ok(BenchResult {
         usages,
         concurrency,
         test_name,
+        wall_secs,
+        prompt_toks_seen,
     })
 }
 
@@ -186,6 +211,100 @@ fn get_ms_tok(result: &BenchResult) -> UncertainTokSec {
     UncertainTokSec { mean, std_dev }
 }
 
+/// A zero is not a measurement.
+///
+/// 🔴 This harness used to print `0.000±0.000` into a results table and exit 0.
+/// That is the house fault: a green result that proves no work happened. Every
+/// row is now required to carry evidence that the engine actually processed the
+/// tokens the row claims to be about, and a row that cannot is a hard stop.
+///
+/// Exits **2** (environment/instrument failure), never 1, so a caller can tell
+/// "the benchmark could not be run" apart from "the benchmark ran and regressed".
+fn assert_results_are_measurements(results: &[BenchResult]) {
+    let mut failures: Vec<String> = Vec::new();
+
+    for r in results {
+        let label = format!("{} @ concurrency {}", r.test_name, r.concurrency);
+
+        if r.usages.is_empty() {
+            failures.push(format!("{label}: no responses at all (0 usages collected)"));
+            continue;
+        }
+
+        for (i, u) in r.usages.iter().enumerate() {
+            let where_ = format!("{label}, response {i}");
+            match r.test_name {
+                TestName::Prompt(n) => {
+                    if u.prompt_tokens == 0 {
+                        failures.push(format!("{where_}: engine reported 0 prompt tokens"));
+                    } else if u.prompt_tokens != n {
+                        failures.push(format!(
+                            "{where_}: asked for {n} prompt tokens, engine processed {}",
+                            u.prompt_tokens
+                        ));
+                    }
+                    if u.total_prompt_time_sec <= 0.0 {
+                        failures.push(format!(
+                            "{where_}: prompt time is {:.6}s -- the prefill step was never timed",
+                            u.total_prompt_time_sec
+                        ));
+                    }
+                    if !u.avg_prompt_tok_per_sec.is_finite() || u.avg_prompt_tok_per_sec <= 0.0 {
+                        failures.push(format!(
+                            "{where_}: prompt rate is {} tok/s",
+                            u.avg_prompt_tok_per_sec
+                        ));
+                    }
+                }
+                TestName::Gen(_) => {
+                    if u.completion_tokens == 0 {
+                        failures.push(format!("{where_}: engine produced 0 completion tokens"));
+                    }
+                    if !u.avg_compl_tok_per_sec.is_finite() || u.avg_compl_tok_per_sec <= 0.0 {
+                        failures.push(format!(
+                            "{where_}: completion rate is {} tok/s",
+                            u.avg_compl_tok_per_sec
+                        ));
+                    }
+                }
+            }
+        }
+    }
+
+    // Which MoE kernel actually ran? The grouped-GEMM variant is chosen at
+    // RUNTIME from an env var, so the binary cannot tell you which one you
+    // measured. `grouped.rs` states the rule outright: "a harness that reports a
+    // per-variant timing MUST show this counter advancing on the variant it
+    // claims to have measured, and NOT advancing on the others, or the number
+    // describes some other kernel." Nothing in the tree printed it, so every
+    // previous variant comparison was unproven. Print it beside the results.
+    let counts = mistralrs_quant::grouped_launch_counts();
+    let total: u64 = counts.iter().sum();
+    eprintln!(
+        "MoE grouped-GEMM launches by variant: {counts:?} (selected variant = {}, total = {total})",
+        mistralrs_quant::grouped_variant()
+    );
+    if total == 0 {
+        eprintln!(
+            "  -> grouped GEMM never launched: this measurement is the per-pair \
+             gather GEMV (or the dequantize fallback), NOT the grouped kernel."
+        );
+    }
+
+    if !failures.is_empty() {
+        eprintln!("\n=== BENCHMARK PRODUCED NO VALID MEASUREMENT ===");
+        for f in &failures {
+            eprintln!("  - {f}");
+        }
+        eprintln!(
+            "\n{} check(s) failed. A zero-token or zero-rate result is an instrument\n\
+             failure, not a score, so no table is printed. Exiting 2.",
+            failures.len()
+        );
+        std::process::exit(2);
+    }
+}
+
 fn print_usage(model: &str, device: &Device, results: Vec<BenchResult>) {
     let backend = match device {
         Device::Cpu => "CPU",
@@ -205,6 +324,11 @@ fn print_usage(model: &str, device: &Device, results: Vec<BenchResult>) {
                 (get_tok_s(&r).mean * r.concurrency as f32)
                     .cell()
                     .justify(Justify::Right),
+                // Second instrument: a clock outside the engine.
+                format!("{:.3}", r.wall_secs).cell().justify(Justify::Right),
+                format!("{:.1}", r.prompt_toks_seen as f32 / r.wall_secs)
+                    .cell()
+                    .justify(Justify::Right),
             ]
         })
         .collect();
@@ -222,6 +346,8 @@ fn print_usage(model: &str, device: &Device, results: Vec<BenchResult>) {
             "ms/t".cell().bold(true),
             "concurrency".cell().bold(true),
             "throughput/s".cell().bold(true),
+            "wall s".cell().bold(true),
+            "prompt t/s (wall)".cell().bold(true),
         ])
         .bold(true);
     print_stdout(table).expect("print table");
@@ -575,6 +701,7 @@ async fn main() -> anyhow::Result<()> {
             results.push(r);
         }
 
+        assert_results_are_measurements(&results);
         print_usage(&model_name, &device, results);
     }
 

mistralrs-core/src/engine/mod.rs

@@ -537,6 +537,45 @@ impl Engine {
                         };
                         let _chunk_guard =
                             (!final_chunk).then(crate::pipeline::mark_prefill_intermediate);
+
+                        // Record prompt timing BEFORE step() so it is available if the
+                        // response is sent from inside step().
+                        //
+                        // 🔴 This is why `pp N` reported `0.000±0.000` for every prompt
+                        // benchmark. A prefill-only request (`max_len = 1`) finishes
+                        // *during* this prompt step: `pipeline/sampling.rs` calls
+                        // `seq.update_time_info()` and then `group.get_usage()` and
+                        // dispatches `Response::Done` from inside `step()`. But on this
+                        // (default, non-paged) path `prompt_timestamp` was only stamped
+                        // AFTER `step()` returned, ~100 lines below. So at the moment the
+                        // usage was built, `prompt_timestamp` was still `None`,
+                        // `update_time_info` skipped `group.total_prompt_time`
+                        // (`sequence.rs`, `if let Some(ts) = self.prompt_timestamp`), and
+                        // `get_usage` took the `total_prompt_time == 0` branch and
+                        // returned `avg_prompt_tok_per_sec: 0.0`. The harness then printed
+                        // that zero as if it were a measurement.
+                        //
+                        // The PagedAttention arm below already does exactly this, with
+                        // this same comment — the fix was simply never applied to the arm
+                        // that actually serves (PagedAttention is banned here: it shadows
+                        // the graph arm and measures zero tokens).
+                        //
+                        // `set_step_start_instant` gives `update_time_info` its in-flight
+                        // fallback (`start.elapsed()`) while the prompt step is running;
+                        // the post-step block below still overwrites `total_prompt_time`
+                        // with the precise `prompt_exec_time` for sequences that go on to
+                        // decode, so steady-state accounting is unchanged.
+                        {
+                            let now = SystemTime::now()
+                                .duration_since(UNIX_EPOCH)
+                                .expect("Time travel has occurred!")
+                                .as_millis();
+                            for seq in scheduled.prompt.iter_mut() {
+                                seq.prompt_timestamp = Some(now);
+                                seq.set_step_start_instant();
+                            }
+                        }
+
                         let prompt_exec_time = {
                             let mut pipeline = {
                                 let _s = arc_profiler::span("pipeline.lock");

mistralrs-core/src/lib.rs

@@ -556,6 +556,11 @@ impl MistralRsBuilder {
     }
 }
 
+/// How long `Drop` waits for an engine thread to finish its own teardown before
+/// giving up and detaching it. Generous because the thread may be mid-step on a
+/// long prefill; bounded because a hung engine must not wedge process exit.
+const ENGINE_JOIN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
+
 impl Drop for MistralRs {
     fn drop(&mut self) {
         // Terminate all engines
@@ -565,6 +570,53 @@ impl Drop for MistralRs {
                 let _ = engine.sender.try_send(Request::Terminate);
             }
         }
+
+        // 🔴 …and then WAIT for them. This used to be fire-and-forget.
+        //
+        // `Request::Terminate` is asynchronous: it only asks. Before this join,
+        // `main` returned while the engine thread was still alive and still
+        // releasing device memory, CUDA graph execs and stream objects. libc
+        // `exit()` then ran CUDA's `atexit` handler concurrently with that
+        // thread. Two threads tearing down the same CUDA state with no ordering
+        // is a textbook `corrupted double-linked list` / SIGSEGV, which is
+        // exactly how this binary died — *after* printing valid results, so the
+        // exit code libelled good runs. One chain already discarded a good,
+        // expensive measurement because its harness gated on `rc == 0`.
+        //
+        // The handle had no `.join()` anywhere in the workspace; its only uses
+        // were `is_finished()` polls. Joining here is what makes CUDA objects
+        // get freed on a live, context-bound thread before `exit()` runs.
+        //
+        // Bounded and poll-based rather than a bare `join()`: a wedged engine
+        // must not be able to hang process exit forever, and `JoinHandle` has no
+        // timed join in std.
+        let handles: Vec<(String, JoinHandle<()>)> = match self.engines.write() {
+            Ok(mut engines) => engines
+                .drain()
+                .map(|(id, instance)| (id, instance.engine_handler))
+                .collect(),
+            Err(_) => return,
+        };
+
+        let deadline = std::time::Instant::now() + ENGINE_JOIN_TIMEOUT;
+        for (id, handle) in handles {
+            while !handle.is_finished() && std::time::Instant::now() < deadline {
+                std::thread::sleep(std::time::Duration::from_millis(2));
+            }
+            if handle.is_finished() {
+                if handle.join().is_err() {
+                    tracing::warn!("Engine thread for `{id}` panicked during shutdown.");
+                }
+            } else {
+                // Detach rather than block forever. Say so: a silent detach here
+                // would reintroduce exactly the race this join exists to close.
+                tracing::warn!(
+                    "Engine thread for `{id}` did not finish within {:?}; detaching. \
+                     CUDA teardown may race with process exit.",
+                    ENGINE_JOIN_TIMEOUT
+                );
+            }
+        }
     }
 }