diff --git a/.github/workflows/cuda_compile_check.yaml b/.github/workflows/cuda_compile_check.yaml index 2aa62730b..2bd1c949c 100644 --- a/.github/workflows/cuda_compile_check.yaml +++ b/.github/workflows/cuda_compile_check.yaml @@ -305,6 +305,16 @@ jobs: - name: QTIP beam spill gate (LOCAL:0, all geometries) run: bash arc-tools/qtip_beam_res_usage_check.sh sm_${{ matrix.compute_cap }} + # Same free-instrument class as the QTIP gate above, for the blockwise- + # FP8 dense-GEMM lane: asserts all 4 fp8_matmul_tiled instantiations + # exist (2 dtypes x 2 POW2_SCALE scale-index arms), prints the LDS width + # census that decides whether the bank-conflict cost model describes the + # compiled kernel, and HARD-FAILS if the never-executed WMMA GEMM + # contains no tensor-core (HMMA/HGMMA) instructions. nvcc + cuobjdump + # only; no GPU. + - name: FP8 GEMM SASS probe (LDS census + WMMA HMMA presence) + run: bash arc-tools/fp8_gemm_sass_check.sh sm_${{ matrix.compute_cap }} + - name: Compile QTIP CUDA kernels (mistralrs-quant) run: cargo build -p mistralrs-quant --lib --features cuda diff --git a/.typos.toml b/.typos.toml index 7de3b414e..01c67f7c6 100644 --- a/.typos.toml +++ b/.typos.toml @@ -60,6 +60,13 @@ extend-ignore-re = [ # optin : "opt-in" without the hyphen — CUDA's # cudaDevAttrMaxSharedMemoryPerBlockOptin, and filenames / # branches that cannot carry a hyphen there +# arange : candle's `Tensor::arange` (numpy naming), which also +# appears as a word segment inside larger identifiers +# (`cached_arange_f32`, `ARANGE_CACHE`) — the +# extend-identifiers entries below only cover the exact +# identifier `arange`, not segments +"arange" = "arange" +"ARANGE" = "ARANGE" "nd" = "nd" "clen" = "clen" "tpos" = "tpos" diff --git a/Cargo.lock b/Cargo.lock index 5a68f4570..778418644 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -692,7 +692,7 @@ checksum = "ade8366b8bd5ba243f0a58f036cc0ca8a2f069cff1a2351ef1cac6b083e16fc0" [[package]] name = "candle-core" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "accelerate-src", "byteorder", @@ -723,7 +723,7 @@ dependencies = [ [[package]] name = "candle-flash-attn" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "anyhow", "candle-core", @@ -734,7 +734,7 @@ dependencies = [ [[package]] name = "candle-flash-attn-v3" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "anyhow", "candle-core", @@ -747,7 +747,7 @@ dependencies = [ [[package]] name = "candle-kernels" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "cudaforge", ] @@ -755,7 +755,7 @@ dependencies = [ [[package]] name = "candle-metal-kernels" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "half", "objc2", @@ -769,7 +769,7 @@ dependencies = [ [[package]] name = "candle-nn" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "accelerate-src", "candle-core", @@ -788,7 +788,7 @@ dependencies = [ [[package]] name = "candle-ug" version = "0.9.2" -source = "git+https://github.com/aeonmindai/candle.git?rev=89ab14ef1216331a782539f533b99f6816708604#89ab14ef1216331a782539f533b99f6816708604" +source = "git+https://github.com/aeonmindai/candle.git?rev=9586979db9031df6377f7dab1233c4548d976b76#9586979db9031df6377f7dab1233c4548d976b76" dependencies = [ "ug", "ug-cuda", @@ -3600,6 +3600,7 @@ dependencies = [ "clap", "cli-table", "mistralrs-core", + "mistralrs-quant", "serde", "serde_json", "tokio", diff --git a/Cargo.toml b/Cargo.toml index b8852835d..ef611a2af 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -44,16 +44,31 @@ rust-version = "1.88" # RUN-161: forked candle for the opt-in CUDA caching allocator (graph-capture-safe). # Default-off; identical to upstream c3bb5bf when ARC_CANDLE_ALLOC_CACHE unset. # Fork of huggingface/candle 0.9.2 (rev c3bb5bf), mirrored to the org for CI. -# PINNED TO AN EXPLICIT REV, not a branch. This change depends on -# `drain_alloc_cache_and_free` (fork commit 88d86a20), and a moving branch -# pointer has broken CI on this repo before — a path/branch dependency that -# advances under a merged PR turns a green build into a red one with no commit -# here to blame. -candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } -candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } -candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } -candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } -candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "89ab14ef1216331a782539f533b99f6816708604" } +# PINNED TO AN EXPLICIT REV, not a branch. This tree needs BOTH lines of the +# fork's allocator work, and a moving branch pointer has broken CI on this repo +# before — a path/branch dependency that advances under a merged PR turns a +# green build into a red one with no commit here to blame. +# +# The rev below is `arcgraph/bounded-alloc-plus-prewarm` (9586979db), which is +# the MERGE of the two, and a strict superset of each: +# +# * 89ab14ef (`perf/alloc-cache-bounded`) — `drain_alloc_cache_and_free` +# (88d86a20) plus the bounded caching allocator #213 depends on, its cap +# contract test, and "eviction must not run inside the capture window". +# * 1fa534db (`arcgraph/prewarm-profile`) — the capture-miss ledger and +# profile-driven pre-warm this file's pre-warm caller needs +# (`prewarm_alloc_cache`, `reset_capture_misses`, `capture_misses`). +# +# Those two revs had DIVERGED: 1fa534db is missing three commits 89ab14ef has. +# Pinning straight to 1fa534db, as this change originally did, would therefore +# have silently reverted the bounded allocator and the capture-window eviction +# fix. Verified by `gh api repos/aeonmindai/candle/compare/...`: 9586979db is +# ahead of 89ab14ef by 6 and of 1fa534db by 8, behind neither. +candle-core = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } +candle-nn = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } +candle-flash-attn-v3 = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } +candle-flash-attn = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } +candle-metal-kernels = { git = "https://github.com/aeonmindai/candle.git", rev = "9586979db9031df6377f7dab1233c4548d976b76" } axum = "0.8.8" anyhow = "1.0.100" diff --git a/README.md b/README.md index 15297bc17..ba023c3dc 100644 --- a/README.md +++ b/README.md @@ -28,7 +28,7 @@ Forked from [mistral.rs](https://github.com/EricLBuehler/mistral.rs). Apache 2.0 | Feature | What it does | |---|---| | **QTIP 2-bit weights** | Cornell ICLR'25 — trellis quantization with Hadamard incoherence rotation. 8× weight compression vs FP16. Two rungs: `qtip2` (LUT) and `qtip2b` (computed codebook), the latter being what Arc serves. Works on any model. | -| **TurboQuant KV** *(paged default at head_dim 128)* | Zandieh et al. (Google Research, [arXiv:2504.19874](https://arxiv.org/abs/2504.19874), ICLR'26) — Arc Rust implementation of WHT + Lloyd-Max KV coding. **On by default on the paged path**: a standard-layout head_dim-128 model on CUDA gets K4/V3 KV with no flag, and loses prefix caching with it (`--pa-cache-type auto` opts out). Every other geometry falls back to `auto` with a warning; there is no kernel at head_dim 512, so DeepSeek V4 does not take it. Served end-to-end on a B200 at **55 tok/s with correct output**; **quality has never been evaluated**. Eager (non-paged) KV is a separate path, opt-in via `ARC_TURBOQUANT_KV=1`. | +| **TurboQuant KV** *(paged default at head_dim 128)* | Zandieh et al. (Google Research, [arXiv:2504.19874](https://arxiv.org/abs/2504.19874), ICLR'26) — Arc Rust implementation of WHT + Lloyd-Max KV coding. **On by default on the paged path**: a standard-layout head_dim-128 model on CUDA gets K4/V3 KV with no flag, and loses prefix caching with it (`--pa-cache-type auto` opts out). The kernels are now instantiated at head_dim 64/128/256/512, but only 128 has ever executed on hardware and only 128 is reachable without asking — every other geometry falls back to `auto` with a warning, so DeepSeek V4's head_dim of 512 is opt-in via an explicit `--pa-cache-type turboquant`. Served end-to-end on a B200 at **55 tok/s with correct output** at head_dim 128; **quality has never been evaluated**, at any width. Eager (non-paged) KV is a separate path, opt-in via `ARC_TURBOQUANT_KV=1`. | | **TD-MoE Tucker + whitening** | Whitened Tucker decomposition of the MoE expert pool. The **"lossless 20%"** figure is the **paper's** (NeurIPS'25), not ours — *published, not reproduced by us*. Works on any MoE. | | **ArcAttention/Indexer** — sparse attention | DeepSeek V4's Lightning Indexer with FlashMLASparse CUDA kernels (MIT, ported from sgl-project) for top-k attention. Dense O(n²) → sparse O(n·k). | | **ArcGraph** — GPU-autonomous decode | Full decode loop (forward + sample + EOS check) on GPU. Zero CPU sync per token. Works on any model. | @@ -79,7 +79,7 @@ Each layer compresses a different axis, and in principle the wins multiply. **Read the block above as arithmetic, because that is all it is.** The weight line is measured end-to-end: 284B → **~74 GB** at 2.09 bits/param (V4 Flash trellis bake, serving on one H200 — see [Results](#results)). **The KV lines are not.** They are bytes-per-token ratios computed from the packed format; no run has ever produced a measured KV compression figure, and the "4.27×" once published here is retracted. -What *has* run: the TurboQuant K4/V3 paged path served Qwen3-32B on a B200 end-to-end with correct output (see [Results](#results)) — so "no TurboQuant forward pass has ever been benchmarked", which this section previously asserted, was false. What that run does **not** establish is any of the ratios above, or any quality number. TurboQuant is the paged default at head_dim 128; it has no kernel at V4's head_dim of 512. The xKV pool is not implemented. +What *has* run: the TurboQuant K4/V3 paged path served Qwen3-32B on a B200 end-to-end with correct output (see [Results](#results)) — so "no TurboQuant forward pass has ever been benchmarked", which this section previously asserted, was false. What that run does **not** establish is any of the ratios above, or any quality number, and it was at head_dim 128 only. TurboQuant is the paged default at head_dim 128; the paged kernels now also cover V4's head_dim of 512, but nothing has ever run at that width, so it stays opt-in. The xKV pool is not implemented. For long context (1M+ tokens), the bottleneck shifts from weights to attention compute + KV bandwidth. Arc handles that via FlashMLASparse (CUDA kernel, MIT-licensed, ported from sgl-project), turning dense attention's O(n²) into sparse top-k O(n·k) on models with native sparse-attention training (DeepSeek V3.2+ family) and via top-k attention + sink preservation on the rest. @@ -144,9 +144,9 @@ Python bindings: `pip install mistralrs-cuda`. > **TurboQuant *is* the paged default** — correcting an earlier note here that read "TurboQuant is not the default anywhere", which was false when written. `defaults::PAGED_CACHE_TYPE` is `PagedCacheType::TurboQuant` and `--pa-cache-type` carries no clap default, so leaving it unset on CUDA (where PagedAttention is itself on by default) gives a standard-layout head_dim-128 model TurboQuant K4/V3 KV with nothing asked for — and **silently drops prefix caching**, which no TurboQuant preset supports. Pass `--pa-cache-type auto`, or `PagedCacheType::Auto` as above, to opt out. > -> Off that envelope — MLA layouts, any other head_dim — the default falls back to `Auto` with a warning, while an explicit `--pa-cache-type turboquant` is a hard error instead. The eager (non-paged) KV path is separate and genuinely opt-in via `ARC_TURBOQUANT_KV=1`. +> Off that envelope the default falls back to `Auto` with a warning: MLA layouts, uninstantiated head dims, and the instantiated-but-unmeasured widths 64/256/512. Those three are reachable only by asking for them — an explicit `--pa-cache-type turboquant` is accepted at any instantiated width, and is a hard error off that set. The eager (non-paged) KV path is separate and genuinely opt-in via `ARC_TURBOQUANT_KV=1`. > -> **Provenance:** this path has served a real model on real hardware — Qwen3-32B, B200, 55 tok/s, correct output — but at b=1, one width, one preset, with **no quality evaluation and no A/B against an unquantized cache**. See [Results](#results) before relying on it. +> **Provenance:** this path has served a real model on real hardware — Qwen3-32B, B200, 55 tok/s, correct output — but at b=1, **head_dim 128 only**, one preset, with **no quality evaluation and no A/B against an unquantized cache**. Widths 64/256/512 compile; none has ever executed. See [Results](#results) before relying on it. ## Supported Models diff --git a/arc-cuda-graph/src/autonomous.rs b/arc-cuda-graph/src/autonomous.rs index 2539b2801..6153d7699 100644 --- a/arc-cuda-graph/src/autonomous.rs +++ b/arc-cuda-graph/src/autonomous.rs @@ -78,6 +78,8 @@ pub struct AutonomousDecodeConfig { pub top_p: f32, pub frequency_penalty: f32, pub presence_penalty: f32, + /// `-1` disables top-k. The previous in-graph sampler had no top-k at all. + pub top_k: i32, pub greedy: bool, } @@ -101,7 +103,19 @@ pub struct AutonomousDecodeRunner { /// True if the graph uses a WHILE conditional node (CUDA 12.4+). /// False if host-driven loop (body graph launched per step). uses_while_node: bool, - rng_offset: u64, + /// The real sampler: temperature + top-k + top-p, with its RNG state held + /// in **device memory** and advanced by the kernel. + /// + /// 🔴 The sampler this replaces (`launch_fused_top_p_bf16`) took + /// `rng_offset` as a **by-value kernel argument**. A captured graph bakes + /// kernel arguments at capture time, so every replayed step drew the + /// identical uniform and the host-side `rng_offset += 1` could not reach + /// an instantiated graph. It was also not nucleus sampling: it walked the + /// vocabulary in **token-id order** accumulating full-distribution mass to + /// `top_p * u`, biasing hard toward low token ids (its own comment said + /// "Not quite right — proper top-p needs sorting"). A device-resident + /// `rng_state` is the only form that survives replay. + sampler: crate::sampling_cuda::CudaSampler, } #[cfg(feature = "cuda")] @@ -171,6 +185,17 @@ impl AutonomousDecodeRunner { std::ptr::write_bytes(ring_write_head_ptr as *mut u8, 0, batch * 4); } + // Bound to the runner's own stream so its kernels land inside the + // captured body rather than on the device-default stream. + let mut sampler = crate::sampling_cuda::CudaSampler::new( + device, + config.padded_batch_size, + config.vocab_size, + candle_core::DType::BF16, + 0x5EED_A11C_E571_2345, + )?; + sampler.set_stream(stream); + Ok(Self { config, device: device.clone(), @@ -183,7 +208,7 @@ impl AutonomousDecodeRunner { ring_size, graph_exec: None, uses_while_node: false, - rng_offset: 0, + sampler, }) } @@ -350,10 +375,30 @@ impl AutonomousDecodeRunner { candle_core::bail!("cuGraphCreate failed: {status}"); } - // 2. Create conditional handle (WHILE, default_value=1 = loop) + // 2. Create conditional handle (WHILE, default_value=1 = loop). + // + // `CU_GRAPH_COND_ASSIGN_DEFAULT` is REQUIRED: without it the default + // is never applied at launch, the condition reads 0, and the WHILE + // body runs zero times while every CUDA call still returns success. + let mut ctx: CUcontext = std::ptr::null_mut(); + let status = unsafe { cuCtxGetCurrent(&mut ctx) }; + if status != CUDA_SUCCESS || ctx.is_null() { + unsafe { + cuGraphDestroy(outer_graph); + } + tracing::warn!("cuCtxGetCurrent failed ({status}), falling back to host-driven loop"); + return self.capture_body_graph(forward_fn, bs, vocab); + } let mut cond_handle: CUgraphConditionalHandle = 0; - let status = - unsafe { cudaGraphConditionalHandleCreate(&mut cond_handle, outer_graph, 1, 0) }; + let status = unsafe { + cuGraphConditionalHandleCreate( + &mut cond_handle, + outer_graph, + ctx, + 1, + CU_GRAPH_COND_ASSIGN_DEFAULT, + ) + }; if status != CUDA_SUCCESS { unsafe { cuGraphDestroy(outer_graph); @@ -365,18 +410,19 @@ impl AutonomousDecodeRunner { // 3. Add conditional WHILE node to outer graph // This creates an empty body graph that we populate via stream capture. let mut while_node: CUgraphNode = std::ptr::null_mut(); - let mut body_graph: CUgraph = std::ptr::null_mut(); + // `phGraph_out` is an OUT field the driver populates; leave it null. let mut params: CudaGraphNodeParams = unsafe { std::mem::zeroed() }; params.node_type = CudaGraphNodeType::Conditional; params.conditional.handle = cond_handle; params.conditional.cond_type = CUgraphConditionalNodeType::WHILE; params.conditional.size = 1; - params.conditional.body_graph_out = &mut body_graph; + params.conditional.ctx = ctx; let status = unsafe { - cudaGraphAddNode( - outer_graph, + cuGraphAddNode( &mut while_node, + outer_graph, + std::ptr::null(), std::ptr::null(), 0, &mut params, @@ -390,19 +436,51 @@ impl AutonomousDecodeRunner { return self.capture_body_graph(forward_fn, bs, vocab); } - // 4. Populate body graph via stream capture + // 4. Populate the conditional node's OWN body graph. + // + // The driver hands back the body graph in `phGraph_out`; capture + // straight into it with `cuStreamBeginCaptureToGraph`. The previous + // code captured with `cuStreamBeginCapture_v2` — which always creates + // a NEW graph — and then destroyed the result, leaving the conditional + // body EMPTY. An empty body generates no tokens. + let body_graph: CUgraph = unsafe { + if params.conditional.body_graph_out.is_null() { + cuGraphDestroy(outer_graph); + candle_core::bail!("conditional node returned no body graph array"); + } + *params.conditional.body_graph_out + }; + if body_graph.is_null() { + unsafe { + cuGraphDestroy(outer_graph); + } + candle_core::bail!("conditional node body graph is null"); + } + unsafe { - let status = cuStreamBeginCapture_v2(self.stream, CUstreamCaptureMode::THREAD_LOCAL); + // RELAXED, not THREAD_LOCAL: candle's allocator and helper streams + // create cross-stream dependencies that THREAD_LOCAL rejects with + // CUDA_ERROR_STREAM_CAPTURE_ISOLATION (same reason `graph.rs` + // uses RELAXED). + let status = cuStreamBeginCaptureToGraph( + self.stream, + body_graph, + std::ptr::null(), + std::ptr::null(), + 0, + CUstreamCaptureMode::RELAXED, + ); if status != CUDA_SUCCESS { cuGraphDestroy(outer_graph); - candle_core::bail!("cuStreamBeginCapture for WHILE body failed: {status}"); + candle_core::bail!("cuStreamBeginCaptureToGraph for WHILE body failed: {status}"); } } // Capture: forward → sample → step_update → check_done_conditional self.capture_body_kernels(forward_fn, bs, vocab, Some(cond_handle))?; - // End capture into the body graph + // Ends the capture; returns the graph we captured INTO (`body_graph`), + // which is already attached to the conditional node. let mut captured_body: CUgraph = std::ptr::null_mut(); unsafe { let status = cuStreamEndCapture(self.stream, &mut captured_body); @@ -411,29 +489,37 @@ impl AutonomousDecodeRunner { candle_core::bail!("cuStreamEndCapture for WHILE body failed: {status}"); } } - - // The captured_body needs to be merged into the body_graph that - // cudaGraphAddNode created. For CUDA 12.4 conditional nodes, - // the body_graph_out is pre-created and we should have captured - // INTO it by using it as the capture target. However, stream - // capture always creates a new graph. - // - // The correct approach: don't use stream capture for the body. - // Instead, add kernel nodes to body_graph directly. But that's - // extremely complex (need to manually create kernel nodes for - // every cuBLAS call, every custom kernel, etc.). - // - // Alternative: use cudaStreamBeginCaptureToGraph (CUDA 12.3+) - // which captures into an existing graph. - // - // For now: instantiate the outer graph. If the body graph was - // properly populated by the conditional node setup, it works. - // If not, we destroy and fall back. - unsafe { - cuGraphDestroy(captured_body); + if captured_body != body_graph { + unsafe { + cuGraphDestroy(outer_graph); + } + candle_core::bail!( + "capture-to-graph returned a different graph than the conditional body" + ); + } + // Assert the body actually received nodes. A conditional node with an + // empty body is the failure this function exists to prevent, and it is + // invisible at launch: the graph runs and produces nothing. + let mut body_nodes: usize = 0; + let status = unsafe { cuGraphGetNodes(body_graph, std::ptr::null_mut(), &mut body_nodes) }; + if status != CUDA_SUCCESS || body_nodes == 0 { + unsafe { + cuGraphDestroy(outer_graph); + } + candle_core::bail!( + "WHILE body graph has {body_nodes} nodes after capture (status {status}) — \ + the decode body did not record" + ); } + tracing::info!("CUDA graph: WHILE body captured with {body_nodes} nodes"); - // 5. Instantiate outer graph + // 5. Instantiate outer graph. + // + // NOTE: no AUTO_FREE_ON_LAUNCH here. `cuda.h:1971` restricts a + // conditional body to "kernel nodes, empty nodes, child graphs, + // memsets, memcopies, and conditionals" — memory alloc/free nodes are + // NOT in that list, so the captured decode body must be + // allocation-free for this path to instantiate at all. let mut exec: CUgraphExec = std::ptr::null_mut(); let status = unsafe { cuGraphInstantiate_v2( @@ -510,7 +596,7 @@ impl AutonomousDecodeRunner { /// Capture the body kernels: forward → sample → step_update → check_done. /// Called during stream capture (kernels recorded, not executed). fn capture_body_kernels( - &self, + &mut self, forward_fn: &F, bs: i32, vocab: i32, @@ -521,48 +607,54 @@ impl AutonomousDecodeRunner { { // Forward pass let logits = forward_fn()?; - let logits_ptr = tensor_ptr(&logits)? as *const _; - // Sampling - let sampled_ptr = tensor_ptr(&self.decode_state.sampled_tokens)? as *mut _; - - unsafe { - if self.config.greedy { - launch_fused_argmax_bf16( - logits_ptr, - sampled_ptr, - std::ptr::null_mut(), - vocab, - bs, - self.stream, - ); - } else { - if self.config.frequency_penalty != 0.0 || self.config.presence_penalty != 0.0 { - launch_apply_penalties( - logits_ptr as *mut _, - tensor_ptr(&self.decode_state.output_tokens)? as *const i32, - tensor_ptr(&self.decode_state.n_generated)? as *const i32, - self.config.frequency_penalty, - self.config.presence_penalty, - vocab, - self.config.max_tokens as i32, - bs, - self.stream, - ); - } - launch_fused_top_p_bf16( - logits_ptr, - sampled_ptr, - self.config.temperature, - self.config.top_p, + // Cheap handle clones so every immutable borrow of `self` is finished + // before we take `&mut self.sampler`. + let sampled = self.decode_state.sampled_tokens.clone(); + + // Penalties stay a separate in-place pre-pass over the logits. The + // fused sampler applies penalties from a [batch, vocab] count tensor + // (`freq_counts`), which we do not maintain; passing `None` there + // while leaving non-zero penalties in the config would drop them + // SILENTLY (`sampling_kernel.cu:167` — `fcnt = freq_counts ? .. : + // nullptr`). So apply them here and zero them in `cfg`. + if self.config.frequency_penalty != 0.0 || self.config.presence_penalty != 0.0 { + let logits_mut = tensor_ptr(&logits)? as *mut _; + let out_toks = tensor_ptr(&self.decode_state.output_tokens)? as *const i32; + let n_gen = tensor_ptr(&self.decode_state.n_generated)? as *const i32; + unsafe { + launch_apply_penalties( + logits_mut, + out_toks, + n_gen, + self.config.frequency_penalty, + self.config.presence_penalty, vocab, + self.config.max_tokens as i32, bs, - 42, - self.rng_offset, self.stream, ); } + } + + let cfg = crate::sampling_cpu::SamplingConfig { + temperature: self.config.temperature, + top_p: self.config.top_p, + top_k: self.config.top_k, + // applied above, in-place, so the kernel must not re-apply them + frequency_penalty: 0.0, + presence_penalty: 0.0, + greedy: self.config.greedy, + eos_token_id: self.config.eos_token_id, + }; + // Scoped so the mutable borrow ends before the immutable ones below. + { + self.sampler.sample(&logits, None, cfg, &sampled)?; + } + let sampled_ptr = tensor_ptr(&sampled)? as *mut i32; + + unsafe { // Step update launch_decode_step_update( sampled_ptr as *const i32, @@ -621,9 +713,8 @@ impl AutonomousDecodeRunner { candle_core::Error::Msg("Graph not captured — call capture() first".into()) })?; - // Reset state - self.decode_state - .reset(&self.device, self.config.padded_batch_size)?; + // Reset state in place — the captured graph holds these addresses. + self.decode_state.reset(self.stream)?; unsafe { std::ptr::write_bytes( self.ring_write_head_ptr as *mut u8, @@ -631,7 +722,6 @@ impl AutonomousDecodeRunner { self.config.padded_batch_size * 4, ); } - self.rng_offset += 1; if self.uses_while_node { // ============================================================ @@ -656,7 +746,9 @@ impl AutonomousDecodeRunner { } cudaStreamSynchronize(self.stream); } - let cond = self.decode_state.loop_condition.to_vec1::()?; + // I32: the kernel writes `int32_t`. Reading this as i64 + // consumed two adjacent kernel writes as one value. + let cond = self.decode_state.loop_condition.to_vec1::()?; if cond[0] == 0 { break; } @@ -668,7 +760,7 @@ impl AutonomousDecodeRunner { .decode_state .output_tokens .to_dtype(candle_core::DType::I64)?; - let n_gen = self.decode_state.n_generated.to_vec1::()?; + let n_gen = self.decode_state.n_generated.to_vec1::()?; let mut results = Vec::new(); for b in 0..self.config.padded_batch_size { let n = n_gen[b] as usize; diff --git a/arc-cuda-graph/src/buffers.rs b/arc-cuda-graph/src/buffers.rs index 6b8d00f0f..df3ec623a 100644 --- a/arc-cuda-graph/src/buffers.rs +++ b/arc-cuda-graph/src/buffers.rs @@ -1,8 +1,45 @@ //! Pre-allocated GPU buffers for the decode loop. +//! +//! **Every buffer here is read and written by CUDA kernels whose signatures +//! are `int32_t*`** (`cuda/decode_loop.cu`, `cuda/sampling.cu`). The dtype of +//! the Candle tensor therefore has to be 4 bytes wide, or the host and the +//! device disagree about where element `i` lives. #[cfg(feature = "cuda")] use candle_core::{DType, Device, Tensor}; +#[cfg(feature = "cuda")] +use candle_core::cuda::cudarc::driver::sys::CUstream; + +#[cfg(feature = "cuda")] +extern "C" { + fn cudaMemsetAsync( + dst: *mut std::ffi::c_void, + value: i32, + count: usize, + stream: CUstream, + ) -> u32; + fn cudaMemcpyAsync( + dst: *mut std::ffi::c_void, + src: *const std::ffi::c_void, + count: usize, + kind: u32, + stream: CUstream, + ) -> u32; +} + +/// Allocate a zeroed I32 tensor. `Tensor::zeros` on I32 is not supported on +/// every Candle backend (see the same note in `sampling_cuda.rs:296`), so go +/// through `from_vec`. +#[cfg(feature = "cuda")] +fn zeros_i32( + elems: usize, + shape: impl Into, + device: &Device, +) -> candle_core::Result { + Tensor::from_vec(vec![0i32; elems], shape, device) +} + /// Pre-allocated decode input buffers at fixed GPU addresses. #[cfg(feature = "cuda")] pub struct DecodeInputBuffers { @@ -31,33 +68,74 @@ impl DecodeInputBuffers { } /// GPU-side decode state that persists across WHILE loop iterations. +/// +/// 🔴 These were all `I64` while every kernel that touches them declares +/// `int32_t*` (`decode_loop.cu:19-37`, `sampling.cu:21`). The kernel wrote +/// element `b` at byte `4*b`; Candle read element `b` at byte `8*b`. Nothing +/// lined up: `n_generated`, `finished`, `loop_condition` and the whole +/// `output_tokens` matrix were read back as garbage, and `output_tokens` was +/// additionally indexed with the wrong row stride. The path is gated off +/// before it runs, so nothing ever reported it. #[cfg(feature = "cuda")] pub struct DecodeState { - pub sampled_tokens: Tensor, // [padded_bs] i64 - pub n_generated: Tensor, // [padded_bs] i64 - pub output_tokens: Tensor, // [padded_bs, max_tokens] i64 - pub finished: Tensor, // [padded_bs] i64 - pub loop_condition: Tensor, // [1] i64 + pub sampled_tokens: Tensor, // [padded_bs] i32 + pub n_generated: Tensor, // [padded_bs] i32 + pub output_tokens: Tensor, // [padded_bs, max_tokens] i32 + pub finished: Tensor, // [padded_bs] i32 + pub loop_condition: Tensor, // [1] i32 pub max_tokens: usize, } #[cfg(feature = "cuda")] impl DecodeState { pub fn new(padded_bs: usize, max_tokens: usize, device: &Device) -> candle_core::Result { + let _ = DType::I32; // keep the dtype named in this module Ok(Self { - sampled_tokens: Tensor::zeros(padded_bs, DType::I64, device)?, - n_generated: Tensor::zeros(padded_bs, DType::I64, device)?, - output_tokens: Tensor::zeros((padded_bs, max_tokens), DType::I64, device)?, - finished: Tensor::zeros(padded_bs, DType::I64, device)?, - loop_condition: Tensor::ones(1, DType::I64, device)?, + sampled_tokens: zeros_i32(padded_bs, padded_bs, device)?, + n_generated: zeros_i32(padded_bs, padded_bs, device)?, + output_tokens: zeros_i32(padded_bs * max_tokens, (padded_bs, max_tokens), device)?, + finished: zeros_i32(padded_bs, padded_bs, device)?, + loop_condition: Tensor::from_vec(vec![1i32], 1, device)?, max_tokens, }) } - pub fn reset(&mut self, device: &Device, padded_bs: usize) -> candle_core::Result<()> { - self.n_generated = Tensor::zeros(padded_bs, DType::I64, device)?; - self.finished = Tensor::zeros(padded_bs, DType::I64, device)?; - self.loop_condition = Tensor::ones(1, DType::I64, device)?; + /// Reset the per-generation state **in place**, on `stream`. + /// + /// 🔴 This used to re-`Tensor::zeros` each field, which allocates NEW + /// device buffers at NEW addresses. The captured CUDA graph has the OLD + /// addresses baked into its kernel nodes, so resetting between + /// generations pointed the graph at freed memory — the reset silently + /// un-did the capture. A pointer-stable buffer is a hard requirement of + /// anything that gets captured, so zero the existing allocations instead. + pub fn reset(&mut self, stream: CUstream) -> candle_core::Result<()> { + let n_gen = crate::weights::tensor_device_ptr(&self.n_generated)?; + let fin = crate::weights::tensor_device_ptr(&self.finished)?; + let cond = crate::weights::tensor_device_ptr(&self.loop_condition)?; + let bs = self.n_generated.elem_count(); + unsafe { + let s = cudaMemsetAsync(n_gen as *mut _, 0, bs * 4, stream); + if s != 0 { + candle_core::bail!("cudaMemsetAsync(n_generated) failed: {s}"); + } + let s = cudaMemsetAsync(fin as *mut _, 0, bs * 4, stream); + if s != 0 { + candle_core::bail!("cudaMemsetAsync(finished) failed: {s}"); + } + // loop_condition starts at 1 ("keep going"). memset writes a byte + // pattern, so 0x01010101 would be wrong; write the word directly. + let one: i32 = 1; + let s = cudaMemcpyAsync( + cond as *mut _, + &one as *const i32 as *const _, + 4, + 1, // cudaMemcpyHostToDevice + stream, + ); + if s != 0 { + candle_core::bail!("cudaMemcpyAsync(loop_condition) failed: {s}"); + } + } Ok(()) } } diff --git a/arc-cuda-graph/src/cuda/decode_loop.cu b/arc-cuda-graph/src/cuda/decode_loop.cu index b18195029..4a4894d3a 100644 --- a/arc-cuda-graph/src/cuda/decode_loop.cu +++ b/arc-cuda-graph/src/cuda/decode_loop.cu @@ -152,3 +152,115 @@ extern "C" void launch_check_all_done_conditional( extern "C" int arc_has_graph_conditional() { return ARC_HAS_GRAPH_CONDITIONAL; } + +// ============================================================================ +// ArcInfer/ArcGraph — device decode loop step commit +// +// The narrow sibling of `decode_step_update` above, for the CUDA-graph replay +// path (`arc-cuda-graph/src/device_loop.rs`). That kernel is written for the +// paged-attention autonomous runner and unconditionally dereferences +// `finished`, `n_generated`, `slot_mappings` and `block_tables`; the replay +// path has none of those (V4 runs with `cache_config == None`), so passing +// nulls would fault. This kernel touches only what that path actually owns. +// +// It closes the decode loop on-device: +// sampled token -> the PINNED U32 buffer the captured graph reads +// -> device position += 1 +// -> pinned+mapped host ring, for the host to drain later +// +// Ordering is load-bearing. Every range check runs BEFORE any write, and the +// ring write head is published LAST, after `__threadfence_system()`. So a token +// the host can see in the ring is guaranteed to have been committed to the +// graph input buffer, and a refused commit publishes nothing at all rather than +// letting the next launch consume a garbage id. +// ============================================================================ + +#define ARC_COMMIT_FAULT_NONE 0 +#define ARC_COMMIT_FAULT_TOKEN_RANGE 1 +#define ARC_COMMIT_FAULT_POSITION_LIMIT 2 + +__global__ void arc_graph_step_commit( + const int32_t* __restrict__ sampled, // [batch] i32, written by CudaSampler + uint32_t* __restrict__ input_ids, // [batch] U32, PINNED graph input + uint32_t* __restrict__ positions, // [batch] U32, PINNED graph input; nullable + int32_t* __restrict__ ring, // [batch, ring_size] pinned + mapped + int32_t* __restrict__ ring_head, // [batch] pinned + mapped + int32_t* __restrict__ fault, // [1] pinned + mapped, sticky + int32_t ring_size, + int32_t vocab, + int32_t position_limit, // exclusive; <= 0 disables the check + int32_t batch_size +) { + int bid = blockIdx.x; + if (threadIdx.x != 0) return; + if (bid >= batch_size) return; + + // A fault latched by an earlier step means the graph input buffer already + // holds a stale id. Committing on top of that would paper over it. + if (fault != nullptr && ((volatile int32_t*)fault)[0] != ARC_COMMIT_FAULT_NONE) return; + + int token = sampled[bid]; + + // ---- checks first, writes after ------------------------------------- + if (token < 0 || token >= vocab) { + if (fault != nullptr) { + ((volatile int32_t*)fault)[0] = ARC_COMMIT_FAULT_TOKEN_RANGE; + __threadfence_system(); + } + return; + } + if (positions != nullptr && position_limit > 0) { + // `positions[bid]` is the slot this step just wrote KV to. The next + // step writes at +1, and the fixed-capacity window is `position_limit` + // wide, so +1 must stay strictly inside it. + if ((int64_t)positions[bid] + 1 >= (int64_t)position_limit) { + if (fault != nullptr) { + ((volatile int32_t*)fault)[0] = ARC_COMMIT_FAULT_POSITION_LIMIT; + __threadfence_system(); + } + return; + } + } + + // ---- commit --------------------------------------------------------- + // This single store is what removes the host from the loop: the captured + // graph baked THIS address at capture time, so the next `cuGraphLaunch` + // reads the token without any host round trip. + input_ids[bid] = (uint32_t)token; + if (positions != nullptr) { + positions[bid] += 1u; + } + + // ---- publish last ---------------------------------------------------- + if (ring != nullptr && ring_head != nullptr && ring_size > 0) { + int wp = ring_head[bid]; + int slot = wp % ring_size; + if (slot < 0) slot += ring_size; + ((volatile int32_t*)ring)[(int64_t)bid * ring_size + slot] = token; + // Order the token store ahead of the head store as seen by the HOST. + // Without this the host can observe an advanced head and read the + // previous occupant of the slot. + __threadfence_system(); + ((volatile int32_t*)ring_head)[bid] = wp + 1; + } +} + +extern "C" void arc_launch_graph_step_commit( + const int32_t* sampled, + uint32_t* input_ids, + uint32_t* positions, + int32_t* ring, + int32_t* ring_head, + int32_t* fault, + int32_t ring_size, + int32_t vocab, + int32_t position_limit, + int32_t batch_size, + cudaStream_t stream +) { + arc_graph_step_commit<<>>( + sampled, input_ids, positions, + ring, ring_head, fault, + ring_size, vocab, position_limit, batch_size + ); +} diff --git a/arc-cuda-graph/src/cuda/sampling_kernel.cu b/arc-cuda-graph/src/cuda/sampling_kernel.cu index f9d5d9f02..53caaf192 100644 --- a/arc-cuda-graph/src/cuda/sampling_kernel.cu +++ b/arc-cuda-graph/src/cuda/sampling_kernel.cu @@ -127,6 +127,12 @@ __device__ __forceinline__ float block_sum(float val, float* s_vals) { return s_vals[0]; } +// Branch engagement counters for the hybrid sampler. A test that assumes +// "support width 512 must have taken the fallback" is assuming the very thing +// it is trying to establish; these make it observable. +__device__ unsigned int arc_hybrid_branch_a = 0; +__device__ unsigned int arc_hybrid_branch_b = 0; + // Splitmix64 step matching sampling_cpu.rs lines 156-161. // state = state * 0x9E3779B97F4A7C15 + 0xDEADBEEFC0DECAFE // mixed = (state ^ (state >> 30)) * 0xBF58476D1CE4E5B9 @@ -353,6 +359,390 @@ __global__ void arc_sample_kernel( } } + +// ============================================================================= +// Nucleus sampling in a FIXED number of passes (threshold bisection). +// +// `arc_sample_kernel` above pulls the keep-list one argmax at a time, so it +// costs O(vocab x kept): a full vocabulary scan PER KEPT TOKEN. With +// `top_k <= 0` -- which is what mistral.rs's default maps to -- `kept` is +// bounded only by `top_p`, and on a diffuse distribution it reaches thousands. +// MEASURED on an H200 at vocab=129280 with a flat 12928-token support: +// 663,116 us to sample ONE token, roughly 10x an entire V4 decode step. +// +// This kernel selects the same nucleus by binary-searching a probability +// THRESHOLD rather than enumerating the set. For non-negative IEEE-754 floats +// the bit pattern is monotone when compared as u32, so bisecting the integer +// key converges to one ULP in a fixed 32 passes regardless of how many tokens +// the nucleus contains. Cost becomes O(vocab) with a constant factor. +// +// It does NOT narrow the distribution. The kept set is still +// {i : p_i >= t} for the largest t whose mass still covers `top_p` -- that is +// the definition of the nucleus. Where several tokens share the boundary bit +// pattern it keeps ALL of them, so it errs toward keeping MORE than the exact +// nucleus, never fewer. `top_k` is applied as a second, independent bisection +// on the kept COUNT, and the two thresholds combine with max(), so a caller +// setting both still gets the intersection. +// ============================================================================= +template +__global__ void arc_sample_bisect_kernel( + const T* __restrict__ logits, + const uint32_t* __restrict__ freq_counts, + uint64_t* __restrict__ rng_state, + int32_t* __restrict__ token_ids, + float* __restrict__ probs_scratch, + int vocab, + SamplingParams cfg +) { + const int bid = blockIdx.x; + const int tid = threadIdx.x; + const T* __restrict__ row = logits + (int64_t)bid * vocab; + float* __restrict__ probs = probs_scratch + (int64_t)bid * vocab; + const uint32_t* fcnt = freq_counts ? (freq_counts + (int64_t)bid * vocab) : nullptr; + + extern __shared__ char smem_bisect[]; + float* s_vals = reinterpret_cast(smem_bisect); + int* s_idxs = reinterpret_cast(s_vals + BLOCK); + + const bool temp_active = (cfg.temperature > 0.0f && cfg.temperature != 1.0f); + const float inv_temp = temp_active ? (1.0f / cfg.temperature) : 1.0f; + + // --- Phase 1: penalties + temperature; keep the argmax for the fallback. + float local_max = -FLT_MAX; + int local_i = 0; + for (int i = tid; i < vocab; i += BLOCK) { + float l = to_float(row[i]); + if (fcnt) { + uint32_t c = fcnt[i]; + if (c != 0) { + l -= cfg.frequency_penalty * static_cast(c); + l -= cfg.presence_penalty; + } + } + if (temp_active) l *= inv_temp; + probs[i] = l; + if (l > local_max || (l == local_max && i < local_i)) { local_max = l; local_i = i; } + } + float gmax; int gmax_idx; + block_argmax_idx(local_max, local_i, s_vals, s_idxs, gmax, gmax_idx); + __syncthreads(); + + // --- Phase 2/3: exp-shift, sum, normalize to probabilities. + float local_sum = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + float p = __expf(probs[i] - gmax); + probs[i] = p; + local_sum += p; + } + const float gsum = block_sum(local_sum, s_vals); + __syncthreads(); + const float inv_sum = (gsum > 0.0f) ? (1.0f / gsum) : 1.0f; + float local_pmax = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + float p = probs[i] * inv_sum; + probs[i] = p; + if (p > local_pmax) local_pmax = p; + } + const float pmax = block_max(local_pmax, s_vals); + __syncthreads(); + + // Degenerate distribution: fall back to the mode, as the CPU reference does. + if (!(pmax > 0.0f)) { + if (tid == 0) token_ids[bid] = gmax_idx; + return; + } + + float target = cfg.top_p; + if (target > 1.0f) target = 1.0f; + if (target < 0.0f) target = 0.0f; + + // --- Phase 4a: largest threshold key whose kept mass still covers top_p. + uint32_t lo = 0u, hi = __float_as_uint(pmax); + while (lo < hi) { + const uint32_t mid = lo + ((hi - lo + 1u) >> 1); + const float threshold = __uint_as_float(mid); + float m = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + const float p = probs[i]; + if (p >= threshold) m += p; + } + const float gm = block_sum(m, s_vals); + __syncthreads(); + if (gm >= target) lo = mid; else hi = mid - 1u; + } + uint32_t key = lo; + + // --- Phase 4b: top_k, as an independent bisection on the kept COUNT. + if (cfg.top_k > 0 && cfg.top_k < vocab) { + uint32_t klo = 0u, khi = __float_as_uint(pmax); + while (klo < khi) { + const uint32_t mid = klo + ((khi - klo) >> 1); + const float threshold = __uint_as_float(mid); + float c = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + if (probs[i] >= threshold) c += 1.0f; + } + const float gc = block_sum(c, s_vals); + __syncthreads(); + if (gc <= static_cast(cfg.top_k)) khi = mid; else klo = mid + 1u; + } + if (klo > key) key = klo; + } + + const float threshold = __uint_as_float(key); + + // --- Phase 5: CDF walk, parallel across the block. + // Each thread sums the kept mass in its own strided subset; thread 0 + // exclusive-scans those 256 partials so every thread knows the CDF offset + // its subset begins at. Only the thread whose interval contains `u` then + // walks, and it walks vocab/BLOCK items rather than the whole vocabulary. + float part = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { + const float p = probs[i]; + if (p >= threshold) part += p; + } + s_vals[tid] = part; + __syncthreads(); + + __shared__ float s_u; + if (tid == 0) { + float acc = 0.0f; + for (int t = 0; t < BLOCK; ++t) { const float v = s_vals[t]; s_vals[t] = acc; acc += v; } + uint64_t st = rng_state[bid]; + const float u = splitmix_uniform(st); + rng_state[bid] = st; + s_u = u * acc; + // Deterministic fallback if float error leaves no owning interval. + token_ids[bid] = gmax_idx; + } + __syncthreads(); + + const float base = s_vals[tid]; + if (s_u >= base && s_u < base + part) { + float acc = base; + for (int i = tid; i < vocab; i += BLOCK) { + const float p = probs[i]; + if (p >= threshold) { + acc += p; + if (acc >= s_u) { token_ids[bid] = i; break; } + } + } + } +} + + +// ============================================================================= +// Hybrid nucleus sampler: exact enumeration while the nucleus is small, +// threshold bisection once it isn't. One kernel, no host involvement. +// +// MEASURED on an H200 at vocab=129280 (support width -> legacy us / bisect us): +// 1 -> 134 / 1994 | 8 -> 503 / 1996 | 64 -> 3132 / 1991 +// 512 -> 24287 / 1997 | 4096 -> 205446 / 1998 | 12928 -> 665694 / 1998 +// Bisection is FLAT; enumeration is linear in the kept count. Neither wins +// everywhere: enumeration is 3-15x faster on the peaked distributions real +// models produce, bisection is 333x faster on the diffuse tail. Shipping +// either alone trades one regression for another. +// +// THE SWITCH POINT IS NOT A TUNING CONSTANT. Bisection resolves one bit of the +// float32 probability key per vocabulary pass, so it costs 32 passes (plus 32 +// more when top_k forces a second search). Enumeration costs one vocabulary +// pass per kept token. Enumeration is therefore cheaper exactly while +// kept < (the number of passes bisection would spend) +// so the budget IS that pass count: ARC_ENUM_BUDGET = 32, derived from the +// width of the key being searched, the same way the 32 passes are. Spending +// the budget and then falling back costs at most budget + bisection, so the +// hybrid is never worse than ~2x the better branch, and the measured crossover +// (~64) sits within a factor of two of the derived budget -- the derivation +// and the measurement agree. +// +// Why the fallback is INSIDE the kernel: a captured CUDA graph bakes its +// kernel arguments and cannot branch on device state, so "launch enumeration, +// read a flag, maybe launch bisection" would need a device->host round trip +// per step -- exactly the host dependency the autonomous decode loop exists to +// remove. The branch is therefore taken block-uniformly on device. +// +// No cap and no narrowing: the fallback is exact, not a truncation. Exceeding +// the budget changes which ALGORITHM selects the nucleus, never which tokens +// the nucleus contains. +// ============================================================================= +template +__global__ void arc_sample_hybrid_kernel( + const T* __restrict__ logits, + const uint32_t* __restrict__ freq_counts, + uint64_t* __restrict__ rng_state, + int32_t* __restrict__ token_ids, + float* __restrict__ probs_scratch, + int32_t* __restrict__ keep_idx_scratch, + float* __restrict__ keep_p_scratch, + int vocab, + SamplingParams cfg +) { + constexpr int ARC_ENUM_BUDGET = 32; // == bisection's pass count + + const int bid = blockIdx.x; + const int tid = threadIdx.x; + const T* __restrict__ row = logits + (int64_t)bid * vocab; + float* __restrict__ probs = probs_scratch + (int64_t)bid * vocab; + int32_t* __restrict__ keep_idx = keep_idx_scratch + (int64_t)bid * vocab; + float* __restrict__ keep_p = keep_p_scratch + (int64_t)bid * vocab; + const uint32_t* fcnt = freq_counts ? (freq_counts + (int64_t)bid * vocab) : nullptr; + + extern __shared__ char smem_hy[]; + float* s_vals = reinterpret_cast(smem_hy); + int* s_idxs = reinterpret_cast(s_vals + BLOCK); + + const bool temp_active = (cfg.temperature > 0.0f && cfg.temperature != 1.0f); + const float inv_temp = temp_active ? (1.0f / cfg.temperature) : 1.0f; + + float target = cfg.top_p; + if (target > 1.0f) target = 1.0f; + if (target < 0.0f) target = 0.0f; + +#define ARC_BUILD_PROBS(GMAX_IDX_OUT) \ + do { \ + float _lmax = -FLT_MAX; int _li = 0; \ + for (int i = tid; i < vocab; i += BLOCK) { \ + float l = to_float(row[i]); \ + if (fcnt) { uint32_t c = fcnt[i]; \ + if (c != 0) { l -= cfg.frequency_penalty * (float)c; \ + l -= cfg.presence_penalty; } } \ + if (temp_active) l *= inv_temp; \ + probs[i] = l; \ + if (l > _lmax || (l == _lmax && i < _li)) { _lmax = l; _li = i; } \ + } \ + float _gv; int _gi; \ + block_argmax_idx(_lmax, _li, s_vals, s_idxs, _gv, _gi); \ + __syncthreads(); \ + (GMAX_IDX_OUT) = _gi; \ + float _lsum = 0.0f; \ + for (int i = tid; i < vocab; i += BLOCK) { \ + float p = __expf(probs[i] - _gv); probs[i] = p; _lsum += p; } \ + float _gsum = block_sum(_lsum, s_vals); \ + __syncthreads(); \ + float _inv = (_gsum > 0.0f) ? (1.0f / _gsum) : 1.0f; \ + for (int i = tid; i < vocab; i += BLOCK) probs[i] *= _inv; \ + __syncthreads(); \ + } while (0) + + int gmax_idx = 0; + ARC_BUILD_PROBS(gmax_idx); + + // ---- Branch A: bounded exact enumeration (the legacy algorithm). + __shared__ int s_kept; + __shared__ float s_cum; + __shared__ int s_done; // 1 = nucleus fully determined here + if (tid == 0) { s_kept = 0; s_cum = 0.0f; s_done = 0; } + __syncthreads(); + + while (true) { + if (s_cum >= target) { if (tid == 0) s_done = 1; __syncthreads(); break; } + if (cfg.top_k > 0 && s_kept >= cfg.top_k) { if (tid == 0) s_done = 1; __syncthreads(); break; } + if (s_kept >= ARC_ENUM_BUDGET) { __syncthreads(); break; } // spend budget, fall back + + float lv = -FLT_MAX; int li = 0; + for (int i = tid; i < vocab; i += BLOCK) { + float v = probs[i]; + if (v > lv || (v == lv && i < li)) { lv = v; li = i; } + } + float gv; int gi; + block_argmax_idx(lv, li, s_vals, s_idxs, gv, gi); + __syncthreads(); + if (gv <= 0.0f) { if (tid == 0) s_done = 1; __syncthreads(); break; } + if (tid == 0) { + keep_idx[s_kept] = gi; + keep_p[s_kept] = gv; + s_kept += 1; + s_cum += gv; + probs[gi] = -FLT_MAX; // tombstone + } + __syncthreads(); + } + + if (s_done) { + if (tid == 0) atomicAdd(&arc_hybrid_branch_a, 1u); + if (s_kept <= 0) { if (tid == 0) token_ids[bid] = gmax_idx; return; } + if (tid == 0) { + uint64_t st = rng_state[bid]; + float u = splitmix_uniform(st); + rng_state[bid] = st; + float t = u * s_cum; + float acc = 0.0f; + int sel = keep_idx[s_kept - 1]; + for (int k = 0; k < s_kept; ++k) { + acc += keep_p[k]; + if (acc >= t) { sel = keep_idx[k]; break; } + } + token_ids[bid] = sel; + } + return; + } + + // ---- Branch B: budget spent, nucleus is diffuse. probs[] was tombstoned + // by the enumeration above, so rebuild it, then bisect. + if (tid == 0) atomicAdd(&arc_hybrid_branch_b, 1u); + ARC_BUILD_PROBS(gmax_idx); + + float lpmax = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { float p = probs[i]; if (p > lpmax) lpmax = p; } + const float pmax = block_max(lpmax, s_vals); + __syncthreads(); + if (!(pmax > 0.0f)) { if (tid == 0) token_ids[bid] = gmax_idx; return; } + + uint32_t lo = 0u, hi = __float_as_uint(pmax); + while (lo < hi) { + const uint32_t mid = lo + ((hi - lo + 1u) >> 1); + const float threshold_m = __uint_as_float(mid); + float m = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { float p = probs[i]; if (p >= threshold_m) m += p; } + const float gm = block_sum(m, s_vals); + __syncthreads(); + if (gm >= target) lo = mid; else hi = mid - 1u; + } + uint32_t key = lo; + + if (cfg.top_k > 0 && cfg.top_k < vocab) { + uint32_t klo = 0u, khi = __float_as_uint(pmax); + while (klo < khi) { + const uint32_t mid = klo + ((khi - klo) >> 1); + const float threshold_m = __uint_as_float(mid); + float c = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { if (probs[i] >= threshold_m) c += 1.0f; } + const float gc = block_sum(c, s_vals); + __syncthreads(); + if (gc <= (float)cfg.top_k) khi = mid; else klo = mid + 1u; + } + if (klo > key) key = klo; + } + + const float threshold = __uint_as_float(key); + float part = 0.0f; + for (int i = tid; i < vocab; i += BLOCK) { float p = probs[i]; if (p >= threshold) part += p; } + s_vals[tid] = part; + __syncthreads(); + + __shared__ float s_u; + if (tid == 0) { + float acc = 0.0f; + for (int t = 0; t < BLOCK; ++t) { const float v = s_vals[t]; s_vals[t] = acc; acc += v; } + uint64_t st = rng_state[bid]; + const float u = splitmix_uniform(st); + rng_state[bid] = st; + s_u = u * acc; + token_ids[bid] = gmax_idx; + } + __syncthreads(); + + const float base = s_vals[tid]; + if (s_u >= base && s_u < base + part) { + float acc = base; + for (int i = tid; i < vocab; i += BLOCK) { + const float p = probs[i]; + if (p >= threshold) { acc += p; if (acc >= s_u) { token_ids[bid] = i; break; } } + } + } +#undef ARC_BUILD_PROBS +} + } // namespace arc_sampler // ============================================================================= @@ -499,4 +889,70 @@ void arc_launch_sampler_f16( } } + +// Fixed-pass nucleus sampler. Same signature as `arc_launch_sampler_*` so it +// is a drop-in; `keep_idx_scratch` / `keep_p_scratch` are unused by it. +#define ARC_DEFINE_BISECT_LAUNCHER(SUFFIX, CTYPE) \ +void arc_launch_sampler_##SUFFIX##_bisect( \ + const void* logits, const uint32_t* freq_counts, uint64_t* rng_state, \ + int32_t* token_ids, float* probs_scratch, int32_t* keep_idx_scratch, \ + float* keep_p_scratch, int vocab, int batch, \ + arc_sampler::SamplingParams cfg, cudaStream_t stream) { \ + (void)keep_idx_scratch; (void)keep_p_scratch; \ + const int B = 256; \ + size_t smem = B * (sizeof(float) + sizeof(int)); \ + if (cfg.greedy) { \ + arc_sampler::arc_greedy_kernel<<>>( \ + reinterpret_cast(logits), freq_counts, token_ids, vocab, \ + cfg); \ + return; \ + } \ + arc_sampler::arc_sample_bisect_kernel<<>>( \ + reinterpret_cast(logits), freq_counts, rng_state, \ + token_ids, probs_scratch, vocab, cfg); \ +} + +ARC_DEFINE_BISECT_LAUNCHER(f32, float) +ARC_DEFINE_BISECT_LAUNCHER(bf16, __nv_bfloat16) +ARC_DEFINE_BISECT_LAUNCHER(f16, __half) + + +void arc_hybrid_branch_counts(unsigned int* a_out, unsigned int* b_out) { + cudaMemcpyFromSymbol(a_out, arc_sampler::arc_hybrid_branch_a, + sizeof(unsigned int), 0, cudaMemcpyDeviceToHost); + cudaMemcpyFromSymbol(b_out, arc_sampler::arc_hybrid_branch_b, + sizeof(unsigned int), 0, cudaMemcpyDeviceToHost); +} + +void arc_hybrid_branch_reset(void) { + unsigned int z = 0; + cudaMemcpyToSymbol(arc_sampler::arc_hybrid_branch_a, &z, sizeof(z), 0, + cudaMemcpyHostToDevice); + cudaMemcpyToSymbol(arc_sampler::arc_hybrid_branch_b, &z, sizeof(z), 0, + cudaMemcpyHostToDevice); +} + +#define ARC_DEFINE_HYBRID_LAUNCHER(SUFFIX, CTYPE) \ +void arc_launch_sampler_##SUFFIX##_hybrid( \ + const void* logits, const uint32_t* freq_counts, uint64_t* rng_state, \ + int32_t* token_ids, float* probs_scratch, int32_t* keep_idx_scratch, \ + float* keep_p_scratch, int vocab, int batch, \ + arc_sampler::SamplingParams cfg, cudaStream_t stream) { \ + const int B = 256; \ + size_t smem = B * (sizeof(float) + sizeof(int)); \ + if (cfg.greedy) { \ + arc_sampler::arc_greedy_kernel<<>>( \ + reinterpret_cast(logits), freq_counts, token_ids, vocab, \ + cfg); \ + return; \ + } \ + arc_sampler::arc_sample_hybrid_kernel<<>>( \ + reinterpret_cast(logits), freq_counts, rng_state, \ + token_ids, probs_scratch, keep_idx_scratch, keep_p_scratch, vocab, cfg); \ +} + +ARC_DEFINE_HYBRID_LAUNCHER(f32, float) +ARC_DEFINE_HYBRID_LAUNCHER(bf16, __nv_bfloat16) +ARC_DEFINE_HYBRID_LAUNCHER(f16, __half) + } // extern "C" diff --git a/arc-cuda-graph/src/dedicated.rs b/arc-cuda-graph/src/dedicated.rs index 555ea4fb5..b8675e067 100644 --- a/arc-cuda-graph/src/dedicated.rs +++ b/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); diff --git a/arc-cuda-graph/src/device_loop.rs b/arc-cuda-graph/src/device_loop.rs new file mode 100644 index 000000000..757c31a7f --- /dev/null +++ b/arc-cuda-graph/src/device_loop.rs @@ -0,0 +1,2027 @@ +//! Parent system: ArcInfer / ArcGraph +//! +//! Device-resident decode loop: **N decode steps cost N `cuGraphLaunch` calls +//! and zero `cudaStreamSynchronize`.** +//! +//! # What was wrong with plain replay +//! +//! [`crate::graph::CudaGraphRunner::replay`] launches the captured graph and +//! then calls `cudaStreamSynchronize` (`graph.rs:362`) before handing the +//! output tensor back. The caller samples the next token on the **host**, which +//! copies logits (or an argmax scalar) down over PCIe, and only then launches +//! the next graph. So a decode step is +//! +//! ```text +//! GPU runs step -> host blocks on sync -> D2H -> host samples -> launch +//! ``` +//! +//! Capture removed the cost of *issuing* the work — 3,961 launch APIs per token +//! became 3.8 — and left that serialized round trip completely intact. Deleting +//! the sync on its own buys nothing, because the host-side sample forces a +//! device-to-host copy that synchronizes anyway. **The sync and the host sample +//! have to go together.** +//! +//! # The loop this module closes +//! +//! Per step, all on one stream, no host in the middle: +//! +//! 1. `cuGraphLaunch` — the captured forward. +//! 2. [`crate::sampling_cuda::CudaSampler`] — argmax/top-p **on device**, token +//! id written to an `i32[batch]` device buffer. Never copied to the host. +//! 3. `arc_graph_step_commit` — scatters that token into the *pinned* graph +//! input buffer the captured graph reads (so the next launch consumes it), +//! bumps the device position, and publishes the token into a pinned+mapped +//! host ring. +//! +//! Nothing in that sequence blocks. The host enqueues all `burst` steps and the +//! GPU works through them while the host is already ahead. +//! +//! # How the host still gets its tokens +//! +//! The ring lives in `cudaHostAlloc(..., cudaHostAllocMapped)` memory. The GPU +//! pushes each token straight into it, `__threadfence_system()`, then publishes +//! a monotonic write head. The host reads both with **plain volatile loads** — +//! no CUDA API call, no stream synchronize, no D2H copy. Draining is a host +//! memory poll, which is categorically different from a stream sync: with a +//! sync the host cannot enqueue step N+1 until step N retires; here every step +//! is already enqueued and the host is only watching them land. +//! +//! # Why the aliasing logits buffer is safe here +//! +//! The tensor `replay()` hands back **aliases graph-owned storage that the next +//! `cuGraphLaunch` overwrites**. Holding it across a launch, or draining it +//! lazily, reads data the next step already clobbered — silently, producing +//! plausible tokens. That alias/clone class has already caused three +//! verification bugs in this repo. +//! +//! This design is safe against it for one specific reason, and it is worth +//! stating rather than assuming: **every access is enqueued on the same +//! stream, in order, and none of them is a host read.** Per step the stream +//! carries `launch N -> sample -> commit -> launch N+1`. CUDA stream ordering +//! guarantees the sampler kernel starts only after launch N has written the +//! logits and finishes before launch N+1 begins overwriting them. The host +//! never reads the logits at all, so there is no window in which it could read +//! a clobbered buffer. +//! +//! This is strictly stronger than the host path, which has to *synchronize* to +//! read those logits safely. What the host does drain — the token ring — is a +//! **separate, pinned, host-owned allocation** that no graph launch touches, so +//! draining it lazily is sound however far behind the host falls. The one thing +//! that would break the argument is issuing any of these on a different stream; +//! that is why `CudaDeviceOps` captures the graph's own stream and uses it for +//! the sampler and the commit kernel too. +//! +//! # Detecting a fault without synchronizing +//! +//! A fault *during* graph execution is asynchronous — `cuGraphLaunch` returns +//! SUCCESS and only a later API call surfaces it. `replay()` used a full sync +//! for exactly this reason and the comment there says so. Dropping the check +//! outright would turn an illegal access into a silently poisoned context and +//! an output buffer full of garbage, a bug class this codebase has shipped more +//! than once. So the check is *replaced*, not removed, by three signals that +//! together need no blocking call: +//! +//! * **`cudaStreamQuery`** — non-blocking. NVIDIA's CUDA Runtime API reference +//! for it states: *"Note that this function may also return error codes from +//! previous, asynchronous launches."* So a sticky async fault surfaces on a +//! call that never blocks. +//! * **Idle-but-short** — `cudaStreamQuery` returning `cudaSuccess` means every +//! enqueued step has retired. If all of them retired and fewer tokens than +//! steps reached the ring, the work definitively did not do what it was asked +//! to. That is a fault with no timeout heuristic in it. +//! * **The device fault word** — the commit kernel range-checks the sampled +//! token and the position before it writes anything, and latches a code into +//! a pinned word. It travels back with the ring, so reading it is free. +//! +//! `max_spins` is a last-resort liveness backstop for a kernel that hangs +//! without ever erroring. It is not the primary detector. +//! +//! # KNOWN LIMITATION — the burst holds the scheduler +//! +//! A burst enqueues `burst` decode steps back-to-back and does not return to +//! the engine until they have all landed. **For that whole window the engine +//! admits no new request**, so a request arriving mid-burst waits up to +//! `burst` decode steps (default 4) before it is even considered. +//! +//! At batch 1 — the only shape [`admit`] accepts — that is harmless: there is +//! one sequence, nothing to interleave with, and the alternative is a host +//! round trip per token. It is **not** harmless above batch 1, where a burst +//! would straddle sequences with different stop conditions and different +//! admission urgency, and where the wasted tail after one sequence's EOS is +//! multiplied by the batch. That is why `admit` refuses `batch_size != 1` +//! outright rather than treating it as a tuning question. +//! +//! Anyone widening this past batch 1 owns that problem first. Raising `burst` +//! trades admission latency for host round trips one-for-one, and the +//! trade is only free while there is exactly one sequence in flight. +//! +//! # Layering +//! +//! Everything above the CUDA calls is expressed against the [`DeviceOps`] +//! trait, and that trait deliberately has **no `synchronize` method** — the +//! driver cannot block the stream because it is not given a way to. The +//! resulting [`BurstDriver`] is plain Rust, compiles without the `cuda` +//! feature, and is unit-tested against a scripted fake GPU. The CUDA-gated +//! [`CudaDeviceOps`] is pointer plumbing and nothing else. + +use std::fmt; + +// =========================================================================== +// Configuration +// =========================================================================== + +/// Opt-in switch. Absent or `0` means this whole module is inert and decode +/// takes the existing `replay()` + host-sample path. +pub const ENV_ENABLE: &str = "ARC_GRAPH_DEVICE_LOOP"; +/// Steps enqueued per burst before the host drains. Default [`DEFAULT_BURST`]. +pub const ENV_BURST: &str = "ARC_GRAPH_DEVICE_LOOP_BURST"; + +/// Steps launched back-to-back before the host drains the ring. +/// +/// The whole burst is enqueued without any host round trip, so this is also +/// how far the GPU may run ahead of the host's `Sequence` bookkeeping — and +/// therefore the worst-case detection latency, in steps, for a stop string, an +/// EOS, or an async fault. 4 keeps the wasted-work tail small while still +/// amortising the drain over four launches. +pub const DEFAULT_BURST: usize = 4; + +/// Tokens the pinned ring holds per row. Must exceed `burst`, or a burst could +/// overwrite tokens the host has not read yet. +pub const DEFAULT_RING_SIZE: usize = 1024; + +/// Host spins between `cudaStreamQuery` calls while draining. Pure host +/// memory polls are ~free; the driver query costs a syscall-ish amount, so it +/// is not worth doing on every iteration. +pub const DEFAULT_POLL_EVERY: u32 = 64; + +/// Emit [`DeviceDecodeLoop::status_line`] every this many bursts. The counter +/// it carries (`launches`) is the only proof the loop actually ran. +pub const STATUS_EVERY_BURSTS: u64 = 64; + +/// Liveness backstop: give up after this many spins even if the stream never +/// reports an error. Only trips on a kernel that hangs without faulting. +pub const DEFAULT_MAX_SPINS: u64 = 200_000_000; + +/// Is the device decode loop switched on? +/// +/// Default **off**. `ARC_GRAPH_DEVICE_LOOP=1` opts in; `0`, `false` and `off` +/// are explicit refusals so an operator can pin it off in an environment that +/// sets it upstream. +pub fn device_loop_enabled() -> bool { + match std::env::var(ENV_ENABLE) { + Ok(v) => { + let v = v.trim().to_ascii_lowercase(); + !(v.is_empty() || v == "0" || v == "false" || v == "off" || v == "no") + } + Err(_) => false, + } +} + +/// Burst length from [`ENV_BURST`], clamped to at least 1. +pub fn burst_from_env() -> usize { + std::env::var(ENV_BURST) + .ok() + .and_then(|v| v.trim().parse::().ok()) + .filter(|v| *v > 0) + .unwrap_or(DEFAULT_BURST) +} + +/// Everything the driver needs that does not change from burst to burst. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct DeviceLoopConfig { + /// Steps per burst. Also the run-ahead bound and the detection latency. + pub burst: usize, + /// Tokens per row in the pinned ring. + pub ring_size: usize, + /// Host spins between stream queries while draining. + pub poll_every: u32, + /// Spin budget before declaring the stream stalled. + pub max_spins: u64, + /// Stop the burst's token stream here. `< 0` disables EOS truncation. + pub eos_token_id: i32, + /// Vocabulary size; the commit kernel range-checks against it. + pub vocab: usize, + /// Exclusive upper bound on the device position. The fixed-capacity graph + /// decode arm reads a constant `capacity`-wide window and writes KV at + /// `position`; advancing past it is an out-of-bounds write, so both the + /// driver and the kernel refuse to. + pub position_limit: u32, +} + +impl DeviceLoopConfig { + /// Config with the shipped defaults for everything but the model facts. + pub fn new(vocab: usize, eos_token_id: i32, position_limit: u32) -> Self { + Self { + burst: burst_from_env(), + ring_size: DEFAULT_RING_SIZE, + poll_every: DEFAULT_POLL_EVERY, + max_spins: DEFAULT_MAX_SPINS, + eos_token_id, + vocab, + position_limit, + } + } + + /// Reject a config that cannot be run safely, rather than discovering it + /// as corruption at step `ring_size + 1`. + pub fn validate(&self) -> Result<(), DeviceLoopError> { + if self.burst == 0 { + return Err(DeviceLoopError::BadConfig("burst must be non-zero")); + } + if self.vocab == 0 { + return Err(DeviceLoopError::BadConfig("vocab must be non-zero")); + } + if self.poll_every == 0 { + // The drain does `spins % poll_every`; zero would divide by zero. + return Err(DeviceLoopError::BadConfig("poll_every must be non-zero")); + } + if self.ring_size <= self.burst { + // A burst longer than the ring overwrites its own unread tokens. + return Err(DeviceLoopError::BadConfig( + "ring_size must exceed burst, or a burst overwrites unread tokens", + )); + } + Ok(()) + } + + /// Steps this burst may launch given how far the sequence has already got. + /// + /// Clamped by the position limit so the burst can never walk the device + /// position past the fixed-capacity window. Returns 0 when the sequence has + /// no room left, which the caller must treat as "fall back to eager", not + /// as "launch zero steps and carry on". + pub fn plan_steps(&self, current_position: u32, max_new_tokens: usize) -> usize { + let room = self.position_limit.saturating_sub(current_position) as usize; + self.burst.min(room).min(max_new_tokens) + } +} + +// =========================================================================== +// Errors +// =========================================================================== + +/// Codes the commit kernel latches into the pinned fault word. Kept in sync +/// with `arc_graph_step_commit` in `src/cuda/decode_loop.cu`. +pub const FAULT_NONE: i32 = 0; +/// Sampled token id was outside `[0, vocab)`. +pub const FAULT_TOKEN_RANGE: i32 = 1; +/// Advancing the position would leave the fixed-capacity window. +pub const FAULT_POSITION_LIMIT: i32 = 2; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DeviceLoopError { + BadConfig(&'static str), + /// A CUDA call reported a failure. `at` names the call site. + Cuda { + at: &'static str, + code: u32, + }, + /// The stream retired every enqueued step but the ring is short. The work + /// completed without producing the tokens it was launched to produce. + LostTokens { + expected: i32, + got: i32, + }, + /// The commit kernel refused to write and latched a code. + DeviceFault(i32), + /// The write head ran more than `ring_size` ahead of the read cursor, so + /// unread tokens were overwritten. Reported rather than silently returning + /// whatever survived. + RingOverrun { + pending: i64, + ring_size: usize, + }, + /// Spin budget exhausted with the stream still busy and no error reported. + Stalled { + spins: u64, + }, +} + +impl fmt::Display for DeviceLoopError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::BadConfig(m) => write!(f, "device decode loop misconfigured: {m}"), + Self::Cuda { at, code } => { + write!(f, "device decode loop: {at} failed: CUDA code {code}") + } + Self::LostTokens { expected, got } => write!( + f, + "device decode loop: the stream retired every enqueued step but only {got} of \ + {expected} tokens reached the ring. The graph ran and did not commit its tokens \ + — treat this as an async fault, not as slow progress" + ), + Self::DeviceFault(code) => { + let what = match *code { + FAULT_TOKEN_RANGE => "sampled token id outside [0, vocab)", + FAULT_POSITION_LIMIT => { + "advancing the position would leave the fixed-capacity KV window" + } + _ => "unrecognised code", + }; + write!( + f, + "device decode loop: commit kernel refused to write (code {code}: {what}). No \ + token was scattered into the graph input buffer, so the next launch would \ + have read a stale id" + ) + } + Self::RingOverrun { pending, ring_size } => write!( + f, + "device decode loop: {pending} tokens pending in a {ring_size}-slot ring — the GPU \ + overwrote tokens the host had not read. Tokens were lost; this is a correctness \ + stop, not a throughput warning" + ), + Self::Stalled { spins } => write!( + f, + "device decode loop: stream still busy after {spins} spins with no error reported. \ + The graph is hung rather than faulted" + ), + } + } +} + +impl std::error::Error for DeviceLoopError {} + +impl From for candle_core::Error { + fn from(e: DeviceLoopError) -> Self { + candle_core::Error::Msg(e.to_string()) + } +} + +// =========================================================================== +// The device interface +// =========================================================================== + +/// Result of a non-blocking stream probe. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamState { + /// Every enqueued step has retired (`cudaSuccess`). + Idle, + /// Work still outstanding (`cudaErrorNotReady`). + Busy, +} + +/// The device operations a burst needs. +/// +/// **There is deliberately no `synchronize` method.** The driver is structurally +/// incapable of blocking the stream, which is the property this whole module +/// exists to guarantee — it cannot be regressed by an edit that "just adds a +/// sync to be safe", because there is nothing to call. +pub trait DeviceOps { + /// One `cuGraphLaunch` of the captured decode graph. Must not block. + fn launch_graph(&mut self) -> Result<(), DeviceLoopError>; + + /// Sample on device: graph logits -> the `i32[batch]` token buffer. Must + /// not block and must not copy the token to the host. + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError>; + + /// Commit on device: token -> pinned graph input ids, position += 1, and + /// publish into the pinned ring. Must not block. + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError>; + + /// Non-blocking stream probe. `Err` for a sticky async fault. + fn poll_stream(&mut self) -> Result; + + /// Monotonic write head for `row`, read from pinned mapped host memory. + /// A plain volatile load — no CUDA API call. + fn ring_head(&self, row: usize) -> i32; + + /// Token at absolute index `index` for `row` (the impl applies the modulo). + fn ring_slot(&self, row: usize, index: i32) -> i32; + + /// The sticky fault word, read from pinned mapped host memory. + fn device_fault(&self) -> i32; + + /// Rows the ring carries. + fn batch(&self) -> usize; +} + +// =========================================================================== +// The driver +// =========================================================================== + +/// What one burst produced. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct BurstOutcome { + /// Tokens per row, in generation order, truncated at and including EOS. + pub tokens: Vec>, + /// True when some row emitted EOS inside this burst. + pub hit_eos: bool, + /// Steps actually launched. When `hit_eos` is true this exceeds the number + /// of tokens returned: the burst was already enqueued and the GPU computed + /// the tail before the host could know. That work is discarded, and this + /// field is how much of it there was. + pub steps_launched: usize, +} + +/// Drives bursts against a [`DeviceOps`], owning the ring read cursors. +pub struct BurstDriver { + cfg: DeviceLoopConfig, + /// Absolute index of the next unread ring slot, per row. + cursor: Vec, + launches: u64, + bursts: u64, + spins: u64, +} + +impl BurstDriver { + pub fn new(cfg: DeviceLoopConfig, batch: usize) -> Result { + cfg.validate()?; + Ok(Self { + cfg, + cursor: vec![0; batch.max(1)], + launches: 0, + bursts: 0, + spins: 0, + }) + } + + pub fn config(&self) -> &DeviceLoopConfig { + &self.cfg + } + + /// Total `cuGraphLaunch` calls issued. With `bursts()` this is the honest + /// answer to "did the device loop do anything": `launches == steps` and + /// zero stream synchronizes is the entire claim. + pub fn launches(&self) -> u64 { + self.launches + } + + pub fn bursts(&self) -> u64 { + self.bursts + } + + /// Host spins spent waiting for tokens to land. Rises when the host + /// out-runs the GPU, which is the healthy direction. + pub fn spins(&self) -> u64 { + self.spins + } + + /// The read cursor is per-sequence state. Reset it when the ring is reset + /// (a new request reuses the buffers). + pub fn reset_cursors(&mut self) { + for c in self.cursor.iter_mut() { + *c = 0; + } + } + + /// Enqueue `steps` decode steps and drain their tokens. + /// + /// Issues exactly `steps` `cuGraphLaunch` calls and **zero** stream + /// synchronizations. The drain is a host-memory poll on pinned mapped + /// memory plus a periodic non-blocking `cudaStreamQuery`. + pub fn run_burst( + &mut self, + ops: &mut D, + steps: usize, + ) -> Result { + if steps == 0 { + return Err(DeviceLoopError::BadConfig( + "run_burst called with zero steps; the caller must fall back to eager instead", + )); + } + if steps >= self.cfg.ring_size { + return Err(DeviceLoopError::BadConfig( + "burst is not smaller than the ring; it would overwrite its own tokens", + )); + } + + // ── enqueue ────────────────────────────────────────────────────── + // No blocking call in this loop, by construction: `DeviceOps` has no + // way to synchronize. `poll_stream` is the async-fault probe and it + // returns immediately whatever the stream is doing. + for _ in 0..steps { + ops.launch_graph()?; + ops.sample_on_device()?; + ops.commit_on_device()?; + self.launches += 1; + // Surface a sticky fault from an EARLIER step while later steps are + // already in flight. Costs nothing and shortens detection latency + // from "end of burst" to "next step". + ops.poll_stream()?; + } + self.bursts += 1; + + // ── drain ──────────────────────────────────────────────────────── + let batch = ops.batch().min(self.cursor.len()); + let want: Vec = self.cursor[..batch] + .iter() + .map(|c| c.saturating_add(steps as i32)) + .collect(); + for (row, &target) in want.iter().enumerate() { + self.await_row(ops, row, target)?; + } + + // The commit kernel latches its refusal instead of writing, so this + // must be checked even when every token arrived: a fault on step 3 of 4 + // means steps 4.. read a stale input id. + let fault = ops.device_fault(); + if fault != FAULT_NONE { + return Err(DeviceLoopError::DeviceFault(fault)); + } + + // ── collect ────────────────────────────────────────────────────── + let mut tokens = Vec::with_capacity(batch); + let mut hit_eos = false; + for (row, &target) in want.iter().enumerate() { + let mut row_tokens = Vec::with_capacity(steps); + let mut idx = self.cursor[row]; + while idx < target { + let tok = ops.ring_slot(row, idx); + idx += 1; + let is_eos = self.cfg.eos_token_id >= 0 && tok == self.cfg.eos_token_id; + row_tokens.push(tok); + if is_eos { + hit_eos = true; + break; + } + } + // The cursor advances over the WHOLE burst even when EOS truncated + // the returned tokens. Those slots were written; leaving the cursor + // behind would re-read them as if they were the next burst's. + self.cursor[row] = target; + tokens.push(row_tokens); + } + + Ok(BurstOutcome { + tokens, + hit_eos, + steps_launched: steps, + }) + } + + /// Wait for `row`'s write head to reach `target`, without synchronizing. + fn await_row( + &mut self, + ops: &mut D, + row: usize, + target: i32, + ) -> Result<(), DeviceLoopError> { + let mut spins: u64 = 0; + loop { + let head = ops.ring_head(row); + if head >= target { + // Overrun check: more unread tokens than the ring can hold + // means the GPU lapped the host and older tokens are gone. + let pending = (head as i64) - (self.cursor[row] as i64); + if pending > self.cfg.ring_size as i64 { + return Err(DeviceLoopError::RingOverrun { + pending, + ring_size: self.cfg.ring_size, + }); + } + return Ok(()); + } + + if spins.is_multiple_of(self.cfg.poll_every as u64) { + match ops.poll_stream()? { + StreamState::Busy => {} + StreamState::Idle => { + // Every enqueued step retired. Re-read once: the head + // store and the stream's retirement are not ordered + // against each other from the host's point of view. + let head = ops.ring_head(row); + if head >= target { + return Ok(()); + } + return Err(DeviceLoopError::LostTokens { + expected: target, + got: head, + }); + } + } + } + + spins += 1; + self.spins += 1; + if spins >= self.cfg.max_spins { + return Err(DeviceLoopError::Stalled { spins }); + } + std::hint::spin_loop(); + } + } +} + +// =========================================================================== +// CUDA implementation +// =========================================================================== + +#[cfg(feature = "cuda")] +mod cuda_impl { + use super::*; + use crate::ffi::*; + use crate::sampling_cpu::SamplingConfig; + use crate::sampling_cuda::CudaSampler; + use candle_core::cuda::cudarc::driver::sys::CUstream; + use candle_core::{DType, Device, Tensor}; + + /// Device address of a buffer the captured graph reads. + /// + /// Refuses a non-contiguous tensor rather than silently taking the address + /// of the copy `contiguous()` would make — that copy is not the address the + /// graph baked, so writing to it would look fine and change nothing. + fn pinned_ptr(t: &Tensor, what: &str) -> candle_core::Result { + if !t.is_contiguous() { + candle_core::bail!( + "device decode loop: the {what} buffer is not contiguous, so its address is not \ + the one the captured graph baked. Refusing rather than writing into a copy." + ); + } + crate::weights::tensor_device_ptr(t) + } + + /// `cudaErrorNotReady` — the stream still has outstanding work. Not a + /// failure, and the only non-success code `cudaStreamQuery` returns that + /// is not a fault. + const CUDA_ERROR_NOT_READY: u32 = 600; + + /// Pinned, mapped host memory the GPU writes and the host reads with plain + /// volatile loads. Freed with `cudaFreeHost` on drop. + struct PinnedRing { + tokens: *mut i32, + head: *mut i32, + fault: *mut i32, + batch: usize, + ring_size: usize, + } + + impl PinnedRing { + fn new(batch: usize, ring_size: usize) -> Result { + let mut tokens: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut head: *mut std::ffi::c_void = std::ptr::null_mut(); + let mut fault: *mut std::ffi::c_void = std::ptr::null_mut(); + unsafe { + let s = cudaHostAlloc(&mut tokens, batch * ring_size * 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(ring tokens)", + code: s, + }); + } + let s = cudaHostAlloc(&mut head, batch * 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + cudaFreeHost(tokens); + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(ring head)", + code: s, + }); + } + let s = cudaHostAlloc(&mut fault, 4, CUDA_HOST_ALLOC_MAPPED); + if s != CUDA_SUCCESS { + cudaFreeHost(tokens); + cudaFreeHost(head); + return Err(DeviceLoopError::Cuda { + at: "cudaHostAlloc(fault word)", + code: s, + }); + } + std::ptr::write_bytes(tokens as *mut u8, 0, batch * ring_size * 4); + std::ptr::write_bytes(head as *mut u8, 0, batch * 4); + std::ptr::write_bytes(fault as *mut u8, 0, 4); + } + Ok(Self { + tokens: tokens as *mut i32, + head: head as *mut i32, + fault: fault as *mut i32, + batch, + ring_size, + }) + } + + fn reset(&self) { + unsafe { + std::ptr::write_bytes(self.head as *mut u8, 0, self.batch * 4); + std::ptr::write_bytes(self.fault as *mut u8, 0, 4); + } + } + } + + impl Drop for PinnedRing { + fn drop(&mut self) { + unsafe { + if !self.tokens.is_null() { + cudaFreeHost(self.tokens as *mut _); + } + if !self.head.is_null() { + cudaFreeHost(self.head as *mut _); + } + if !self.fault.is_null() { + cudaFreeHost(self.fault as *mut _); + } + } + } + } + + /// The real [`DeviceOps`]: sampler + commit kernel on the graph's stream. + /// + /// Holds `&mut` to nothing outside itself except through raw pointers taken + /// once at construction, because every buffer it touches must be at a + /// **stable device address** — the captured graph baked those addresses and + /// a reallocation silently invalidates the graph. + pub struct CudaDeviceOps { + stream: CUstream, + sampler: CudaSampler, + sampling_cfg: SamplingConfig, + /// The graph's output logits tensor, reshaped to `[batch, vocab]`. + logits: Tensor, + /// `i32[batch]`, written by the sampler, read by the commit kernel. + token_ids: Tensor, + /// The PINNED `U32` token-id buffer the captured graph reads. + input_ids_ptr: *mut u32, + /// The PINNED `U32` position buffer, or null when positions are owned + /// elsewhere (e.g. by a device-side RoPE position path). + positions_ptr: *mut u32, + ring: PinnedRing, + cfg: DeviceLoopConfig, + /// `cuGraphLaunch` closure state: the exec handle to replay. + exec: CUgraphExec, + } + + impl CudaDeviceOps { + /// * `logits` — the captured graph's output tensor. Its storage address + /// is graph-owned and stable; the sampler reads it in place. + /// * `input_ids_buf` / `positions_buf` — the *pinned* graph input + /// buffers (`layers::GRAPH_MODE_INPUT_IDS` / `GRAPH_MODE_POSITIONS`). + /// These are the addresses the captured graph reads, which is why a + /// device-side write into them closes the loop with no host copy. + #[allow(clippy::too_many_arguments)] + pub fn new( + device: &Device, + exec: CUgraphExec, + stream: CUstream, + sampler: CudaSampler, + sampling_cfg: SamplingConfig, + logits: &Tensor, + input_ids_buf: &Tensor, + positions_buf: Option<&Tensor>, + cfg: DeviceLoopConfig, + ) -> candle_core::Result { + cfg.validate()?; + let batch = sampler.batch(); + + // `[1, 1, vocab]` (or `[1, vocab]`) -> `[batch, vocab]`. Reshaping a + // contiguous tensor is metadata-only, so this does not move the + // graph-owned storage the sampler is about to read. + let logits = logits.reshape((batch, sampler.vocab()))?; + if logits.dtype() != sampler.dtype() { + candle_core::bail!( + "device decode loop: graph logits dtype {:?} != sampler dtype {:?}", + logits.dtype(), + sampler.dtype() + ); + } + + if input_ids_buf.dtype() != DType::U32 { + candle_core::bail!( + "device decode loop: the graph input-ids buffer must be U32 (got {:?}); the \ + commit kernel writes it as uint32_t", + input_ids_buf.dtype() + ); + } + if input_ids_buf.elem_count() < batch { + candle_core::bail!( + "device decode loop: input-ids buffer holds {} elements, need {batch}", + input_ids_buf.elem_count() + ); + } + if let Some(p) = positions_buf { + if p.dtype() != DType::U32 { + candle_core::bail!( + "device decode loop: the graph positions buffer must be U32 (got {:?})", + p.dtype() + ); + } + if p.elem_count() < batch { + candle_core::bail!( + "device decode loop: positions buffer holds {} elements, need {batch}", + p.elem_count() + ); + } + } + + // I32 via `from_vec`: `Tensor::zeros` on I32 is not universally + // supported across backends, the same reason `CudaSampler::new` + // builds `keep_idx_scratch` this way. + let token_ids = Tensor::from_vec(vec![0i32; batch], batch, device)?; + + let input_ids_ptr = pinned_ptr(input_ids_buf, "graph input-ids")? as *mut u32; + let positions_ptr = match positions_buf { + Some(p) => pinned_ptr(p, "graph positions")? as *mut u32, + None => std::ptr::null_mut(), + }; + let ring = PinnedRing::new(batch, cfg.ring_size)?; + + Ok(Self { + stream, + sampler, + sampling_cfg, + logits, + token_ids, + input_ids_ptr, + positions_ptr, + ring, + cfg, + exec, + }) + } + + /// Zero the ring head and fault word. Call when a new sequence starts + /// reusing these buffers; pair it with [`BurstDriver::reset_cursors`]. + pub fn reset_ring(&self) { + self.ring.reset(); + } + } + + impl DeviceOps for CudaDeviceOps { + fn launch_graph(&mut self) -> Result<(), DeviceLoopError> { + // The launch and nothing else. `graph.rs::replay` follows this with + // `cudaStreamSynchronize`; the entire point here is that we do not. + let s = unsafe { cuGraphLaunch(self.exec, self.stream) }; + if s != CUDA_SUCCESS { + return Err(DeviceLoopError::Cuda { + at: "cuGraphLaunch", + code: s, + }); + } + Ok(()) + } + + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.sampler + .sample(&self.logits, None, self.sampling_cfg, &self.token_ids) + .map_err(|e| { + tracing::error!("device decode loop: on-device sample failed: {e}"); + DeviceLoopError::Cuda { + at: "CudaSampler::sample", + code: u32::MAX, + } + }) + } + + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError> { + let token_ptr = match crate::weights::tensor_device_ptr(&self.token_ids) { + Ok(p) => p as *const i32, + Err(e) => { + tracing::error!("device decode loop: token buffer pointer unavailable: {e}"); + return Err(DeviceLoopError::Cuda { + at: "tensor_device_ptr(token_ids)", + code: u32::MAX, + }); + } + }; + unsafe { + arc_launch_graph_step_commit( + token_ptr, + self.input_ids_ptr, + self.positions_ptr, + self.ring.tokens, + self.ring.head, + self.ring.fault, + self.cfg.ring_size as i32, + self.cfg.vocab as i32, + self.cfg.position_limit as i32, + self.ring.batch as i32, + self.stream, + ); + } + Ok(()) + } + + fn poll_stream(&mut self) -> Result { + // Non-blocking. NVIDIA's reference for cudaStreamQuery: "Note that + // this function may also return error codes from previous, + // asynchronous launches." That is what makes a sticky fault from an + // earlier graph step visible here without ever blocking. + let s = unsafe { cudaStreamQuery(self.stream) }; + match s { + CUDA_SUCCESS => Ok(StreamState::Idle), + CUDA_ERROR_NOT_READY => Ok(StreamState::Busy), + code => Err(DeviceLoopError::Cuda { + at: "cudaStreamQuery (async fault from an earlier graph step)", + code, + }), + } + } + + fn ring_head(&self, row: usize) -> i32 { + unsafe { std::ptr::read_volatile(self.ring.head.add(row)) } + } + + fn ring_slot(&self, row: usize, index: i32) -> i32 { + let slot = + row * self.ring.ring_size + (index.rem_euclid(self.ring.ring_size as i32) as usize); + unsafe { std::ptr::read_volatile(self.ring.tokens.add(slot)) } + } + + fn device_fault(&self) -> i32 { + unsafe { std::ptr::read_volatile(self.ring.fault) } + } + + fn batch(&self) -> usize { + self.ring.batch + } + } + + /// The whole device loop behind one handle: sampler, commit kernel, pinned + /// ring and burst driver. + /// + /// Built once per captured graph and kept alive across decode steps, + /// because every buffer it holds a pointer to must stay at the address the + /// graph baked. Dropping and rebuilding it mid-generation would hand the + /// commit kernel a fresh `token_ids` allocation while the graph still reads + /// the old input-ids address — so callers keep it, they do not recreate it. + pub struct DeviceDecodeLoop { + ops: CudaDeviceOps, + driver: BurstDriver, + /// 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 + // CUDA handles and pinned host allocations, both valid from any thread, and + // the pipeline that owns this is held behind a mutex. Without these the + // `NormalPipeline` field would make the whole pipeline non-`Send`. + unsafe impl Send for CudaDeviceOps {} + unsafe impl Sync for CudaDeviceOps {} + unsafe impl Send for DeviceDecodeLoop {} + unsafe impl Sync for DeviceDecodeLoop {} + + impl DeviceDecodeLoop { + #[allow(clippy::too_many_arguments)] + pub fn new( + device: &Device, + exec: CUgraphExec, + stream: CUstream, + logits: &Tensor, + input_ids_buf: &Tensor, + positions_buf: Option<&Tensor>, + cfg: DeviceLoopConfig, + ) -> candle_core::Result { + let batch = 1usize; + let vocab = cfg.vocab; + // Greedy, always. [`admit`] refuses every non-greedy sequence, so + // taking the mode as a parameter would only create a way for a + // caller to contradict the guard. The seed is irrelevant under + // greedy (pure argmax draws no random numbers) and is fixed so the + // sampler's RNG buffer is deterministic if the mode ever widens. + let sampling_cfg = SamplingConfig::greedy(); + let sampler = CudaSampler::new(device, batch, vocab, logits.dtype(), 0)?; + let ops = CudaDeviceOps::new( + device, + exec, + stream, + sampler, + sampling_cfg, + logits, + input_ids_buf, + positions_buf, + cfg, + )?; + let driver = BurstDriver::new(cfg, batch)?; + Ok(Self { + ops, + driver, + 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 + // running counters periodically so `launches` is observable in a + // live run, not just in a debugger. + if self.driver.bursts().is_multiple_of(STATUS_EVERY_BURSTS) { + tracing::info!("{}", self.status_line()); + } + Ok(out) + } + + /// Start a fresh sequence on the same buffers. + pub fn reset(&mut self) { + self.ops.reset_ring(); + self.driver.reset_cursors(); + clear_pending_tokens(); + } + + /// `cuGraphLaunch` calls issued so far — the numerator of the claim. + pub fn launches(&self) -> u64 { + self.driver.launches() + } + + pub fn bursts(&self) -> u64 { + self.driver.bursts() + } + + /// The one line that separates a working device loop from an inert one + /// (D18). `launches=0` is a subsystem that did nothing and says so. + pub fn status_line(&self) -> String { + format!( + "ARCGRAPH DEVICE LOOP: launches={} bursts={} burst_len={} host_syncs=0 \ + host_spins={} pending={}", + self.driver.launches(), + self.driver.bursts(), + self.driver.config().burst, + self.driver.spins(), + pending_token_count(), + ) + } + } +} + +#[cfg(feature = "cuda")] +pub use cuda_impl::{CudaDeviceOps, DeviceDecodeLoop}; + +// =========================================================================== +// Admission +// =========================================================================== + +/// Why the device loop refused a request. +/// +/// Refusing is the normal, safe outcome — the caller falls back to the existing +/// `replay()` + host-sample path. Each variant names something the device +/// sampler cannot reproduce *exactly*, because a token that differs from the +/// host path is a silent behaviour change, not an optimisation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Refusal { + /// `ARC_GRAPH_DEVICE_LOOP` is not set. The default. + NotEnabled, + /// Only batch 1 is wired: the burst holds the scheduler for its whole + /// length, so it must not straddle sequences with different stop conditions. + NotBatchOne, + /// Prefill, not decode. + NotDecodeStep, + /// No captured graph to replay, or its output is not yet trusted. + GraphNotReady, + /// Non-greedy sampling. The device sampler runs Splitmix64 per row while + /// the host runs Isaac64, so for any stochastic config the two draw + /// *different, equally valid* tokens — which silently breaks seeded + /// reproducibility. Greedy is deterministic, so device and host agree + /// exactly, which is also what makes the path verifiable. + NotGreedy, + /// The caller wants logprobs, which the device sampler does not return. + WantsLogprobs, + /// A logits processor, grammar or penalty must see the logits on the host. + HostSideLogitsWork, + /// No room left in the fixed-capacity KV window for even one step. + NoPositionRoom, +} + +impl fmt::Display for Refusal { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let s = match self { + Self::NotEnabled => "ARC_GRAPH_DEVICE_LOOP is not set", + Self::NotBatchOne => "batch size is not 1", + Self::NotDecodeStep => "this is a prefill step, not decode", + Self::GraphNotReady => "no captured graph, or its output is not yet trusted", + Self::NotGreedy => { + "sampling is not greedy; the device and host RNGs differ, so tokens would \ + diverge from the seeded host path" + } + Self::WantsLogprobs => { + "the request asks for logprobs, which the device sampler \ + does not produce" + } + Self::HostSideLogitsWork => { + "a logits processor, grammar or repetition penalty needs the logits on the host" + } + Self::NoPositionRoom => "no room left in the fixed-capacity KV window", + }; + write!(f, "{s}") + } +} + +/// Everything the admission check needs, gathered by the caller. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct AdmissionFacts { + pub batch_size: usize, + pub seq_len: usize, + pub graph_ready: bool, + pub greedy: bool, + pub return_logprobs: bool, + pub host_side_logits_work: bool, + pub current_position: u32, +} + +/// May the device loop drive this step? +/// +/// Kept separate from the CUDA plumbing, and not feature-gated, so the decision +/// is unit-tested on any host. Every refusal is a fall-back to the existing +/// path, never an error. +pub fn admit( + facts: &AdmissionFacts, + cfg: &DeviceLoopConfig, + enabled: bool, +) -> Result { + if !enabled { + return Err(Refusal::NotEnabled); + } + if facts.seq_len != 1 { + return Err(Refusal::NotDecodeStep); + } + if facts.batch_size != 1 { + return Err(Refusal::NotBatchOne); + } + if !facts.graph_ready { + return Err(Refusal::GraphNotReady); + } + if facts.return_logprobs { + return Err(Refusal::WantsLogprobs); + } + if facts.host_side_logits_work { + return Err(Refusal::HostSideLogitsWork); + } + if !facts.greedy { + return Err(Refusal::NotGreedy); + } + let steps = cfg.plan_steps(facts.current_position, cfg.burst); + if steps == 0 { + return Err(Refusal::NoPositionRoom); + } + Ok(steps) +} + +// =========================================================================== +// Host-side token handoff +// =========================================================================== + +// Tokens a burst produced that the engine has not consumed yet. +// +// A burst computes `steps` tokens in one go, but the engine's `Sequence` +// bookkeeping — stop strings, streaming, EOS — is per token. So the burst +// parks its tokens here and the engine drains them one at a time on the +// following steps, during which no GPU work is launched at all. +// +// Thread-local because at batch 1 `sample_and_add_toks` runs the sampler +// inline on the engine thread (`use_async_pool = seqs_len > 1` is false), the +// same thread that ran the forward. The `GRAPH_MODE_*` graph inputs in +// `mistralrs-core/src/layers.rs` are scoped the same way for the same reason. +std::thread_local! { + static PENDING_TOKENS: std::cell::RefCell> = + const { std::cell::RefCell::new(std::collections::VecDeque::new()) }; + /// 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. 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 { + if t >= 0 { + q.push_back(t as u32); + } else { + tracing::error!( + "device decode loop: dropping negative token id {t} from the ring; the \ + commit kernel range-checks before writing, so this means the ring was \ + misread rather than that the model produced it" + ); + } + } + }) +} + +/// 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. +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. +// +// `Pipeline::forward_inputs` decides whether to run a burst, but it is handed +// only `inputs` and `return_raw_logits` — it never sees the `Sequence`, so it +// cannot ask whether sampling is greedy or whether logprobs were requested. +// `sample_sequence` does see it. So eligibility is published there and read by +// the *next* forward. +// +// The one-step lag is safe in the only direction that matters: a sequence's +// sampling parameters are fixed for its lifetime, and the flag starts `false`, +// so the first decode step always falls back and the loop can only ever engage +// after a step has confirmed the parameters. It can never engage for a +// sequence it has not seen. +std::thread_local! { + static DEVICE_LOOP_ELIGIBLE: std::cell::Cell = const { std::cell::Cell::new(false) }; +} + +/// Publish whether the sequence being sampled may be driven by the device +/// loop, 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)) +} + +/// Did the last sampled sequence qualify? `false` until one has. +pub fn device_loop_eligible() -> bool { + DEVICE_LOOP_ELIGIBLE.with(|c| c.get()) +} + +/// Drop parked tokens. +/// +/// MUST be called when a sequence finishes, is preempted, or the path falls +/// back — a token left here would be handed to the *next* sequence as if it +/// were its own. That is the "+1-step divergence" failure mode with a different +/// name, so it gets an explicit call rather than a lifetime. +/// +/// 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_OWNER.with(|o| o.set(None)); + PENDING_TOKENS.with(|q| q.borrow_mut().clear()) +} + +/// 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(); + 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. +/// +/// Without this a construction or burst failure would be retried on every +/// decode step forever — and because a refusal falls back to the *eager* +/// forward rather than to `replay()`, that would quietly make every step slower +/// than doing nothing at all. One failure, one warning, off for good. +static DEVICE_LOOP_KILLED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); + +/// Latch the device loop off for the rest of the process, loudly and once. +pub fn kill_device_loop(reason: &str) { + if !DEVICE_LOOP_KILLED.swap(true, std::sync::atomic::Ordering::SeqCst) { + tracing::error!( + "ArcGraph device decode loop DISABLED for the rest of this process: {reason}. Decode \ + falls back to graph replay with a host sync and a host sample. This is a stop, not a \ + retry — a loop that fails once and is retried every step is slower than never having \ + engaged." + ); + } + stand_down(); +} + +/// Has the device loop been latched off? +pub fn device_loop_killed() -> bool { + DEVICE_LOOP_KILLED.load(std::sync::atomic::Ordering::SeqCst) +} + +#[cfg(feature = "cuda")] +extern "C" { + /// `arc_graph_step_commit` — `src/cuda/decode_loop.cu`. + /// + /// Scatters the device-sampled token into the pinned graph input buffer, + /// advances the device position, and publishes the token into the pinned + /// mapped ring. Range-checks before it writes anything, latching a code + /// into `fault` instead of committing a bad token. + pub fn arc_launch_graph_step_commit( + sampled: *const i32, + input_ids: *mut u32, + positions: *mut u32, + ring: *mut i32, + ring_head: *mut i32, + fault: *mut i32, + ring_size: i32, + vocab: i32, + position_limit: i32, + batch: i32, + stream: candle_core::cuda::cudarc::driver::sys::CUstream, + ); +} + +// =========================================================================== +// Tests +// =========================================================================== + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::VecDeque; + + /// A scripted fake GPU. + /// + /// It models the property that matters: work is enqueued without blocking + /// and lands *later*. `launch_graph` pushes the next scripted token into an + /// in-flight queue; `poll_stream` retires one in-flight step per call and + /// publishes its token into the ring. So a driver that expects tokens to be + /// present the instant it launches will not find them, and a driver that + /// never polls will spin forever — both real failure modes. + struct FakeGpu { + script: VecDeque, + in_flight: VecDeque, + ring: Vec, + head: i32, + fault: i32, + ring_size: usize, + launches: usize, + samples: usize, + commits: usize, + polls: usize, + /// Set to drop the commit instead of publishing — the "no-op + /// implementation" this test suite must not accept. + commit_is_noop: bool, + /// Injected async fault, surfaced from `poll_stream`. + fault_after: Option, + /// Report Idle immediately even with work outstanding. + lie_idle: bool, + /// A hung kernel: work is enqueued, the stream reports Busy forever and + /// nothing ever retires. No error is ever raised, so only the liveness + /// backstop can end it. + never_retires: bool, + } + + impl FakeGpu { + fn new(script: Vec, ring_size: usize) -> Self { + Self { + script: script.into(), + in_flight: VecDeque::new(), + ring: vec![0; ring_size], + head: 0, + fault: FAULT_NONE, + ring_size, + launches: 0, + samples: 0, + commits: 0, + polls: 0, + commit_is_noop: false, + fault_after: None, + lie_idle: false, + never_retires: false, + } + } + + fn retire_one(&mut self) { + if let Some(tok) = self.in_flight.pop_front() { + let slot = (self.head as usize) % self.ring_size; + self.ring[slot] = tok; + self.head += 1; + } + } + } + + impl DeviceOps for FakeGpu { + fn launch_graph(&mut self) -> Result<(), DeviceLoopError> { + self.launches += 1; + Ok(()) + } + fn sample_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.samples += 1; + Ok(()) + } + fn commit_on_device(&mut self) -> Result<(), DeviceLoopError> { + self.commits += 1; + if self.commit_is_noop { + return Ok(()); + } + let tok = self.script.pop_front().unwrap_or(0); + self.in_flight.push_back(tok); + Ok(()) + } + fn poll_stream(&mut self) -> Result { + self.polls += 1; + if let Some(after) = self.fault_after { + if self.polls > after { + return Err(DeviceLoopError::Cuda { + at: "cudaStreamQuery (async fault from an earlier graph step)", + code: 700, + }); + } + } + if self.never_retires { + return Ok(StreamState::Busy); + } + self.retire_one(); + if self.lie_idle || self.in_flight.is_empty() { + Ok(StreamState::Idle) + } else { + Ok(StreamState::Busy) + } + } + fn ring_head(&self, _row: usize) -> i32 { + self.head + } + fn ring_slot(&self, _row: usize, index: i32) -> i32 { + self.ring[index.rem_euclid(self.ring_size as i32) as usize] + } + fn device_fault(&self) -> i32 { + self.fault + } + fn batch(&self) -> usize { + 1 + } + } + + fn cfg(burst: usize, ring: usize) -> DeviceLoopConfig { + DeviceLoopConfig { + burst, + ring_size: ring, + poll_every: 1, + max_spins: 10_000, + eos_token_id: 2, + vocab: 128, + position_limit: 4096, + } + } + + /// THE claim: N decode steps cost exactly N `cuGraphLaunch` calls, and the + /// tokens that come back are the ones the GPU produced. + /// + /// Fails on a no-op: with fewer launches the count is wrong, and with a + /// commit that never publishes the drain reports `LostTokens`. + #[test] + fn n_steps_cost_exactly_n_launches_and_return_the_gpu_tokens() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let out = d.run_burst(&mut gpu, 4).unwrap(); + + assert_eq!(gpu.launches, 4, "N steps must cost exactly N cuGraphLaunch"); + assert_eq!(gpu.samples, 4, "every step must sample on device"); + assert_eq!(gpu.commits, 4, "every step must commit on device"); + assert_eq!(out.tokens, vec![vec![11, 12, 13, 14]]); + assert_eq!(d.launches(), 4); + assert!(!out.hit_eos); + } + + /// The control the previous test needs: an implementation that launches and + /// samples but never publishes the token MUST fail, and must fail as + /// `LostTokens` rather than by spinning out. This is the "test that cannot + /// fail" guard — if the driver accepted a silent no-op commit, the whole + /// module would be decoration. + #[test] + fn a_commit_that_publishes_nothing_is_rejected_as_lost_tokens() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.commit_is_noop = true; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + assert!( + matches!( + err, + DeviceLoopError::LostTokens { + expected: 4, + got: 0 + } + ), + "expected LostTokens, got {err:?}" + ); + // It still launched: the failure is in the commit, and the error must + // not be confused with "the driver never ran". + assert_eq!(gpu.launches, 4); + } + + /// An async fault reported by the non-blocking query must propagate as a + /// CUDA error — not be swallowed, and not degrade into `Stalled`. + #[test] + fn async_fault_from_stream_query_fails_loudly_as_cuda() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.fault_after = Some(2); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + match err { + DeviceLoopError::Cuda { code, .. } => assert_eq!(code, 700), + other => panic!("async fault must surface as Cuda, got {other:?}"), + } + } + + /// Detection latency: the in-burst probe finds the fault before the burst + /// finishes enqueuing, so a fault on step 1 does not wait for step 4. + #[test] + fn in_burst_probe_detects_a_fault_before_the_burst_completes() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14, 15, 16, 17, 18], 32); + gpu.fault_after = Some(1); // second poll errors -> during step 2 + let mut d = BurstDriver::new(cfg(8, 32), 1).unwrap(); + assert!(d.run_burst(&mut gpu, 8).is_err()); + assert!( + gpu.launches < 8, + "the burst must stop early on a fault, launched {}", + gpu.launches + ); + } + + /// The no-timeout fault detector: stream reports every step retired, but + /// the ring is short. That is definitively wrong and must be an error, not + /// an infinite spin and not a silent short read. + #[test] + fn stream_idle_with_a_short_ring_is_a_fault_not_a_wait() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.commit_is_noop = true; + gpu.lie_idle = true; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + match d.run_burst(&mut gpu, 4).unwrap_err() { + DeviceLoopError::LostTokens { expected, got } => { + assert_eq!((expected, got), (4, 0)); + } + other => panic!("expected LostTokens, got {other:?}"), + } + assert!( + d.spins() < 100, + "idle-but-short must be detected immediately, not spun on ({} spins)", + d.spins() + ); + } + + /// The commit kernel's range check travels back in a pinned word. A driver + /// that forgets to read it returns `Ok` with plausible tokens — exactly the + /// silent-corruption class this repo keeps shipping. So: tokens all arrive, + /// fault word set, and the burst must still fail. + #[test] + fn a_latched_device_fault_fails_the_burst_even_though_every_token_arrived() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.fault = FAULT_TOKEN_RANGE; + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let err = d.run_burst(&mut gpu, 4).unwrap_err(); + assert_eq!(err, DeviceLoopError::DeviceFault(FAULT_TOKEN_RANGE)); + assert_eq!(gpu.head, 4, "the tokens did land; the fault is orthogonal"); + } + + /// EOS truncates the returned tokens but not the launch count. The burst + /// was already enqueued, so the tail was computed and is discarded — the + /// outcome must report both numbers honestly. + #[test] + fn eos_truncates_tokens_but_the_wasted_steps_are_still_reported() { + let mut gpu = FakeGpu::new(vec![11, 12, 2, 14], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let out = d.run_burst(&mut gpu, 4).unwrap(); + assert_eq!(out.tokens, vec![vec![11, 12, 2]]); + assert!(out.hit_eos); + assert_eq!( + out.steps_launched, 4, + "the discarded tail must be visible, not rounded away" + ); + } + + /// Consecutive bursts must not re-read the previous burst's slots. The + /// cursor advances over the full burst even when EOS truncated the output. + #[test] + fn cursor_advances_over_the_whole_burst_across_consecutive_bursts() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14, 21, 22, 23, 24], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + let a = d.run_burst(&mut gpu, 4).unwrap(); + let b = d.run_burst(&mut gpu, 4).unwrap(); + assert_eq!(a.tokens, vec![vec![11, 12, 13, 14]]); + assert_eq!(b.tokens, vec![vec![21, 22, 23, 24]]); + assert_eq!(gpu.launches, 8); + } + + /// Ring wraparound is arithmetic, and arithmetic is where silent corruption + /// lives. Drive past the ring size and check the tokens are still right. + #[test] + fn tokens_survive_ring_wraparound() { + let script: Vec = (100..140).collect(); + let mut gpu = FakeGpu::new(script.clone(), 8); + let mut d = BurstDriver::new(cfg(4, 8), 1).unwrap(); + let mut seen = Vec::new(); + for _ in 0..10 { + let out = d.run_burst(&mut gpu, 4).unwrap(); + seen.extend(out.tokens[0].iter().copied()); + } + assert_eq!(seen, script, "wraparound must not permute or drop tokens"); + } + + /// A burst that could lap the ring is refused up front rather than + /// discovered as scrambled tokens. + #[test] + fn a_burst_that_would_lap_the_ring_is_refused() { + assert!(matches!( + BurstDriver::new(cfg(16, 8), 1), + Err(DeviceLoopError::BadConfig(_)) + )); + let mut gpu = FakeGpu::new(vec![1; 32], 8); + let mut d = BurstDriver::new(cfg(4, 8), 1).unwrap(); + assert!(matches!( + d.run_burst(&mut gpu, 8), + Err(DeviceLoopError::BadConfig(_)) + )); + } + + /// A kernel that hangs *without* faulting reports Busy forever and lands + /// nothing. No CUDA error is ever raised and the stream never goes idle, so + /// neither of the two real detectors can fire — only the liveness backstop + /// can end the burst. Without it this is an unkillable spin in the engine. + #[test] + fn a_hung_stream_trips_the_liveness_backstop() { + let mut gpu = FakeGpu::new(vec![11, 12, 13, 14], 16); + gpu.never_retires = true; + let mut c = cfg(4, 16); + c.max_spins = 500; + let mut d = BurstDriver::new(c, 1).unwrap(); + match d.run_burst(&mut gpu, 4).unwrap_err() { + DeviceLoopError::Stalled { spins } => assert_eq!(spins, 500), + other => panic!("expected Stalled, got {other:?}"), + } + } + + /// `poll_every == 0` would divide by zero in the drain. Refuse it at + /// construction rather than panicking mid-burst. + #[test] + fn zero_poll_every_is_refused_rather_than_dividing_by_zero() { + let mut c = cfg(4, 16); + c.poll_every = 0; + assert!(matches!( + BurstDriver::new(c, 1), + Err(DeviceLoopError::BadConfig(_)) + )); + } + + /// `plan_steps` must never walk the device position out of the + /// fixed-capacity KV window, and must report 0 rather than a short burst + /// the caller might mistake for success. + #[test] + fn plan_steps_clamps_to_the_position_window_and_the_token_budget() { + let mut c = cfg(8, 64); + c.position_limit = 100; + assert_eq!(c.plan_steps(0, 1000), 8); + assert_eq!(c.plan_steps(95, 1000), 5, "must not pass position_limit"); + assert_eq!(c.plan_steps(100, 1000), 0, "no room left"); + assert_eq!( + c.plan_steps(120, 1000), + 0, + "already past: saturating, not wrapping" + ); + assert_eq!(c.plan_steps(0, 3), 3, "must not exceed the token budget"); + } + + /// Zero steps is a caller bug, not a quiet no-op that reports success. + #[test] + fn zero_steps_is_an_error_not_a_silent_success() { + let mut gpu = FakeGpu::new(vec![], 16); + let mut d = BurstDriver::new(cfg(4, 16), 1).unwrap(); + assert!(matches!( + d.run_burst(&mut gpu, 0), + Err(DeviceLoopError::BadConfig(_)) + )); + assert_eq!(gpu.launches, 0); + } + + fn ok_facts() -> AdmissionFacts { + AdmissionFacts { + batch_size: 1, + seq_len: 1, + graph_ready: true, + greedy: true, + return_logprobs: false, + host_side_logits_work: false, + current_position: 10, + } + } + + /// Admission must refuse every case where the device sampler would produce + /// a token the host path would not. Each of these is a silent behaviour + /// change if it slips through, so each gets its own assertion rather than + /// one "returns Err" catch-all. + #[test] + fn admission_refuses_every_case_the_device_sampler_cannot_reproduce() { + let c = cfg(4, 64); + // The happy path first — otherwise a guard that refuses everything + // would pass every negative case below. + assert_eq!(admit(&ok_facts(), &c, true), Ok(4)); + + assert_eq!(admit(&ok_facts(), &c, false), Err(Refusal::NotEnabled)); + + let mut f = ok_facts(); + f.seq_len = 32; + assert_eq!(admit(&f, &c, true), Err(Refusal::NotDecodeStep)); + + let mut f = ok_facts(); + f.batch_size = 2; + assert_eq!(admit(&f, &c, true), Err(Refusal::NotBatchOne)); + + let mut f = ok_facts(); + f.graph_ready = false; + assert_eq!(admit(&f, &c, true), Err(Refusal::GraphNotReady)); + + let mut f = ok_facts(); + f.greedy = false; + assert_eq!( + admit(&f, &c, true), + Err(Refusal::NotGreedy), + "stochastic sampling must fall back: device Splitmix64 and host Isaac64 draw \ + different tokens, which breaks seeded reproducibility" + ); + + let mut f = ok_facts(); + f.return_logprobs = true; + assert_eq!(admit(&f, &c, true), Err(Refusal::WantsLogprobs)); + + let mut f = ok_facts(); + f.host_side_logits_work = true; + assert_eq!(admit(&f, &c, true), Err(Refusal::HostSideLogitsWork)); + } + + /// Admission must also clamp the burst to the KV window, and refuse rather + /// than return a zero-step "success" the caller could mistake for a launch. + #[test] + fn admission_clamps_the_burst_to_the_kv_window() { + let mut c = cfg(4, 64); + c.position_limit = 100; + let mut f = ok_facts(); + f.current_position = 98; + assert_eq!(admit(&f, &c, true), Ok(2), "clamped to the window"); + f.current_position = 100; + assert_eq!(admit(&f, &c, true), Err(Refusal::NoPositionRoom)); + } + + /// The parked tokens are handed out in generation order, exactly once, 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() { + 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(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); + } + + /// A leftover token must never reach the next sequence. `clear` is the only + /// thing standing between a fall-back and one sequence's token appearing in + /// another's output. + #[test] + fn clearing_pending_tokens_prevents_leaking_them_into_the_next_sequence() { + stand_down(); + set_device_loop_eligible(11, true); + push_pending_tokens(&[7, 8, 9]); + clear_pending_tokens(); + 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() { + stand_down(); + set_device_loop_eligible(11, true); + push_pending_tokens(&[5, -1, 6]); + 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 + /// sequence always falls back rather than engaging on parameters nothing + /// has checked yet. + #[test] + fn eligibility_starts_false_and_stand_down_resets_everything() { + stand_down(); + assert!(!device_loop_eligible()); + + 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(11), + PendingTake::Empty, + "stand_down must clear tokens" + ); + } + + /// Default-off, and explicitly pinnable off. + #[test] + fn the_opt_in_defaults_to_off() { + // Not using the process env here: these run in parallel with other + // tests. The parsing rule is what matters and it is exercised directly. + for (val, want) in [ + ("", false), + ("0", false), + ("false", false), + ("off", false), + ("no", false), + ("1", true), + ("true", true), + ] { + let v = val.trim().to_ascii_lowercase(); + let enabled = !(v.is_empty() || v == "0" || v == "false" || v == "off" || v == "no"); + assert_eq!(enabled, want, "for {val:?}"); + } + // And with the variable absent entirely. + std::env::remove_var(ENV_ENABLE); + assert!(!device_loop_enabled()); + } +} diff --git a/arc-cuda-graph/src/ffi.rs b/arc-cuda-graph/src/ffi.rs index ae78c6f76..813a2a2f4 100644 --- a/arc-cuda-graph/src/ffi.rs +++ b/arc-cuda-graph/src/ffi.rs @@ -14,6 +14,21 @@ pub type CUgraphConditionalHandle = u64; #[cfg(feature = "cuda")] pub type CUmemoryPool = *mut std::ffi::c_void; #[cfg(feature = "cuda")] +pub type CUcontext = *mut std::ffi::c_void; +/// `cuda.h:1944` (CUDA 13.1) — +/// `#define CU_GRAPH_COND_ASSIGN_DEFAULT 0x1` "Default value is applied when +/// graph is launched." +/// +/// `cuda.h:21919` documents `defaultLaunchValue` as "Applied at the beginning +/// of each graph execution **if CU_GRAPH_COND_ASSIGN_DEFAULT is set in +/// flags**". Creating the handle with `flags = 0` therefore leaves the +/// condition at 0 on entry and a WHILE body executes **zero** times — the +/// graph launches, returns success, and generates nothing. Measured on +/// `arc-v4-stack` (H200, CUDA 13.1): `flags=0` → body ran 0 times; +/// `flags=CU_GRAPH_COND_ASSIGN_DEFAULT` → body ran exactly N times. +#[cfg(feature = "cuda")] +pub const CU_GRAPH_COND_ASSIGN_DEFAULT: u32 = 0x1; +#[cfg(feature = "cuda")] pub type CUdevice = i32; #[cfg(feature = "cuda")] @@ -70,6 +85,14 @@ pub enum CudaGraphNodeType { Conditional = 13, } +/// `CUDA_CONDITIONAL_NODE_PARAMS`, `cuda.h:1958` (CUDA 13.1). +/// +/// `phGraph_out` is an **OUT** field: "CUDA-owned array populated with +/// conditional node child graphs during creation of the node." Callers leave +/// it null and read it back after `cuGraphAddNode`; assigning a caller-owned +/// pointer to it accomplishes nothing because the driver overwrites it. +/// +/// `ctx` (the 5th field) was missing entirely from this struct. #[cfg(feature = "cuda")] #[repr(C)] pub struct CudaConditionalNodeParams { @@ -77,6 +100,7 @@ pub struct CudaConditionalNodeParams { pub cond_type: CUgraphConditionalNodeType, pub size: u32, pub body_graph_out: *mut CUgraph, + pub ctx: CUcontext, } #[cfg(feature = "cuda")] @@ -162,6 +186,7 @@ extern "C" { pub fn cuGraphLaunch(exec: CUgraphExec, stream: CUstream) -> u32; pub fn cuGraphExecDestroy(exec: CUgraphExec) -> u32; pub fn cuGraphDestroy(graph: CUgraph) -> u32; + pub fn cuGraphGetNodes(graph: CUgraph, nodes: *mut CUgraphNode, num_nodes: *mut usize) -> u32; // ========== Memory pools ========== pub fn cuMemPoolCreate(pool: *mut CUmemoryPool, props: *const CUmemPoolProps) -> u32; @@ -179,22 +204,62 @@ extern "C" { pub fn cuStreamCreate(stream: *mut CUstream, flags: u32) -> u32; pub fn cuStreamDestroy_v2(stream: CUstream) -> u32; pub fn cudaStreamSynchronize(stream: CUstream) -> u32; + /// Non-blocking completion probe. Returns `cudaSuccess` when every + /// enqueued item has retired and `cudaErrorNotReady` (600) while work is + /// outstanding. NVIDIA's reference adds: *"Note that this function may + /// also return error codes from previous, asynchronous launches."* That + /// is what lets `device_loop` surface an async graph fault without ever + /// blocking the stream — see `device_loop::CudaDeviceOps::poll_stream`. + pub fn cudaStreamQuery(stream: CUstream) -> u32; // ========== Conditional nodes (CUDA 12.4+) ========== - pub fn cudaGraphConditionalHandleCreate( + /// `cuda.h:21932` (CUDA 13.1) — the driver form takes a `CUcontext`, + /// which the conditional node params must match ("Context on which to run + /// the node. Must match context used to create the handle and all body + /// nodes", `cuda.h:1986`). + pub fn cuGraphConditionalHandleCreate( handle: *mut CUgraphConditionalHandle, graph: CUgraph, - default_value: u32, + ctx: CUcontext, + default_launch_value: u32, flags: u32, ) -> u32; - pub fn cudaGraphSetConditional(handle: CUgraphConditionalHandle, value: u32) -> u32; - pub fn cudaGraphAddNode( - graph: CUgraph, + pub fn cuCtxGetCurrent(ctx: *mut CUcontext) -> u32; + /// `cuda.h:21829` (CUDA 13.1): + /// ```text + /// CUresult cuGraphAddNode(CUgraphNode *phGraphNode, CUgraph hGraph, + /// const CUgraphNode *dependencies, + /// const CUgraphEdgeData *dependencyData, + /// size_t numDependencies, + /// CUgraphNodeParams *nodeParams); + /// ``` + /// The OUT node pointer is the **first** argument and there are **six** of + /// them. This was previously declared as + /// `cudaGraphAddNode(graph, node_out, deps, num_deps, params)` — five + /// arguments with the first two transposed — so the driver received the + /// `CUgraph` handle in the slot where it expected `CUgraphNode *` and + /// wrote the new node handle through that value, while `numDependencies` + /// received a pointer and `nodeParams` an uninitialised register. + pub fn cuGraphAddNode( node_out: *mut CUgraphNode, + graph: CUgraph, dependencies: *const CUgraphNode, + dependency_data: *const std::ffi::c_void, num_dependencies: usize, params: *mut CudaGraphNodeParams, ) -> u32; + /// `cuda.h:15665` (CUDA 13.1). `cuda.h:1976` names this as a supported way + /// to populate a conditional node's body graph. Ordinary + /// `cuStreamBeginCapture_v2` always creates a **new** graph, so a body + /// captured that way cannot be attached to the conditional node. + pub fn cuStreamBeginCaptureToGraph( + stream: CUstream, + graph: CUgraph, + dependencies: *const CUgraphNode, + dependency_data: *const std::ffi::c_void, + num_dependencies: usize, + mode: CUstreamCaptureMode, + ) -> u32; // ========== Pinned host memory ========== pub fn cudaHostAlloc(ptr: *mut *mut std::ffi::c_void, size: usize, flags: u32) -> u32; diff --git a/arc-cuda-graph/src/graph.rs b/arc-cuda-graph/src/graph.rs index c288998f7..e5ff6aefa 100644 --- a/arc-cuda-graph/src/graph.rs +++ b/arc-cuda-graph/src/graph.rs @@ -108,6 +108,10 @@ pub struct CudaGraphRunner { /// context-dependent and the real fix is shape-constant buffers, not more /// warmup. deferred_passes_remaining: u32, + /// Deferred passes still available to `grant_extra_deferred_pass`. Counts + /// down independently of `deferred_passes_remaining` so that "keep warming + /// while the profile grows" cannot become an unbounded loop. + deferred_extra_budget: u32, } #[cfg(feature = "cuda")] @@ -121,13 +125,40 @@ impl CudaGraphRunner { Self::new_with_passes(device, warmup_steps, Self::default_deferred_passes()) } - /// Number of deferred-free warmup passes, from `ARC_GRAPH_DEFERRED_PASSES` - /// (default 1, the historical behaviour). See `deferred_passes_remaining`. + /// Number of deferred-free warmup passes, from `ARC_GRAPH_DEFERRED_PASSES`. + /// + /// Default 4, not the historical 1. One pass only suffices if every decode + /// step allocates the same set of sizes, and V4's does not: the rolling + /// compressor's retained tail is rebuilt at width `tokens - base`, and + /// `base` jumps a whole `ratio` at a group boundary while `tokens` climbs by + /// one, so the size cycles through `ratio` consecutive values (measured: + /// `4096 × {18,19,20,21}`). One pass warms one phase of that cycle; the + /// captured step lands on another phase and allocates a size the pool has + /// never held. `ratio` is 4 on the shipped V4 configuration. + /// + /// This is a floor, not the answer — `grant_extra_deferred_pass` keeps + /// extending while the observed size profile is still growing, so a model + /// with a longer cycle is covered without retuning this number. + /// See `deferred_passes_remaining`. pub fn default_deferred_passes() -> u32 { std::env::var("ARC_GRAPH_DEFERRED_PASSES") .ok() .and_then(|v| v.parse::().ok()) - .unwrap_or(1) + .unwrap_or(4) + } + + /// Upper bound on passes granted by `grant_extra_deferred_pass`, from + /// `ARC_GRAPH_DEFERRED_MAX` (default 24). The bound is what keeps a model + /// whose allocation sizes grow monotonically with context length from + /// warming forever: it will exhaust the budget, and the capture-miss assert + /// will then refuse the capture and name the sizes. That refusal is the + /// useful result — it says the fix is shape-constant buffers (or allocator + /// size-class bucketing), not more warmup. + pub fn default_deferred_max() -> u32 { + std::env::var("ARC_GRAPH_DEFERRED_MAX") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(24) } pub fn new_with_passes( @@ -176,6 +207,7 @@ impl CudaGraphRunner { verify_remaining: 0, verify_failed: false, deferred_passes_remaining: 0, + deferred_extra_budget: 0, }); } @@ -198,6 +230,7 @@ impl CudaGraphRunner { verify_remaining: Self::default_verify_replays(), verify_failed: false, deferred_passes_remaining: deferred_passes, + deferred_extra_budget: Self::default_deferred_max(), }) } @@ -273,6 +306,29 @@ output_trusted={} verify_remaining={} verify_failed={}", } } + /// Ask for one more deferred-free warm pass, spending from a fixed budget. + /// Returns whether one was granted. + /// + /// The caller grants a pass whenever the last one taught the allocator a + /// size it had not seen, and stops when a pass adds nothing. That turns + /// "how many passes cover the allocation cycle?" from a constant that has + /// to be right into an observation that terminates on its own, while + /// `deferred_extra_budget` guarantees termination even when the sizes never + /// converge. + pub fn grant_extra_deferred_pass(&mut self) -> bool { + if self.deferred_extra_budget == 0 { + return false; + } + self.deferred_extra_budget -= 1; + self.deferred_passes_remaining += 1; + true + } + + /// Deferred passes still available to `grant_extra_deferred_pass`. + pub fn deferred_extra_budget(&self) -> u32 { + self.deferred_extra_budget + } + /// How many replays must be proven against an eager forward before the /// graph's output is trusted. `ARC_GRAPH_VERIFY_REPLAYS`, default 3. pub fn default_verify_replays() -> u32 { @@ -728,6 +784,49 @@ output_trusted={} verify_remaining={} verify_failed={}", } } + /// The handles the device decode loop needs to replay this graph **without + /// the synchronize above**: the exec to launch and the graph-owned output + /// tensor to sample from, in place, on device. + /// + /// Handing these out is the whole difference between `replay()` and + /// `device_loop`. `replay()` launches, blocks on `cudaStreamSynchronize` + /// (line 362) and returns the output for the *host* to sample — which + /// copies logits down and serialises the step. `device_loop` takes the same + /// two handles, launches, samples on device and scatters the token back + /// into the pinned input buffer, so the next launch needs no host at all. + /// + /// The returned `Tensor`'s storage was allocated from the graph's private + /// pool during capture, so its address is stable for the life of the + /// `CapturedGraph` — which is exactly why sampling it in place is sound. + /// The clone is an `Arc` bump and does **not** extend that lifetime past + /// the runner's; callers must not hold it after dropping the runner. + /// + /// 🔴 The same BOOBY TRAP `replay` documents applies here verbatim: this + /// output aliases storage the next launch overwrites, so it is not a + /// snapshot. `device_loop` is safe against it for a specific reason — + /// every access is enqueued on this same stream, in order, and none of + /// them is a host read — and `normal.rs` enforces at runtime the invariant + /// that lets the aliasing tensor be handed back at all. Anything else + /// taking these handles must copy. + pub fn replay_handles(&self, batch_size: usize) -> Option<(CUgraphExec, Tensor)> { + let captured = self.graphs.get(&batch_size)?; + let out = captured.output.clone()?; + Some((captured.exec, out)) + } + + /// The stream capture and replay run on. `device_loop` launches its commit + /// kernel here so it is ordered behind the graph without an event. + pub fn stream(&self) -> CUstream { + self.stream + } + + /// Count a decode step driven by the device loop rather than by + /// [`replay`](Self::replay), so `status_line` stays honest about how much + /// work this runner actually did. + pub fn record_device_loop_steps(&mut self, steps: u64) { + self.replays += steps; + } + pub fn disable(&mut self) { self.enabled = false; } diff --git a/arc-cuda-graph/src/lib.rs b/arc-cuda-graph/src/lib.rs index 8e534ed6d..1760ad210 100644 --- a/arc-cuda-graph/src/lib.rs +++ b/arc-cuda-graph/src/lib.rs @@ -12,6 +12,11 @@ pub mod autonomous; pub mod buffers; pub mod decode_forward; pub mod dedicated; +/// Device-resident decode loop: N steps, N `cuGraphLaunch` calls, zero stream +/// synchronizations. Opt-in via `ARC_GRAPH_DEVICE_LOOP`. The driver and its +/// policy compile without CUDA so they are unit-tested on any host; only the +/// pointer plumbing is gated. +pub mod device_loop; pub mod ffi; pub mod gemv_ffi; pub mod graph; @@ -82,6 +87,15 @@ pub use buffers::{DecodeInputBuffers, DecodeState}; pub use decode_forward::{decode_forward, DecodeBuffers, LayerKvCache, PagedAttentionState}; #[cfg(feature = "cuda")] pub use dedicated::DedicatedDecodePath; +pub use device_loop::{ + admit, clear_pending_tokens, device_loop_eligible, device_loop_enabled, device_loop_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}; #[cfg(feature = "cuda")] pub use graph::CudaGraphRunner; #[cfg(feature = "cuda")] diff --git a/arc-cuda-graph/src/sampling_cuda.rs b/arc-cuda-graph/src/sampling_cuda.rs index 8d6feb95a..67d52a939 100644 --- a/arc-cuda-graph/src/sampling_cuda.rs +++ b/arc-cuda-graph/src/sampling_cuda.rs @@ -209,6 +209,59 @@ extern "C" { cfg: SamplingParams, stream: CUstream, ); + + // Hybrid: exact enumeration while the nucleus is small, threshold + // bisection past a fixed budget, chosen block-uniformly ON DEVICE so a + // captured graph never has to ask the host which branch to take. + // + // MEASURED, H200, vocab=129280, host+GPU verified exclusive before and + // after (support width -> legacy / hybrid us): + // 1 -> 134.9/166.6 | 8 -> 503.0/624.2 | 64 -> 3127.8/4178.5 + // 512 -> 24320.7/4185.0 | 4096 -> 205125.2/4187.3 + // 12928 -> 664200.0/4188.5 + // Costs +24% on the peaked distributions real models produce (32 us, i.e. + // 0.05% of a 66.68 ms V4 decode step) and removes a 664 ms cliff -- 158x + // at the diffuse tail, where the enumerating sampler is ~10x an entire + // decode step for a single token. + pub fn arc_launch_sampler_f32_hybrid( + logits: *const f32, + freq_counts: *const u32, + rng_state: *mut u64, + token_ids: *mut i32, + probs_scratch: *mut f32, + keep_idx_scratch: *mut i32, + keep_p_scratch: *mut f32, + vocab: i32, + batch: i32, + cfg: SamplingParams, + stream: CUstream, + ); + pub fn arc_launch_sampler_bf16_hybrid( + logits: *const std::ffi::c_void, + freq_counts: *const u32, + rng_state: *mut u64, + token_ids: *mut i32, + probs_scratch: *mut f32, + keep_idx_scratch: *mut i32, + keep_p_scratch: *mut f32, + vocab: i32, + batch: i32, + cfg: SamplingParams, + stream: CUstream, + ); + pub fn arc_launch_sampler_f16_hybrid( + logits: *const std::ffi::c_void, + freq_counts: *const u32, + rng_state: *mut u64, + token_ids: *mut i32, + probs_scratch: *mut f32, + keep_idx_scratch: *mut i32, + keep_p_scratch: *mut f32, + vocab: i32, + batch: i32, + cfg: SamplingParams, + stream: CUstream, + ); } /// Raw device pointer for a CUDA tensor, byte-offset aware. @@ -311,6 +364,15 @@ impl CudaSampler { }) } + /// Bind the sampler to `stream`. + /// + /// `new()` picks up the device-default stream. Anything that is going to + /// be **captured** into a CUDA graph must launch on the capture stream, or + /// the sampler's kernels are simply not recorded into the graph. + pub fn set_stream(&mut self, stream: CUstream) { + self.stream = stream; + } + /// Run the sampler. Logits: [batch, vocab] of the configured dtype. /// `freq_counts`: optional [batch, vocab] u32 token-count tensor (one per /// batch row); pass `None` when no penalties are needed. @@ -370,7 +432,7 @@ impl CudaSampler { unsafe { match self.dtype { - DType::F32 => arc_launch_sampler_f32( + DType::F32 => arc_launch_sampler_f32_hybrid( logits_ptr as *const f32, freq_ptr, rng_ptr, @@ -383,7 +445,7 @@ impl CudaSampler { params, self.stream, ), - DType::BF16 => arc_launch_sampler_bf16( + DType::BF16 => arc_launch_sampler_bf16_hybrid( logits_ptr, freq_ptr, rng_ptr, @@ -396,7 +458,7 @@ impl CudaSampler { params, self.stream, ), - DType::F16 => arc_launch_sampler_f16( + DType::F16 => arc_launch_sampler_f16_hybrid( logits_ptr, freq_ptr, rng_ptr, diff --git a/arc-cuda-graph/src/weights.rs b/arc-cuda-graph/src/weights.rs index f8891c311..bdf77b4b0 100644 --- a/arc-cuda-graph/src/weights.rs +++ b/arc-cuda-graph/src/weights.rs @@ -610,6 +610,16 @@ pub fn tensor_device_ptr(tensor: &Tensor) -> candle_core::Result { // fed with (radix top-k `seq_lens`, `CudaSampler` token ids and // keep-list scratch). Candle's `Tensor::from_vec(Vec, …)` // produces it, so leaving it out disabled those paths wholesale. + // + // Concretely, from #130: `CudaSampler::sample` REQUIRES + // `token_ids` to be I32 (`sampling_cuda.rs:347`) and allocates + // an I32 `keep_idx_scratch` (`:298`). Both resolve through this + // function, so before I32 was listed both fell to the catch-all + // `bail!` and `CudaSampler::sample()` returned `unsupported + // dtype I32` on EVERY call — the replay-safe top-k/top-p + // sampler could never run on a GPU. Its test suite is a CPU + // simulator (`gpu_algorithm_simulate`), which by construction + // cannot observe that. DType::I32 => { let s = cuda_storage.as_cuda_slice::()?; let (p, _) = s.device_ptr(s.stream()); diff --git a/arc-tools/arcspec_perseq_ladder.sh b/arc-tools/arcspec_perseq_ladder.sh new file mode 100755 index 000000000..062eb2217 --- /dev/null +++ b/arc-tools/arcspec_perseq_ladder.sh @@ -0,0 +1,544 @@ +#!/bin/bash +# ArcSpec — per-sequence KV advance: THROUGHPUT on the ladder B = 1, 8, 32, 128. +# +# ───────────────────────────────────────────────────────────────────────────── +# THE QUESTION +# +# #116 proved per-sequence advance *works* at B=8 (it returned 1 token with no +# `finish_reason` before; it returns the full 48 after). It did not prove the +# mechanism is worth anything. MTP was measured at 1.93 tok/step at one user +# collapsing to 1.06 at 128, and this stack targets exactly that collapse. +# +# ⚠️ REPORT BOTH NUMBERS OR THE RESULT IS UNREADABLE. tok/step rising while +# aggregate throughput falls has already happened on this chain (+9%/+13% +# tok/step, −5×/−20% aggregate). A tok/step number alone is not a throughput +# result, and an aggregate number alone cannot say what moved. So every cell +# reports: +# +# aggregate tok/s (wall clock, the number that matters) +# tok_per_step (per-sequence multiplier) +# tok_per_batch_step (tokens out per target forward, whole batch) +# batch_steps/s (the decomposition: aggregate ≈ tok_per_batch_step +# × batch_steps/s — when aggregate does not move, +# this says which factor ate it, in ms) +# +# ───────────────────────────────────────────────────────────────────────────── +# ⚠️ TWO PROMPT REGIMES, MEAN-MATCHED. Both are needed, for opposite reasons. +# +# spread 144 24 320 64 260 40 200 96 words, cycled across workers +# uniform 144 words for every worker (= the spread's mean, and its first +# element, so B=1 is byte-identical in both regimes) +# +# **spread** is the regime this stack exists for. Uniform prompts hide the +# failure mode: the `xs` window defect #116 fixed could not even be reached with +# equal-length prompts, because the retained window is sized for the greediest +# row and a short row only becomes shorter-than-the-window when arrivals differ. +# A benchmark that sends one prompt to every worker measures the case that +# already worked. +# +# **uniform** is the only regime where the batch is WHOLE, and that is why it is +# not optional. The scheduler A/B has now measured the bucketing law directly: +# +# running bucket size = B / (distinct cache lengths) +# +# holding at 8/8=1 and 32/8=4, with `1 running, 7 waiting` sustained. Under a +# spread of 8 distinct lengths, B=8 therefore runs ONE sequence at a time — that +# chain measured B=8 spread at 7.91 tok/s against B=1's 15.36, i.e. batching is +# *negative* on realistic traffic. So a flat aggregate in the spread arm is +# unreadable on its own: it could be this fix not paying, or it could be that +# the batch never existed. Uniform is where a whole batch actually forms, so it +# is the only place this fix's effect on aggregate is visible without the +# serialisation swamping it. +# +# Mean-matched on purpose: same total prompt tokens in both regimes, so the only +# thing that differs is the *spread*, and a difference is attributable to +# raggedness rather than to prompt length. +# +# Session 8 measured V4 serving dying outright at ~1,055-word prompts and this +# branch does NOT carry #97 (`pin V4 against length`), so 320 keeps a >3x margin +# to that cliff while still being a real spread. +# +# ───────────────────────────────────────────────────────────────────────────── +# ⚠️ THE SCHEDULER CONFOUND, MEASURED RATHER THAN ASSUMED. +# +# Both schedulers bucket the running set by EXACT cache length and run ONE +# bucket per step, preempting the rest (`default_scheduler.rs`, +# `paged_attention/scheduler.rs`). So "B=128" in this harness can be a 3-wide +# step in the engine, and an aggregate number measured inside that is a number +# about the scheduler, not about KV advance. +# +# `SCHED[agg]` (added with this harness) is emitted on the SAME log fence as +# `MTP[agg]`, so each cell reports, over exactly its own window: +# +# buckets_per_step 1.0 = the batch was never split +# running_bucket_size the width the engine ACTUALLY ran +# offered_per_step the width this harness thinks it asked for +# +# If `running_bucket_size` << `offered_per_step`, the throughput cell is a +# scheduler measurement and says so. +# +# ───────────────────────────────────────────────────────────────────────────── +# THE DESIGN +# +# Two server processes, one per arm; four concurrency cells inside each. +# OFF no flags (control) +# ON ARC_V4_XS_PER_SEQ=1 ARC_MTP_PER_SEQ_KV=1 (treatment) +# Same binary, same model, same prompts, same MTP depth, same --max-seqs. One +# server per arm rather than one per cell because the engine counters are +# cumulative and monotone: differencing them across a wall-clock fence gives a +# cell's own numbers without paying the ~2m30s V4 load eight times. +# +# ⚠️ ENGAGEMENT IS ASSERTED, NOT ASSUMED (D18). "per-sequence KV advance is ON" +# and "Ragged batch admission is ON" must be PRESENT in the ON log and ABSENT +# in OFF, "cannot honour it" must be absent from ON, and `per_seq_steps` must be +# non-zero — granted-but-inert has to show as a visible zero, not as silence. +# +# Environment failure exits 2. A gate that ran and failed exits 1. +set -u + +ARCENV="${ARCENV:-/root/arcenv.sh}" +WT="${WT:-/root/wt-ladder}" +BIN="${BIN:-$WT/target/release/mistralrs}" +MODEL="${MODEL:-/root/models/v4-src}" +UQFF="${UQFF:-/root/models/v4-uqff/qtip2b-0.uqff}" +ARCH="${ARCH:-deepseekv4}" +PREFLIGHT="${PREFLIGHT:-/root/arc-tools/gpu_box_preflight.sh}" +OUT="${OUT:-/root/logs/ladder}" +STATUS="${STATUS:-/root/logs/ladder.status}" +PORT="${PORT:-1243}" +MTP_DEPTH="${MTP_DEPTH:-3}" +MAX_SEQS="${MAX_SEQS:-256}" +SEED="${SEED:-1234}" +BATCHES="${BATCHES:-1 8 32 128}" +WARMUP="${WARMUP:-20}" +STEADY="${STEADY:-45}" +MAX_TOKENS="${MAX_TOKENS:-4096}" +# 144 first so B=1 is byte-identical between the two regimes; mean is 144. +WORD_SPREAD="${WORD_SPREAD:-144 24 320 64 260 40 200 96}" +UNIFORM_WORDS="${UNIFORM_WORDS:-144}" +REGIMES="${REGIMES:-uniform spread}" + +mkdir -p "$OUT" "$(dirname "$STATUS")" +: > "$STATUS" +say() { echo "[$(date -u +%H:%M:%SZ)] $*" | tee -a "$STATUS"; } +die() { say "ABORT: $*"; echo "ABORT: $*" >> "$OUT/summary.txt"; exit 2; } + +say "ArcSpec per-seq throughput ladder — start" +[ -f "$ARCENV" ] || die "$ARCENV not found" +# shellcheck disable=SC1090 +. "$ARCENV" +say "env: CUDA_HOME=${CUDA_HOME:-unset} LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-unset}" +[ -x "$PREFLIGHT" ] || die "$PREFLIGHT missing — refusing to measure" +"$PREFLIGHT" > "$OUT/preflight.log" 2>&1 || die "preflight refused the box" +say "preflight OK" +[ -x "$BIN" ] || die "$BIN not found" +SHA=$(git -C "$WT" rev-parse HEAD 2>/dev/null) || die "no git ref at $WT" +say "binary ref: $SHA" +nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv > "$OUT/gpu.txt" 2>&1 + +# ── The load driver. Streaming, so a cell yields tokens even when no request +# finishes inside it (at B=128 V4 needs minutes per request). +cat > "$OUT/drive.py" <<'PYEOF' +import argparse, json, threading, time, urllib.request + +FILLER = ("alpha bravo charlie delta echo foxtrot golf hotel india juliet " + "kilo lima mike november oscar papa quebec romeo sierra tango").split() + +def prompt_of(words): + return ("Continue this list, one item per line. " + + " ".join(FILLER[i % len(FILLER)] for i in range(words))) + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--port", type=int, required=True) + ap.add_argument("--k", type=int, required=True) + ap.add_argument("--warmup", type=float, required=True) + ap.add_argument("--steady", type=float, required=True) + ap.add_argument("--spread", required=True, help="comma-separated word counts") + ap.add_argument("--max-tokens", type=int, default=4096) + ap.add_argument("--seed", type=int, default=1234) + a = ap.parse_args() + + # 🔑 RAGGED BY CONSTRUCTION: worker i takes spread[i % len(spread)], so the + # arrivals differ in length the way real ones do. Deterministic (no RNG), so + # the two arms see byte-identical prompt sets. + spread = [int(x) for x in a.spread.split(",") if x] + prompts = [prompt_of(spread[i % len(spread)]) for i in range(a.k)] + url = f"http://127.0.0.1:{a.port}/v1/completions" + stop = threading.Event() + lock = threading.Lock() + n = {"chunks": 0, "errors": 0, "started": 0, "finished": 0} + counting = [False] + + def worker(wid): + body_prompt = prompts[wid] + while not stop.is_set(): + body = json.dumps({ + "model": "default", "prompt": body_prompt, + "max_tokens": a.max_tokens, "temperature": 0, + "seed": a.seed, "stream": True, + }).encode() + req = urllib.request.Request( + url, data=body, headers={"Content-Type": "application/json"}) + try: + with lock: + n["started"] += 1 + with urllib.request.urlopen(req, timeout=1800) as r: + for raw in r: + if stop.is_set(): + break + line = raw.decode("utf-8", "ignore").strip() + if not line.startswith("data:"): + continue + payload = line[5:].strip() + if payload == "[DONE]": + with lock: + n["finished"] += 1 + break + try: + ch = json.loads(payload) + except Exception: + continue + txt = (ch.get("choices") or [{}])[0].get("text", "") + if txt: + with lock: + if counting[0]: + n["chunks"] += 1 + except Exception: + with lock: + n["errors"] += 1 + time.sleep(0.5) + + threads = [threading.Thread(target=worker, args=(i,), daemon=True) + for i in range(a.k)] + for t in threads: + t.start() + time.sleep(a.warmup) + with lock: + n["chunks"] = 0 + counting[0] = True + t0 = time.time() + fence = {"t_start": time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime())} + time.sleep(a.steady) + with lock: + counting[0] = False + snap = dict(n) + fence["t_end"] = time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime()) + stop.set() + elapsed = time.time() - t0 + snap.update(fence) + snap["k"] = a.k + snap["prompt_words"] = [spread[i % len(spread)] for i in range(a.k)] + snap["elapsed_s"] = round(elapsed, 3) + snap["driver_tok_s"] = round(snap["chunks"] / elapsed, 3) if elapsed else 0.0 + print(json.dumps(snap)) + for t in threads: + t.join(timeout=3) + +main() +PYEOF + +# ── The extractor. Differences cumulative MTP[agg] AND SCHED[agg] across the +# cell's wall-clock fence. Everything reported is a delta of raw counts. +cat > "$OUT/fence.py" <<'PYEOF' +"""Delta of the engine's cumulative counters across one cell's fence. + +`MTP[agg] ... steps=S committed=C batch_steps=BS` and `SCHED[agg] ... calls=N +buckets=B offered=O chosen=C` are cumulative for the process and emitted on the +same log fence, so a cell's own numbers are the difference between the last +line at or before its end and the last line at or before its start. + +Reports nothing it did not measure: a cell with <2 usable lines prints NA with +a reason rather than 0.0, because a zero would read as "the engine produced +nothing" when it means "the reporter never fired twice" (D18). +""" +import json, re, sys + +TS = re.compile(r"(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})") +FIELDS = { + "MTP[agg]": ("accepted", "proposed", "steps", "drafted_steps", "committed", + "batch_steps"), + "SCHED[agg]": ("calls", "buckets", "offered", "chosen", "shattered"), +} + +def parse(path, marker): + keys = FIELDS[marker] + out = [] + for line in open(path, errors="ignore"): + if marker not in line: + continue + mt = TS.search(line) + if not mt: + continue + rec, ok = {"ts": mt.group(1)}, True + for k in keys: + m = re.search(rf"\b{k}=(\d+)\b", line) + if not m: + ok = False + break + rec[k] = int(m.group(1)) + if ok: + out.append(rec) + return out + +def bracket(recs, t0, t1): + """(last record at or before t0, last record at or before t1).""" + before = [r for r in recs if r["ts"] <= t0] + within = [r for r in recs if r["ts"] <= t1] + if not before or not within: + return None, None + return before[-1], within[-1] + +def delta(path, marker, t0, t1): + recs = parse(path, marker) + if len(recs) < 2: + return {"NA": f"{marker}: only {len(recs)} line(s) in the whole log"} + a, b = bracket(recs, t0, t1) + if a is None or b is None: + return {"NA": f"{marker}: no line brackets [{t0}, {t1}]"} + d = {k: b[k] - a[k] for k in FIELDS[marker]} + if all(v == 0 for v in d.values()): + return {"NA": f"{marker}: counters did not advance inside the fence"} + return d + +def main(): + log, t0, t1, elapsed, driver_tok_s = ( + sys.argv[1], sys.argv[2], sys.argv[3], float(sys.argv[4]), float(sys.argv[5])) + out = {"t_start": t0, "t_end": t1, "elapsed_s": elapsed, + "aggregate_tok_s": driver_tok_s} + m = delta(log, "MTP[agg]", t0, t1) + s = delta(log, "SCHED[agg]", t0, t1) + out["mtp_raw"], out["sched_raw"] = m, s + r = lambda x: round(x, 4) + if "NA" not in m: + out["tok_per_step"] = r(m["committed"] / m["steps"]) if m["steps"] else None + out["tok_per_batch_step"] = ( + r(m["committed"] / m["batch_steps"]) if m["batch_steps"] else None) + out["mean_batch"] = r(m["steps"] / m["batch_steps"]) if m["batch_steps"] else None + out["batch_steps_per_s"] = r(m["batch_steps"] / elapsed) if elapsed else None + out["engine_tok_s"] = r(m["committed"] / elapsed) if elapsed else None + if m["batch_steps"] and elapsed: + out["ms_per_batch_step"] = r(1000.0 * elapsed / m["batch_steps"]) + if "NA" not in s: + out["buckets_per_step"] = r(s["buckets"] / s["calls"]) if s["calls"] else None + out["running_bucket_size"] = r(s["chosen"] / s["calls"]) if s["calls"] else None + out["offered_per_step"] = r(s["offered"] / s["calls"]) if s["calls"] else None + out["shattered_frac"] = r(s["shattered"] / s["calls"]) if s["calls"] else None + print(json.dumps(out)) + +main() +PYEOF + +run_arm() { + local arm="$1"; shift + local slog="$OUT/server.$arm.log" + say "arm $arm — starting server (${*:-no flags})" + # shellcheck disable=SC2086 + env RUST_LOG=info ARC_MTP_LOG_ACCEPTANCE=1 "$@" \ + "$BIN" serve -p "$PORT" -m "$MODEL" -a "$ARCH" --from-uqff "$UQFF" \ + --max-seqs "$MAX_SEQS" --mtp-depth "$MTP_DEPTH" --prefix-cache-n 0 \ + --paged-attn off --seed "$SEED" \ + < /dev/null > "$slog" 2>&1 & + local spid=$! + local up=0 + for _ in $(seq 1 450); do + curl -s -m 5 "http://127.0.0.1:$PORT/v1/models" >/dev/null 2>&1 && { up=1; break; } + kill -0 "$spid" 2>/dev/null || break + sleep 2 + done + [ "$up" = "1" ] || { say "arm $arm — server never came up"; tail -25 "$slog" | tee -a "$STATUS"; return 1; } + + local logged + logged=$(grep -m1 -o 'git revision: [0-9a-f]*' "$slog" | awk '{print $3}') + if [ -n "$SHA" ] && [ "$logged" != "$SHA" ]; then + say "arm $arm — PROVENANCE FAILED: server $logged != worktree $SHA" + kill -TERM "$spid" 2>/dev/null || true + return 1 + fi + say "arm $arm — server up, revision $logged" + + local regime k spread_csv tag + for regime in $REGIMES; do + if [ "$regime" = "uniform" ]; then + spread_csv="$UNIFORM_WORDS" + else + spread_csv=$(echo "$WORD_SPREAD" | tr ' ' ',') + fi + for k in $BATCHES; do + tag="$arm.$regime.k$k" + say "arm $arm — cell $regime B=$k (warmup ${WARMUP}s, steady ${STEADY}s, words $spread_csv)" + python3 "$OUT/drive.py" --port "$PORT" --k "$k" --warmup "$WARMUP" \ + --steady "$STEADY" --spread "$spread_csv" --max-tokens "$MAX_TOKENS" \ + --seed "$SEED" > "$OUT/$tag.drive.json" 2>"$OUT/$tag.drive.err" \ + || { say "cell $tag — driver failed"; return 1; } + local t0 t1 el tok + t0=$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['t_start'])" "$OUT/$tag.drive.json") + t1=$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['t_end'])" "$OUT/$tag.drive.json") + el=$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['elapsed_s'])" "$OUT/$tag.drive.json") + tok=$(python3 -c "import json,sys;print(json.load(open(sys.argv[1]))['driver_tok_s'])" "$OUT/$tag.drive.json") + python3 "$OUT/fence.py" "$slog" "$t0" "$t1" "$el" "$tok" > "$OUT/$tag.cell.json" \ + || { say "cell $tag — fence extraction failed"; return 1; } + say "$tag -> $(cat "$OUT/$tag.cell.json")" + sleep 5 + done + done + + kill -TERM "$spid" 2>/dev/null || true + for _ in $(seq 1 30); do kill -0 "$spid" 2>/dev/null || break; sleep 2; done + kill -KILL "$spid" 2>/dev/null || true + sleep 5 + say "arm $arm — done" + return 0 +} + +run_arm OFF || die "OFF arm failed" +run_arm ON ARC_V4_XS_PER_SEQ=1 ARC_MTP_PER_SEQ_KV=1 || die "ON arm failed" +# The third arm prices the window pin, which is OPT-IN (default OFF) — so every +# arm above is unpinned and this one is the treatment. Same flags as ON, pin +# enabled, so ON vs ON_PINNED isolates one change: whether the compressor's raw +# window is reallocated every decode step on all 41 compressed layers, or held +# at its documented bound. +# +# This arm IS the experiment named in `xs_pin_window_enabled_from`'s FLIP +# CONDITION. If ON_PINNED wins at uniform B=32 with identical generated tokens, +# `ARC_V4_XS_PIN_WINDOW` defaults to ON and that doc comment is rewritten in the +# same change that records the number. +run_arm ON_PINNED ARC_V4_XS_PER_SEQ=1 ARC_MTP_PER_SEQ_KV=1 ARC_V4_XS_PIN_WINDOW=1 \ + || die "ON_PINNED arm failed" + +cat > "$OUT/report.py" <<'PYEOF' +import json, os, sys +out, batches, regimes = sys.argv[1], sys.argv[2].split(), sys.argv[3].split() +ARMS = ("OFF", "ON", "ON_PINNED") +def cell(arm, regime, k): + p = os.path.join(out, f"{arm}.{regime}.k{k}.cell.json") + try: + return json.load(open(p)) + except Exception as e: + return {"NA": str(e)} +def g(c, k): + v = c.get(k) + return "NA" if v is None else v +print("\n=== LADDER: aggregate throughput and its decomposition ===") +print("aggregate_tok/s is wall-clock and is THE number. tok/batch_step x") +print("batch_steps/s reconstructs it: when aggregate does not move, those two") +print("say which factor ate it. run_bucket is the width the engine ACTUALLY ran") +print("— compare it to offered before reading any cell as a batch result.") +hdr = (f"{'regime':<8} {'B':>4} {'arm':<12} {'aggregate_tok/s':>15} {'tok/step':>9} " + f"{'tok/bstep':>10} {'bsteps/s':>9} {'ms/bstep':>9} " + f"{'buckets/step':>13} {'run_bucket':>11} {'offered':>8}") +print(hdr); print("-" * len(hdr)) +for regime in regimes: + for k in batches: + for arm in ARMS: + c = cell(arm, regime, k) + print(f"{regime:<8} {k:>4} {arm:<12} {str(g(c,'aggregate_tok_s')):>15} " + f"{str(g(c,'tok_per_step')):>9} {str(g(c,'tok_per_batch_step')):>10} " + f"{str(g(c,'batch_steps_per_s')):>9} {str(g(c,'ms_per_batch_step')):>9} " + f"{str(g(c,'buckets_per_step')):>13} {str(g(c,'running_bucket_size')):>11} " + f"{str(g(c,'offered_per_step')):>8}") +print("\n=== ON vs OFF, per regime ===") +for regime in regimes: + print(f"\n-- {regime} --") + for k in batches: + o, n = cell("OFF", regime, k), cell("ON", regime, k) + def ratio(key): + a, b = o.get(key), n.get(key) + if not a or not b: + return "NA" + return f"{b/a:.3f}x ({a} -> {b})" + print(f" B={k:<4} aggregate {ratio('aggregate_tok_s')}") + print(f" tok/step {ratio('tok_per_step')}") + print(f" tok/bstep {ratio('tok_per_batch_step')}") + print(f" bsteps/s {ratio('batch_steps_per_s')}") + print(f" run_bucket {ratio('running_bucket_size')}") +print("\n=== THE WINDOW PIN, isolated (ON = control, ON_PINNED = treatment) ===") +print("Same per-seq flags in both; the only difference is whether the") +print("compressor's raw window is reallocated every decode step on 41 layers.") +print("Ratios are treatment/control, so >1.000x means the pin helped.") +print("This is the number in ARC_V4_XS_PIN_WINDOW's FLIP CONDITION: if the") +print("uniform B=32 row is >1.000x with identical output, the default flips ON.") +for regime in regimes: + for k in batches: + # `u` is the unpinned control, `p_` the pinned treatment; the pin is + # opt-in, so the DEFAULT arm (ON) is the control. This pairing was the + # other way round while the pin defaulted on — the ratio is + # treatment/control either way, and swapping the arms without swapping + # this would have silently inverted every number below. + u, p_ = cell("ON", regime, k), cell("ON_PINNED", regime, k) + def r2(key): + a, b = u.get(key), p_.get(key) + if not a or not b: + return "NA" + return f"{b/a:.3f}x ({a} -> {b})" + print(f" {regime:<8} B={k:<4} aggregate {r2('aggregate_tok_s')} " + f"ms/bstep {r2('ms_per_batch_step')}") + +print("\n=== THE BUCKETING LAW, checked against this run ===") +print("The scheduler A/B measured `running bucket = B / distinct lengths`.") +print("uniform has 1 distinct length, spread has 8, so the prediction is:") +for regime in regimes: + dist = 1 if regime == "uniform" else 8 + for k in batches: + pred = max(1.0, int(k) / dist) + for arm in ARMS: + c = cell(arm, regime, k) + got = c.get("running_bucket_size") + if got is None: + continue + print(f" {regime:<8} B={k:<4} {arm:<12} predicted {pred:>6.2f} measured {got:>7.3f}") +print("\n⚠️ Read every spread cell against its run_bucket. If run_bucket is ~1,") +print("that cell is a measurement of the scheduler serialising the batch, not") +print("of per-sequence KV advance, and its aggregate cannot be attributed here.") +PYEOF + +{ + echo "=== ref ==="; echo "$SHA" + echo "=== prompts ===" + echo " spread : $WORD_SPREAD (8 distinct lengths)" + echo " uniform: $UNIFORM_WORDS (the spread's mean AND its first element, so B=1 matches)" + echo + echo "=== ENGAGEMENT ===" + for m in "per-sequence KV advance is ON" "Ragged batch admission is ON"; do + on=$(grep -c "$m" "$OUT/server.ON.log" 2>/dev/null || echo 0) + off=$(grep -c "$m" "$OUT/server.OFF.log" 2>/dev/null || echo 0) + echo " '$m': ON=$on (want >=1) OFF=$off (want 0)" + done + echo " 'cannot honour it' in ON: $(grep -c 'cannot honour it' "$OUT/server.ON.log" 2>/dev/null || echo 0) (want 0)" + echo " per_seq_steps final, ON: $(grep -o 'per_seq_steps=[0-9]*' "$OUT/server.ON.log" 2>/dev/null | tail -1) (want >0)" + echo " per_seq_steps final, OFF: $(grep -o 'per_seq_steps=[0-9]*' "$OUT/server.OFF.log" 2>/dev/null | tail -1) (want =0)" + echo " 'xs rolling cache' errors ON=$(grep -c 'xs rolling cache' "$OUT/server.ON.log" 2>/dev/null || echo 0) OFF=$(grep -c 'xs rolling cache' "$OUT/server.OFF.log" 2>/dev/null || echo 0) (want 0/0)" + echo + echo " -- window pin engagement (a wrong flag name would make ON and ON_PINNED" + echo " the same build and report a free pin from two identical arms) --" + for a in OFF ON ON_PINNED; do + echo " $a: $(grep -m1 -o 'xs rolling window is [A-Z]*' "$OUT/server.$a.log" 2>/dev/null || echo 'NOT LOGGED')" + done + # ⚠️ Read from the LOG, never from the binary. BOTH mode strings are compiled + # in unconditionally, so `strings | grep RESIZING` would succeed in + # every arm and prove nothing. Only the runtime line says which branch a + # process actually took. `RUST_LOG=info` is set once in `run_arm`, shared by + # every arm, so no arm can have this line filtered out while another keeps it. + # + # ⚠️ The default direction matters here, and it INVERTED when the pin became + # opt-in: unset now means RESIZING (`xs_rolling.rs`, + # `xs_pin_window_enabled_from` is `== Some("1")`). So the CONTROL is the + # default and the TREATMENT is the one carrying the flag. A typo in the flag + # name now yields a treatment that silently ran the control — still a + # perfectly clean-looking 1.000x, which is why this check tests BOTH + # directions rather than inferring one from the other. + pin_off=$(grep -c 'xs rolling window is RESIZING' "$OUT/server.ON.log" 2>/dev/null || echo 0) + pin_on=$(grep -c 'xs rolling window is PINNED' "$OUT/server.ON_PINNED.log" 2>/dev/null || echo 0) + treat_leaked=$(grep -c 'xs rolling window is RESIZING' "$OUT/server.ON_PINNED.log" 2>/dev/null || echo 0) + if [ "$pin_off" -ge 1 ] && [ "$pin_on" -ge 1 ] && [ "$treat_leaked" -eq 0 ]; then + echo " => VALID: ON took the resizing path, ON_PINNED took the pinned path." + echo " The treatment reached the pinned branch; it did not silently run the control." + else + echo " => ⚠️ PIN A/B IS VOID: ON resizing=$pin_off, ON_PINNED pinned=$pin_on," + echo " treatment-leaked-to-resizing=$treat_leaked." + echo " Any ON-vs-ON_PINNED ratio below compares one thing with itself. Do not read it." + fi + python3 "$OUT/report.py" "$OUT" "$BATCHES" "$REGIMES" +} | tee "$OUT/summary.txt" + +say "DONE — $OUT/summary.txt" diff --git a/arc-tools/batch_sweep.py b/arc-tools/batch_sweep.py new file mode 100644 index 000000000..d27065b7b --- /dev/null +++ b/arc-tools/batch_sweep.py @@ -0,0 +1,256 @@ +"""Batch sweep on landed master: per-request vs aggregate decode throughput. + +For each batch size B, fire B CONCURRENT streaming requests and time the +inter-token intervals of each. Two independent aggregate anchors that must +agree, because the interesting number here is capacity and a capacity claim +built on one instrument is not a claim: + + A) derived = per_request_tok_s * B (what the brief asked for) + B) measured = tokens delivered inside the all-B-active window / that window + +(A) assumes every stream runs at the median rate; (B) counts actual deltas in +the steady-state window where all B streams are decoding. If they diverge the +batch never reached steady state and the row is not a throughput measurement. + +Per-request rate is a median of inter-token INTERVALS, so prefill is excluded +and one scheduler hiccup cannot set the headline. + +A row where any request returns zero tokens is a FAILED row, printed as such, +never averaged into the table -- the whole point of the bench gate. + +Exit: 0 sweep complete | 1 a row failed | 2 environment failure. +""" +import json +import statistics +import sys +import subprocess +import threading +import time +import urllib.request + +BASE = sys.argv[1] if len(sys.argv) > 1 else "http://127.0.0.1:1234" +NTOK = int(sys.argv[2]) if len(sys.argv) > 2 else 64 +BATCHES = [int(x) for x in (sys.argv[3].split(",") if len(sys.argv) > 3 else ["1", "8", "32", "64", "128"])] + +TOPICS = [ + "oceans", "mountains", "deserts", "forests", "rivers", "glaciers", "volcanoes", + "coral reefs", "savannas", "tundra", "caves", "islands", "wetlands", "canyons", + "prairies", "fjords", +] + + +def one_stream(idx, ntok, out): + """Stream one request; record a timestamp per delivered token.""" + topic = TOPICS[idx % len(TOPICS)] + prompt = ("Write %d short factual sentences about %s. Sentence set %d. " + "One sentence per line, no preamble." % (max(6, ntok // 6), topic, idx)) + body = json.dumps({ + "model": "default", + "messages": [{"role": "user", "content": prompt}], + "max_tokens": ntok, "temperature": 0.0, "stream": True, + }).encode() + req = urllib.request.Request(BASE + "/v1/chat/completions", data=body, + headers={"Content-Type": "application/json"}) + stamps, chars, err = [], 0, None + t0 = time.perf_counter() + try: + with urllib.request.urlopen(req, timeout=1800) as r: + for raw in r: + line = raw.decode("utf-8", "replace").strip() + if not line.startswith("data:"): + continue + payload = line[5:].strip() + if payload == "[DONE]": + break + try: + d = json.loads(payload) + except Exception: + continue + if "error" in d: + err = str(d["error"])[:120] + break + for c in d.get("choices") or []: + piece = ((c.get("delta") or {}).get("content") or "") + if piece: + stamps.append(time.perf_counter()) + chars += len(piece) + except Exception as e: + err = "%s: %s" % (type(e).__name__, str(e)[:100]) + out[idx] = {"t0": t0, "stamps": stamps, "chars": chars, "err": err} + + +SMI = ["nvidia-smi", "--query-gpu=utilization.gpu,utilization.memory,memory.used,memory.free", + "--format=csv,noheader,nounits"] + + +def sample_gpu(stop, sink): + """Poll the GPU WHILE the batch decodes. + + `utilization.memory` is the memory-controller busy percentage -- the + diagnostic that says whether batching is amortising the weight reads. A + sampler that starts after the run writes an empty file and looks green, so + this one is started before the requests and stamped per sample. + """ + while not stop.is_set(): + try: + out = subprocess.run(SMI, capture_output=True, text=True, timeout=5).stdout.strip() + for line in out.splitlines(): + g, m, used, free = [x.strip() for x in line.split(",")] + sink.append((time.perf_counter(), int(g), int(m), int(used), int(free))) + except Exception: + pass + time.sleep(0.25) + + +def run_batch(b, ntok): + out = {} + samples = [] + stop = threading.Event() + sampler = threading.Thread(target=sample_gpu, args=(stop, samples), daemon=True) + sampler.start() + threads = [threading.Thread(target=one_stream, args=(i, ntok, out)) for i in range(b)] + w0 = time.perf_counter() + for t in threads: + t.start() + for t in threads: + t.join() + w1 = time.perf_counter() + stop.set() + sampler.join(timeout=5) + return out, w1 - w0, samples + + +def grade(b, out, wall, samples): + rows = [out[i] for i in sorted(out)] + zero = [i for i, r in enumerate(rows) if len(r["stamps"]) == 0] + errs = [(i, r["err"]) for i, r in enumerate(rows) if r["err"]] + if len(rows) != b: + return {"B": b, "fail": "instrument: %d rows for B=%d" % (len(rows), b)} + if zero: + return {"B": b, "fail": "%d of %d requests returned ZERO tokens%s" + % (len(zero), b, (" (first err: %s)" % errs[0][1]) if errs else "")} + + # per-request: median over requests of that request's median gap + per_req_rates = [] + for r in rows: + s = r["stamps"] + if len(s) < 4: + return {"B": b, "fail": "a request produced only %d tokens" % len(s)} + gaps = [s[i] - s[i - 1] for i in range(1, len(s))] + per_req_rates.append(1.0 / statistics.median(gaps)) + per_req = statistics.median(per_req_rates) + total_tok_all = sum(len(r["stamps"]) for r in rows) + + # COHORT throughput: every token the cohort delivered, over the wall time + # the cohort spent delivering them. This is the capacity number an operator + # bills against, and unlike a strict all-active window it is defined even + # when admission is staggered -- which it is at large B, because the + # scheduler admits in waves. + first = min(r["stamps"][0] for r in rows) + last = max(r["stamps"][-1] for r in rows) + span = last - first + if span <= 0: + return {"B": b, "fail": "cohort delivered no tokens over time"} + measured = total_tok_all / span + + # Achieved concurrency: mean active streams across the span. If this is far + # below B the row is NOT a measurement of B -- the scheduler never ran B at + # once, and reporting it as B would overstate what the box did. + edges = [] + for r in rows: + edges.append((r["stamps"][0], 1)) + edges.append((r["stamps"][-1], -1)) + edges.sort() + active, prev_t, area = 0, first, 0.0 + for t, d in edges: + area += active * (t - prev_t) + active += d + prev_t = t + concurrency = area / span if span > 0 else 0.0 + + # strict all-active window, when one exists (secondary anchor) + lo = max(r["stamps"][0] for r in rows) + hi = min(r["stamps"][-1] for r in rows) + if hi > lo: + inwin = sum(sum(1 for t in r["stamps"] if lo < t <= hi) for r in rows) + steady = inwin / (hi - lo) + else: + steady = -1.0 + + derived = per_req * b + disagree = abs(measured - derived) / max(measured, derived) + # attribute GPU samples to the cohort span + win = [s for s in samples if first <= s[0] <= last] + if not win: + win = samples + memu = sorted(s[2] for s in win) + gpuu = sorted(s[1] for s in win) + used = sorted(s[3] for s in win) + free = sorted(s[4] for s in win) + n_s = len(win) + ttfts = sorted(r["stamps"][0] - r["t0"] for r in rows) + return { + "B": b, "fail": None, + "per_req": per_req, + "ms_per_tok": 1000.0 / per_req, + "agg_derived": derived, + "agg_measured": measured, + "disagree": disagree, + "total_tok": total_tok_all, + "concurrency": concurrency, + "steady": steady, + "window_s": span, + "wall_s": wall, + "ttft_med": ttfts[len(ttfts) // 2], + "ttft_max": ttfts[-1], + "mem_util_med": (memu[len(memu) // 2] if memu else -1), + "mem_util_max": (memu[-1] if memu else -1), + "gpu_util_med": (gpuu[len(gpuu) // 2] if gpuu else -1), + "vram_used_max": (used[-1] if used else -1), + "vram_free_min": (free[0] if free else -1), + "n_samples": n_s, + } + + +print("BASE=%s NTOK=%d batches=%s" % (BASE, NTOK, BATCHES), flush=True) +results, failed = [], False +for b in BATCHES: + print("--- B=%d : firing %d concurrent streams ---" % (b, b), flush=True) + out, wall, samples = run_batch(b, NTOK) + r = grade(b, out, wall, samples) + results.append(r) + if r["fail"]: + failed = True + print(" B=%-4d FAILED ROW: %s" % (b, r["fail"]), flush=True) + else: + print(" B=%-4d per-req %7.2f | agg %8.2f | conc %6.1f/%d | MEM-CTRL %d%% (max %d%%) | " + "GPU %d%% | VRAMfree %d MiB | tok %d | samples %d" + % (b, r["per_req"], r["agg_measured"], r["concurrency"], b, + r["mem_util_med"], r["mem_util_max"], r["gpu_util_med"], + r["vram_free_min"], r["total_tok"], r["n_samples"]), flush=True) + time.sleep(3) + +print("") +print("%-6s %-13s %-11s %-13s %-13s %-8s %-9s %-9s %-10s" % + ("B", "per-req t/s", "ms/tok", "agg derived", "agg measured", "spread", + "MEMCTRL%", "conc", "VRAMfree")) +for r in results: + if r["fail"]: + print("%-6d %s" % (r["B"], "FAILED: " + r["fail"])) + else: + print("%-6d %-13.2f %-11.2f %-13.2f %-13.2f %-8s %-9s %-9s %-10d" + % (r["B"], r["per_req"], r["ms_per_tok"], r["agg_derived"], r["agg_measured"], + "%.1f%%" % (r["disagree"] * 100), + "%d/%d" % (r["mem_util_med"], r["mem_util_max"]), + "%.0f" % r["concurrency"], r["vram_free_min"])) + +ok = [r for r in results if not r["fail"]] +if len(ok) >= 2: + base = ok[0]["agg_measured"] + print("") + print("aggregate scaling vs B=%d:" % ok[0]["B"]) + for r in ok: + print(" B=%-5d %8.2f tok/s x%.2f (per-request retained: %.0f%%)" + % (r["B"], r["agg_measured"], r["agg_measured"] / base, + 100.0 * r["per_req"] / ok[0]["per_req"])) +sys.exit(1 if failed else 0) diff --git a/arc-tools/fp8_gemm_sass_check.sh b/arc-tools/fp8_gemm_sass_check.sh new file mode 100755 index 000000000..48634f447 --- /dev/null +++ b/arc-tools/fp8_gemm_sass_check.sh @@ -0,0 +1,209 @@ +#!/usr/bin/env bash +# Parent system: ArcKernels + ArcGate. +# +# FREE SASS PROBE for the blockwise-FP8 dense-GEMM lane. NO GPU REQUIRED: +# nvcc + cuobjdump only, so it runs on the GPU-less nvcc CI lane +# (.github/workflows/cuda_compile_check.yaml) and on any dev box with the +# toolkit. +# +# WHAT IT SETTLES +# --------------- +# 1. `fp8_matmul_tiled` (blockwise_fp8_gemm.cu) — HOW ptxas actually loads the +# shared tiles in the inner product `acc += s_input[ty][k] * s_weight[tx][k]`. +# The source-derived cost model for this kernel (shared traffic 4.13x over +# supply x a 4-way bank conflict from the old +4 tile padding) is a claim +# about SCALAR 32-bit LDS. If ptxas had already vectorized those loads to +# LDS.128, the conflict arithmetic would be about instructions that do not +# exist and the model's 7% agreement would be a coincidence. This probe +# prints the per-kernel LDS width census so that question is settled by the +# compiled artifact, not by reading C++. +# NOTE the padding interaction: the +1 pad (stride 33 floats, PR #215) +# that kills the 4-way bank conflict ALSO forecloses LDS.128 on tile rows +# (rows are only 4-byte aligned at stride 33). Scalar LDS here is +# therefore the EXPECTED outcome on current source, not a defect — the +# probe exists so a future retune that trades padding for vector width +# does it with the census in front of them. +# 2. `blockwise_fp8_gemm_wmma.cu` — whether the "tensor-core" GEMM actually +# contains tensor-core instructions (HMMA/HGMMA). That kernel has NEVER +# executed on any GPU; if ptxas lowered it without a single matrix-multiply +# instruction, every derivation in its header is void. HARD FAIL if absent. +# +# TEETH (a parser that matches nothing must not pass) +# --------------------------------------------------- +# * Kernel-count assertion: the tiled TU must yield exactly 4 +# `fp8_matmul_tiled` kernels (2 dtypes x 2 POW2_SCALE instantiations — +# see the launcher in blockwise_fp8_gemm.cu). A count drift means an +# instantiation was dropped and the census is about a kernel not in the +# build. +# * Every tiled kernel must show at least one LDS — a tiled shared-memory GEMM +# with zero shared loads means the SASS extractor lost the function body. +# +# Usage: arc-tools/fp8_gemm_sass_check.sh [sm_90] +set -euo pipefail + +ARCH="${1:-sm_90}" +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +KDIR="${REPO_ROOT}/mistralrs-quant/kernels/blockwise_fp8" +OUT="$(mktemp -d)" +trap 'rm -rf "${OUT}"' EXIT + +# Mirror mistralrs-quant/build.rs exactly: a different flag set is a different +# instruction selection, and therefore a probe of a kernel we do not ship. +NVCC_ARGS=( + -std=c++17 + -O3 + -U__CUDA_NO_HALF_OPERATORS__ + -U__CUDA_NO_HALF_CONVERSIONS__ + -U__CUDA_NO_HALF2_OPERATORS__ + -U__CUDA_NO_BFLOAT16_CONVERSIONS__ + --expt-relaxed-constexpr + --expt-extended-lambda + --use_fast_math +) + +echo "== blockwise-FP8 GEMM SASS probe ==" +nvcc --version | tail -1 +echo "arch: ${ARCH}" +echo + +# --------------------------------------------------------------------------- +# Per-function SASS extraction. +# +# PORTABILITY LESSON (this script's first CI run went red on itself): GNU awk +# does not honor \< \> word-boundary escapes in ERE — it warns "escape +# sequence treated as plain character" and matches nothing — so the original +# awk body-splitter returned zero LDS for every function and the fail-on-zero +# guard below correctly refused to trust it. This revision does no regex body +# splitting at all: cuobjdump lists the mangled kernel names, and +# `cuobjdump -sass -fun ` extracts exactly one function's SASS. All +# counting is grep -cE with POSIX classes and bracket expressions only — +# no \< \> anywhere. +# --------------------------------------------------------------------------- +fn_names() { # $1 = cubin, $2 = fixed-string name filter + cuobjdump -sass "$1" | grep 'Function : ' | awk '{print $NF}' | grep -F "$2" || true +} +count() { # $1 = file, $2 = POSIX ERE + grep -cE "$2" "$1" || true +} +# SASS instruction lines carry a /**/ comment; body-vs-empty is +# judged on those, so "extractor returned nothing" and "zero LDS" cannot be +# confused with each other. +RE_INSN='/[*][0-9a-fA-F]+[*]/' +# LDSM ('M' after LDS) must not count; the char after LDS is '.' or space. +RE_LDS_ALL='[[:space:]]LDS[[:space:].]' +RE_LDS_128='[[:space:]]LDS[.](U[.])?128[[:space:]]' +RE_LDS_64='[[:space:]]LDS[.](U[.])?64[[:space:]]' +RE_MMA='[[:space:]](HMMA|HGMMA)[[:space:].]' + +# --------------------------------------------------------------------------- +# 1. The scalar tiled GEMM: LDS width census. +# --------------------------------------------------------------------------- +echo "--- blockwise_fp8_gemm.cu : fp8_matmul_tiled ---" +nvcc -cubin -arch="${ARCH}" "${NVCC_ARGS[@]}" \ + -o "${OUT}/gemm.cubin" "${KDIR}/blockwise_fp8_gemm.cu" + +if command -v c++filt >/dev/null 2>&1; then + cuobjdump -sass "${OUT}/gemm.cubin" | grep 'Function : ' | c++filt | grep fp8_matmul_tiled || true +fi + +fail=0 +extract_ok=1 +n_tiled=0 +sum128=0 +names="$(fn_names "${OUT}/gemm.cubin" fp8_matmul_tiled)" +if [ -z "${names}" ]; then + echo "FAIL[extractor]: no fp8_matmul_tiled function names in the cubin's" + echo " SASS listing — the name pass found nothing, so there is no census." + fail=1 + extract_ok=0 +fi + +echo "kernel insns LDS(all) LDS.128 LDS.64 scalar" +for name in ${names}; do + n_tiled=$((n_tiled + 1)) + sfile="${OUT}/tiled_${n_tiled}.sass" + cuobjdump -sass -fun "${name}" "${OUT}/gemm.cubin" > "${sfile}" 2>/dev/null || true + insns="$(count "${sfile}" "${RE_INSN}")" + lds_all="$(count "${sfile}" "${RE_LDS_ALL}")" + w128="$(count "${sfile}" "${RE_LDS_128}")" + w64="$(count "${sfile}" "${RE_LDS_64}")" + scalar=$((lds_all - w128 - w64)) + sum128=$((sum128 + w128)) + printf '%-48s %6s %8s %7s %6s %6s\n' "${name:0:48}" "${insns}" "${lds_all}" "${w128}" "${w64}" "${scalar}" + if [ "${insns}" -eq 0 ]; then + echo "FAIL[extractor]: ${name}: 'cuobjdump -sass -fun' returned no" + echo " instruction lines. This is an EXTRACTION failure, not a zero-LDS" + echo " kernel — this function's census row is unverified." + fail=1 + extract_ok=0 + elif [ "${lds_all}" -eq 0 ]; then + echo "FAIL[zero-LDS]: ${name}: body extracted (${insns} instructions) but" + echo " it contains no shared-memory loads at all. A shared-tile GEMM" + echo " without LDS means the wrong function was extracted or the kernel" + echo " no longer stages tiles — either way the conflict model has no" + echo " subject here." + fail=1 + fi +done + +if [ "${n_tiled}" -ne 4 ]; then + echo "FAIL: expected 4 fp8_matmul_tiled kernels (2 dtypes x 2 POW2_SCALE" + echo " instantiations), found ${n_tiled}. An instantiation was dropped;" + echo " the census is about kernels that are not in the build." + fail=1 +fi + +# TU-level cross-check: computed on the WHOLE cubin dump, independent of the +# per-function extractor, so it survives an extractor failure and gives the +# reader a signal either way. Informational — the per-function census is the +# only thing the VERDICT may rest on (the TU includes the GEMV/MoE kernels). +cuobjdump -sass "${OUT}/gemm.cubin" > "${OUT}/gemm.sass" 2>/dev/null || true +echo "TU-level cross-check (whole cubin, informational):" \ + "LDS(all)=$(count "${OUT}/gemm.sass" "${RE_LDS_ALL}")" \ + "LDS.128=$(count "${OUT}/gemm.sass" "${RE_LDS_128}")" \ + "LDS.64=$(count "${OUT}/gemm.sass" "${RE_LDS_64}")" + +# The VERDICT prints ONLY on a complete, non-empty per-function census. +# Printing it off TU-level counts (or off an empty extraction) would be a +# lying instrument: confident prose about numbers that were never measured. +if [ "${extract_ok}" -eq 1 ] && [ "${n_tiled}" -eq 4 ] && [ "${fail}" -eq 0 ]; then + if [ "${sum128}" -eq 0 ]; then + echo "VERDICT[tiled]: scalar LDS only, no LDS.128 — the shared loads are" + echo " 32-bit, so the bank-conflict arithmetic (stride mod 32) is about" + echo " instructions that really exist, and the +1-pad fix is live." + else + echo "VERDICT[tiled]: ptxas emitted LDS.128 — the scalar-LDS conflict model" + echo " does NOT describe this build; re-derive before quoting it." + fi +else + echo "VERDICT[tiled]: NOT ESTABLISHED — the per-function census is empty or" + echo " partial (see FAILs above). Do not quote any LDS conclusion from this" + echo " run; the TU-level line above is a cross-check, not a verdict." +fi +echo + +# --------------------------------------------------------------------------- +# 2. The WMMA GEMM: tensor-core instructions must exist. +# --------------------------------------------------------------------------- +echo "--- blockwise_fp8_gemm_wmma.cu : HMMA presence ---" +nvcc -cubin -arch="${ARCH}" "${NVCC_ARGS[@]}" \ + -o "${OUT}/wmma.cubin" "${KDIR}/blockwise_fp8_gemm_wmma.cu" + +cuobjdump -sass "${OUT}/wmma.cubin" > "${OUT}/wmma.sass" 2>/dev/null || true +hmma="$(count "${OUT}/wmma.sass" "${RE_MMA}")" +echo "HMMA/HGMMA instructions in TU: ${hmma}" +if [ "${hmma}" -eq 0 ]; then + echo "FAIL: the tensor-core GEMM contains no HMMA/HGMMA. It has never run on" + echo " hardware; if it also has no tensor-core instructions, every" + echo " derivation in its header is void and the three-arm sweep would" + echo " be comparing cuBLASLt against two scalar kernels." + fail=1 +fi +echo + +if [ "${fail}" -ne 0 ]; then + echo "RESULT: FAILED — see above." + exit 1 +fi +echo "RESULT: census complete (${ARCH}); 4 tiled kernels present (per-function" +echo " extraction verified non-empty), WMMA TU carries ${hmma} HMMA/HGMMA." diff --git a/arc-tools/fp8_threshold_sweep.sh b/arc-tools/fp8_threshold_sweep.sh new file mode 100755 index 000000000..05883dd13 --- /dev/null +++ b/arc-tools/fp8_threshold_sweep.sh @@ -0,0 +1,412 @@ +#!/bin/bash +# Parent system: ArcKernels + ArcQuant/Legacy (blockwise FP8). +# +# THREE-ARM FP8 DENSE-GEMM THRESHOLD SWEEP — the measurement that is allowed +# to move `ARC_FP8_CUBLAS_MIN_M` off 512. +# +# ───────────────────────────────────────────────────────────────────────────── +# WHY THREE ARMS, AND WHY A TWO-ARM RESULT MUST NOT BE MERGED +# +# `arc_fp8_cublas_min_m`'s own doc (blockwise_fp8/mod.rs) and the WMMA +# kernel's doc (blockwise_fp8/ops.rs, `fp8_wmma_enabled`) both record the +# standing instruction: the only sweep ever taken compared (a) dequantize + +# cuBLASLt against (c) the scalar tiled kernel, BEFORE #200 landed the +# (b) tensor-core WMMA GEMM — and before the rank-3 flatten, so arm (a) was a +# B-way batched GEMV wearing a GEMM's name. Freezing a threshold from that +# sweep would leave (b) dead code at every M that matters, silently. This +# script is the three-arm re-run those docs demand: +# +# (a) dequantize_w() + cuBLASLt ARC_FP8_CUBLAS_MIN_M=5 +# (b) native FP8 tensor-core WMMA ARC_FP8_CUBLAS_MIN_M=1000000, ARC_NO_FP8_WMMA=0 +# (c) native scalar fp8_matmul_tiled ARC_FP8_CUBLAS_MIN_M=1000000, ARC_NO_FP8_WMMA=1 +# +# ONE binary. Every leg differs only in environment; a leg that needed a +# rebuild would be measuring a different artifact. The env gates are latched +# per process (OnceLock/LazyLock), so every leg gets a FRESH server process — +# re-exporting into a running server changes nothing and would silently +# measure the previous leg's setting. +# +# ───────────────────────────────────────────────────────────────────────────── +# THE LADDER: M VARIES BY OFFERED CONCURRENCY, NOT ONLY BY THRESHOLD +# +# The deliverable is "the first M at which (a) beats min((b), (c))". A grid +# that sweeps only ARC_FP8_CUBLAS_MIN_M at one fixed batch cannot name that M: +# every leg with MIN_M <= batch routes the same decode shape to the same +# kernel and the legs replicate each other. So the threshold candidates +# {5, 8, 64, 256, 512} are swept as the OFFERED DECODE BATCH — M concurrent +# greedy completions, T=1 per step, so the dense projections see m_rows = M — +# and each ladder point runs all three arms. 15 legs. +# +# ───────────────────────────────────────────────────────────────────────────── +# WHAT MAKES A LEG COUNT (all three, or the leg is VOID) +# +# 1. ENGAGEMENT, asserted per leg. Every leg runs with ARC_LOG_FP8_DISPATCH=1; +# the dispatch sites print `[arc-fp8-dispatch] path=` once per path +# per process. The leg's expected kernel must be PRESENT in the server log +# and both rival kernels ABSENT. A leg that cannot prove which kernel +# served it is measuring an assumption (`ARC_NO_DEDICATED_DECODE`'s two +# void A/Bs are the record of what that costs). In particular: if the WMMA +# arm's log shows `path=tiled`, WMMA eligibility silently failed and the +# arm is the control wearing the treatment's name — VOID, loudly. +# 2. A MINIMUM-TOKEN FLOOR. Rates on fewer than MIN_TOKENS_FLOOR generated +# tokens are warm-up noise. The rate divides by SUMMED +# `usage.completion_tokens` — the tokens the server says it generated — +# NEVER by requested max_tokens (a leg that stalls and returns short would +# otherwise report a healthy rate for work it did not do). +# 3. A TOKEN CANARY vs the baseline leg (arm (c) at M=5 — the shipped scalar +# kernel, nearest to b=1). Greedy, fixed seed, fixed prompt. The arms are +# DIFFERENT NUMERICS by design (f32 scalar vs tensor-core promotion vs +# dequant-BF16 GEMM), so divergence at some index is expected and REPORTED, +# not failed; what hard-fails a leg is a degenerate stream (< 5 tokens or +# no finish_reason) — a failure to generate, which no rate can offset. +# +# The final table names, per ladder M, the three clean rates and min((b),(c)), +# and then the FIRST M at which (a) beats that min — the value +# ARC_FP8_CUBLAS_MIN_M should default to. If any arm at that M is void, the +# crossover is NOT ESTABLISHED and the default stays 512. +# +# USAGE (on the GPU box; do not run without a GPU — D14) +# setsid nohup bash arc-tools/fp8_threshold_sweep.sh /root/logs/fp8_sweep.log 2>&1 & +# cat /root/logs/fp8sweep/summary.txt +# +set -u + +ARCENV="${ARCENV:-/root/arcenv.sh}" +WT="${WT:-/root/arc}" +BIN="${BIN:-$WT/target/release/mistralrs}" +MODEL="${MODEL:-/root/models/v4-src}" +UQFF="${UQFF:-}" # optional: --from-uqff when set +ARCH="${ARCH:-deepseekv4}" +PREFLIGHT="${PREFLIGHT:-/root/arc-tools/gpu_box_preflight.sh}" +OUT="${OUT:-/root/logs/fp8sweep}" +STATUS="${STATUS:-/root/logs/fp8_sweep.status}" +PORT="${PORT:-1241}" +SEED="${SEED:-1234}" +MAX_TOKENS="${MAX_TOKENS:-48}" # per request; the floor is on the SUM +MIN_TOKENS_FLOOR="${MIN_TOKENS_FLOOR:-1000}" +MAX_SEQS="${MAX_SEQS:-512}" # must be >= the largest ladder point +LADDER="${LADDER:-5 8 64 256 512}" # offered decode batch = candidate M +LOCK="${LOCK:-/root/locks/gpu.lock}" +LOCK_WAIT="${LOCK_WAIT:-2400}" +OWNER="fp8sweep" + +mkdir -p "$OUT" "$(dirname "$STATUS")" "$(dirname "$LOCK")" +: > "$STATUS" +say() { echo "[$(date -u +%H:%M:%SZ)] $*" | tee -a "$STATUS"; } +die() { say "ABORT: $*"; echo "ABORT: $*" >> "$OUT/summary.txt"; exit 2; } + +say "fp8 three-arm threshold sweep — start" + +# ── GPU lock with OWNER-MATCHED RELEASE and evidence-based stale reclaim +# (same discipline as arcspec_token_identity_b8.sh; see the incident notes +# there for why release refuses locks this process does not own). +release_lock() { + if [ -f "$LOCK" ] && grep -q "^$OWNER pid=$$ " "$LOCK" 2>/dev/null; then + rm -f "$LOCK" + fi +} +STALE_AFTER="${STALE_AFTER:-120}" +lock_is_stale() { + [ -f "$LOCK" ] || return 1 + local age holder pid used procs + age=$(( $(date +%s) - $(stat -c %Y "$LOCK" 2>/dev/null || echo 0) )) + [ "$age" -ge "$STALE_AFTER" ] || return 1 + holder=$(cat "$LOCK" 2>/dev/null) + pid=$(printf '%s' "$holder" | sed -n 's/.*pid=\([0-9]*\).*/\1/p') + if [ -n "$holder" ] && [ -n "$pid" ] && kill -0 "$pid" 2>/dev/null; then + return 1 + fi + used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | head -1) + procs=$(nvidia-smi --query-compute-apps=pid --format=csv,noheader 2>/dev/null | wc -l) + [ "${used:-1}" -lt 1024 ] && [ "${procs:-1}" -eq 0 ] +} +acquire_lock() { + local waited=0 + while :; do + if ( set -o noclobber; echo "$OWNER pid=$$ started=$(date -u +%FT%TZ)" > "$LOCK" ) 2>/dev/null; then + trap 'release_lock' EXIT INT TERM + say "lock: acquired $LOCK" + return 0 + fi + if lock_is_stale; then + say "lock: RECLAIMING stale lock (age>=${STALE_AFTER}s, owner absent or dead, GPU idle)" + rm -f "$LOCK" + continue + fi + if [ "$waited" -ge "$LOCK_WAIT" ]; then + say "lock: HELD BY -> $(cat "$LOCK" 2>/dev/null)" + return 1 + fi + [ "$waited" = 0 ] && say "lock: waiting for -> $(cat "$LOCK" 2>/dev/null)" + sleep 15; waited=$((waited + 15)) + done +} +acquire_lock || die "GPU lock held by another chain for >${LOCK_WAIT}s" + +[ -f "$ARCENV" ] && . "$ARCENV" +if [ -x "$PREFLIGHT" ]; then + "$PREFLIGHT" > "$OUT/preflight.log" 2>&1 || die "preflight refused the box" + say "preflight OK" +else + say "WARN: no preflight at $PREFLIGHT — continuing without it" +fi +[ -x "$BIN" ] || die "$BIN not found — build once, then sweep (one binary, no rebuild)" +SHA=$(git -C "$WT" rev-parse HEAD 2>/dev/null) +say "binary ref: $SHA" + +# ── The prober: one canary + (M-1) filler prompts, all concurrent, all +# greedy. Non-streaming so `usage.completion_tokens` is authoritative. +cat > "$OUT/probe.py" <<'PYEOF' +"""Drive M concurrent greedy completions; report summed completion_tokens +(from usage — NEVER from requested max_tokens), wall seconds over the +concurrent burst, and the canary text (request 0).""" +import argparse, json, sys, threading, time, urllib.request + +CANARY = "List the first twelve prime numbers, separated by commas, then stop." +FILLERS = [ + "Name the seven continents, one per line.", + "Count from one to twenty in words, separated by spaces.", + "Write the first eight letters of the alphabet, separated by dashes.", + "Name eight colours of the visible spectrum, comma separated.", + "List eight consecutive even numbers starting at two.", + "Name the days of the week in order, one per line.", + "List the first eight powers of two, comma separated.", +] + +def one(port, prompt, max_tokens, seed, timeout): + body = json.dumps({ + "model": "default", "prompt": prompt, "max_tokens": max_tokens, + "temperature": 0, "seed": seed, "stream": False, + }).encode() + req = urllib.request.Request( + f"http://127.0.0.1:{port}/v1/completions", data=body, + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=timeout) as r: + d = json.load(r) + ch = (d.get("choices") or [{}])[0] + usage = d.get("usage") or {} + return { + "text": ch.get("text", ""), + "finish_reason": ch.get("finish_reason"), + "completion_tokens": usage.get("completion_tokens", 0), + } + +def burst(port, m, max_tokens, seed, timeout): + prompts = [CANARY] + [FILLERS[i % len(FILLERS)] for i in range(m - 1)] + out = [None] * m + errs = [] + def w(i): + try: + out[i] = one(port, prompts[i], max_tokens, seed, timeout) + except Exception as e: # noqa: BLE001 — a leg must record, not crash + errs.append(f"req[{i}]: {e}") + out[i] = {"text": "", "finish_reason": None, "completion_tokens": 0} + t0 = time.monotonic() + ts = [threading.Thread(target=w, args=(i,)) for i in range(m)] + for t in ts: t.start() + for t in ts: t.join() + wall = time.monotonic() - t0 + return out, wall, errs + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--port", type=int, required=True) + ap.add_argument("--m", type=int, required=True) + ap.add_argument("--max-tokens", type=int, required=True) + ap.add_argument("--seed", type=int, required=True) + ap.add_argument("--timeout", type=int, default=1800) + ap.add_argument("--warmup", action="store_true") + a = ap.parse_args() + + if a.warmup: + # Small unmeasured round: load caches, emit the first-dispatch log + # lines, absorb ramp. Its tokens are never counted. + burst(a.port, min(a.m, 8), 16, a.seed, a.timeout) + json.dump({"warmup": True}, sys.stdout) + return + + out, wall, errs = burst(a.port, a.m, a.max_tokens, a.seed, a.timeout) + total = sum(r["completion_tokens"] for r in out) + json.dump({ + "m": a.m, + "wall_s": wall, + "completion_tokens_sum": total, # from usage — the only numerator + "tok_per_s": (total / wall) if wall > 0 else 0.0, + "canary_text": out[0]["text"], + "canary_finish": out[0]["finish_reason"], + "canary_tokens": out[0]["completion_tokens"], + "errors": errs, + "short_streams": sum(1 for r in out if r["completion_tokens"] < 5), + }, sys.stdout) + +main() +PYEOF + +# ── One leg = one fresh server process + warm-up + measured burst. +# $1 arm name, $2 ladder M, then env pairs for the arm. +run_leg() { + local arm="$1" m="$2"; shift 2 + local leg="${arm}.M${m}" + local slog="$OUT/server.$leg.log" + say "leg $leg — starting server ($*)" + + local uqff_args=() + [ -n "$UQFF" ] && uqff_args=(--from-uqff "$UQFF") + # shellcheck disable=SC2086 + env RUST_LOG=info ARC_LOG_FP8_DISPATCH=1 "$@" \ + "$BIN" serve -p "$PORT" -m "$MODEL" -a "$ARCH" "${uqff_args[@]}" \ + --max-seqs "$MAX_SEQS" --prefix-cache-n 0 --seed "$SEED" \ + < /dev/null > "$slog" 2>&1 & + local spid=$! + local up=0 + for _ in $(seq 1 450); do + curl -s -m 5 "http://127.0.0.1:$PORT/v1/models" >/dev/null 2>&1 && { up=1; break; } + kill -0 "$spid" 2>/dev/null || break + sleep 2 + done + [ "$up" = "1" ] || { say "leg $leg — server never came up"; tail -20 "$slog" | tee -a "$STATUS"; return 1; } + + # Provenance: the running server must be the code we think it is. + local logged + logged=$(grep -m1 -o 'git revision: [0-9a-f]*' "$slog" | awk '{print $3}') + if [ -n "$SHA" ] && [ -n "$logged" ] && [ "$logged" != "$SHA" ]; then + say "leg $leg — PROVENANCE FAILED: server $logged != worktree $SHA" + kill -TERM "$spid" 2>/dev/null || true + return 1 + fi + + say "leg $leg — warm-up (unmeasured)" + python3 "$OUT/probe.py" --port "$PORT" --m "$m" --max-tokens "$MAX_TOKENS" \ + --seed "$SEED" --warmup > /dev/null || { say "leg $leg — warm-up failed"; kill -TERM "$spid" 2>/dev/null; return 1; } + + say "leg $leg — measured burst: M=$m x max_tokens=$MAX_TOKENS" + python3 "$OUT/probe.py" --port "$PORT" --m "$m" --max-tokens "$MAX_TOKENS" \ + --seed "$SEED" > "$OUT/$leg.json" || { say "leg $leg — probe failed"; kill -TERM "$spid" 2>/dev/null; return 1; } + + kill -TERM "$spid" 2>/dev/null || true + for _ in $(seq 1 30); do kill -0 "$spid" 2>/dev/null || break; sleep 2; done + kill -KILL "$spid" 2>/dev/null || true + sleep 5 + + # ENGAGEMENT: expected kernel present, both rivals absent, in the server log. + local want forbid1 forbid2 + case "$arm" in + cublaslt) want="dequant_cublaslt"; forbid1="wmma"; forbid2="tiled" ;; + wmma) want="wmma"; forbid1="dequant_cublaslt"; forbid2="tiled" ;; + tiled) want="tiled"; forbid1="wmma"; forbid2="dequant_cublaslt" ;; + *) say "leg $leg — unknown arm '$arm'"; return 1 ;; + esac + local ok=1 + grep -q "\[arc-fp8-dispatch\] path=$want " "$slog" || ok=0 + grep -q "\[arc-fp8-dispatch\] path=$forbid1 " "$slog" && ok=0 + grep -q "\[arc-fp8-dispatch\] path=$forbid2 " "$slog" && ok=0 + { + echo "leg=$leg" + grep '\[arc-fp8-dispatch\]' "$slog" | sed 's/^/ /' + echo " engagement=$([ "$ok" = 1 ] && echo PASS || echo FAIL)" + } >> "$OUT/engagement.txt" + if [ "$ok" != 1 ]; then + say "leg $leg — ENGAGEMENT FAILED (want $want, forbid $forbid1/$forbid2) — leg VOID" + echo '{"void": "engagement"}' > "$OUT/$leg.void" + fi + say "leg $leg — done" + return 0 +} + +# ── The 15 legs. Baseline (canary reference) FIRST: tiled @ smallest M. +MIN_M_NATIVE=1000000 # keeps every real m on the native path in arms (b),(c) +for m in $LADDER; do + run_leg tiled "$m" ARC_FP8_CUBLAS_MIN_M=$MIN_M_NATIVE ARC_NO_FP8_WMMA=1 \ + || die "tiled M=$m leg failed" +done +for m in $LADDER; do + run_leg wmma "$m" ARC_FP8_CUBLAS_MIN_M=$MIN_M_NATIVE ARC_NO_FP8_WMMA=0 \ + || die "wmma M=$m leg failed" +done +for m in $LADDER; do + run_leg cublaslt "$m" ARC_FP8_CUBLAS_MIN_M=5 \ + || die "cublaslt M=$m leg failed" +done + +# ── Summary: floors, canaries, the crossover table, the verdict. +python3 - "$OUT" "$MIN_TOKENS_FLOOR" $LADDER <<'PYEOF' | tee "$OUT/summary.txt" +import json, os, sys + +out, floor = sys.argv[1], int(sys.argv[2]) +ladder = [int(x) for x in sys.argv[3:]] +ARMS = ["tiled", "wmma", "cublaslt"] + +def load(arm, m): + leg = f"{arm}.M{m}" + if os.path.exists(os.path.join(out, leg + ".void")): + return None, "VOID(engagement)" + p = os.path.join(out, leg + ".json") + if not os.path.exists(p): + return None, "MISSING" + with open(p) as f: + d = json.load(f) + if d.get("errors"): + return d, f"ERRORS({len(d['errors'])})" + if d["completion_tokens_sum"] < floor: + return d, f"VOID(floor {d['completion_tokens_sum']}<{floor})" + if d["canary_tokens"] < 5 or not d.get("canary_finish"): + return d, "VOID(degenerate canary)" + return d, "CLEAN" + +base, base_status = load("tiled", ladder[0]) +base_text = base["canary_text"] if base else "" + +def canary(d): + if d is None or base is None: + return "-" + a, b = base_text, d["canary_text"] + if a == b: + return "identical" + n = min(len(a), len(b)) + i = next((j for j in range(n) if a[j] != b[j]), n) + return f"diverges@char{i}" + +print("=== FP8 three-arm threshold sweep ===") +print(f"floor: rate counted only when summed usage.completion_tokens >= {floor}") +print(f"canary baseline: tiled @ M={ladder[0]} [{base_status}]") +print() +print(f"{'M':>5} {'arm':>9} {'tokens':>7} {'wall_s':>8} {'tok/s':>9} {'status':>26} canary_vs_baseline") +rates = {} +for m in ladder: + for arm in ARMS: + d, status = load(arm, m) + rates[(arm, m)] = d["tok_per_s"] if (d and status == "CLEAN") else None + if d: + print(f"{m:>5} {arm:>9} {d['completion_tokens_sum']:>7} {d['wall_s']:>8.1f} " + f"{d['tok_per_s']:>9.1f} {status:>26} {canary(d)}") + else: + print(f"{m:>5} {arm:>9} {'-':>7} {'-':>8} {'-':>9} {status:>26} -") +print() +print("=== CROSSOVER (the number ARC_FP8_CUBLAS_MIN_M may move to) ===") +crossover = None +for m in ladder: + a = rates[("cublaslt", m)] + b = rates[("wmma", m)] + c = rates[("tiled", m)] + if a is None or b is None or c is None: + print(f"M={m:>4}: NOT ESTABLISHED — a leg is void; no claim about this M") + continue + native = min(b, c) + beat = a > native + print(f"M={m:>4}: cublaslt {a:.1f} vs min(wmma {b:.1f}, tiled {c:.1f}) = {native:.1f}" + f" -> {'cuBLASLt WINS' if beat else 'native wins'}") + if beat and crossover is None: + crossover = m +print() +if crossover is not None: + print(f"VERDICT: first clean M where cuBLASLt beats min(WMMA, tiled): {crossover}.") + print(f" Recommended default: ARC_FP8_CUBLAS_MIN_M={crossover} — merge WITH these") + print(" rows quoted in the PR (a threshold is a claim about every value it") + print(" excludes; these are the rows that back the claim).") +else: + print("VERDICT: no clean crossover on this ladder — the default STAYS 512.") + print(" Do not move the threshold from a partial table; re-run the void legs.") +PYEOF + +say "DONE — $OUT/summary.txt (engagement detail: $OUT/engagement.txt)" diff --git a/arc-tools/wave63_turboquant_hd512_gate.sh b/arc-tools/wave63_turboquant_hd512_gate.sh new file mode 100755 index 000000000..2b6284415 --- /dev/null +++ b/arc-tools/wave63_turboquant_hd512_gate.sh @@ -0,0 +1,153 @@ +#!/usr/bin/env bash +# wave63 — TurboQuant CUDA kernels at head_dim {64,128,256,512}, Hopper + Blackwell. +# +# THE ONE THING THIS PROVES +# `turbo_paged_attention.cu` compiles for sm_90a AND sm_100a/sm_103a with all +# four head-dim instantiations live, and the built archive actually contains +# both cubins. That is the whole D16 claim in PR #94; everything else here is +# secondary. +# +# EXIT CODES (D18 rule 2 — the distinction the runner acts on) +# 0 pass +# 1 genuine failure. Real signal: the code is wrong. Act on this. +# 2 this machine could not answer. Missing preflight, failing box, no +# toolkit, no repo. Says nothing about the code — do not read it as a +# regression. +# +# RUN (from main, per D15 — agents never call runcrate): +# setsid nohup bash /root/arc/arc-tools/wave63_turboquant_hd512_gate.sh \ +# < /dev/null > /root/wave63_gate.log 2>&1 & +# tail -f /root/wave63_gate.log ; cat /root/wave63_gate.status +# +# The preflight is looked up at a branch-independent path FIRST, because it +# lives on wave61/box-preflight-shared-prefix and vanishes the moment a runner +# checks out any other branch. Override with ARC_PREFLIGHT=/path. +# Set SKIP_NATIVE_BUILD=1 to go straight to the multi-arch proof. +set -uo pipefail + +STATUS=${STATUS:-/root/wave63_gate.status} +REPO=${REPO:-/root/arc} +BRANCH=${BRANCH:-feat/turboquant-hd512-v4} + +export LD_LIBRARY_PATH=/usr/local/cuda/compat:${LD_LIBRARY_PATH:-} +export CUDA_HOME=${CUDA_HOME:-/usr/local/cuda-13.1} +export PATH=$CUDA_HOME/bin:$PATH + +: > "$STATUS" +say() { echo "[$(date -u +%H:%M:%S)] $*"; echo "$*" >> "$STATUS"; } +# Exit 2 — the box could not answer. Not a verdict on the code. +noans() { say "NOANSWER $*"; say "EXIT 2 (environment could not answer)"; exit 2; } +# Exit 1 — the code is wrong. This is the only signal worth acting on. +fail() { say "FAIL $*"; say "EXIT 1 (genuine failure)"; exit 1; } + +# --------------------------------------------------------------- preflight +# gpu_box_preflight.sh ends with +# [ "$_arc_pf_sourced" = "1" ] || exit 1 +# so when it is SOURCED on a *failing* box that test succeeds, short-circuits +# the `|| exit 1`, and becomes the last command — making `source pf || handler` +# return 0 and the handler never fire. We therefore ignore its return value +# entirely and read `_ARC_PF_FAILED` directly, which is the flag it actually +# sets. If that flag is missing the script did not reach its own verdict, and +# an unverified box is a box that cannot answer. +PF="" +for cand in \ + "${ARC_PREFLIGHT:-}" \ + /root/arc-tools/gpu_box_preflight.sh \ + /usr/local/lib/arc/gpu_box_preflight.sh \ + "$REPO/arc-tools/gpu_box_preflight.sh" +do + [ -n "$cand" ] && [ -f "$cand" ] && { PF="$cand"; break; } +done +[ -n "$PF" ] || noans "no gpu_box_preflight.sh found. It lives on \ +wave61/box-preflight-shared-prefix, so it disappears whenever a runner checks \ +out another branch — stage it at /root/arc-tools/ (branch-independent) or set \ +ARC_PREFLIGHT=/path. Nothing about the code was tested." + +say "STEP 0 preflight ($PF)" +PF_FLAG=$(bash -c 'set +e; . "$1" >&2; echo "${_ARC_PF_FAILED:-unset}"' _ "$PF") +case "$PF_FLAG" in + 0) say "PASS preflight" ;; + 1) noans "preflight reports this box is unhealthy (_ARC_PF_FAILED=1). \ +See FAIL[...] lines above. The code was never exercised." ;; + *) noans "preflight never set _ARC_PF_FAILED (got '$PF_FLAG') — it did not \ +reach its own verdict, so its silence cannot be read as a pass." ;; +esac + +# ------------------------------------------------------------------- repo +cd "$REPO" 2>/dev/null || noans "no repo at $REPO" +command -v nvcc >/dev/null || noans "nvcc not on PATH (CUDA_HOME=$CUDA_HOME)" +git fetch --all -q 2>/dev/null || noans "cannot reach the git remote" +git checkout -q "$BRANCH" 2>/dev/null || noans "cannot check out $BRANCH" +git pull -q --ff-only 2>/dev/null || true +say "HEAD $(git rev-parse --short HEAD) on $BRANCH" +say "nvcc $(nvcc --version | grep -o 'release [0-9.]*' | head -1)" + +# ============================================================================ +# STEP 1 — THE D16 PROOF. Runs first and needs no full workspace build. +# ============================================================================ +say "STEP 1 compile paged-attn for sm_90a + sm_100a + sm_103a" +if ! ARC_CUDA_ARCHS=90,100,103 cargo build --release \ + --features cuda -p mistralrs-paged-attn 2>&1 | tail -60; then + fail "multi-arch build of mistralrs-paged-attn failed — see nvcc errors above" +fi +say "PASS multi-arch compile" + +# The build succeeding is not the claim; the cubins being present is. +LIB=$(find target/release/build -name 'libmistralrspagedattention.a' 2>/dev/null | head -1) +[ -n "$LIB" ] || noans "built, but no libmistralrspagedattention.a found to inspect" +command -v cuobjdump >/dev/null || noans "cuobjdump missing — cannot verify which \ +arches landed, and the unverified build is exactly the assertion we came to replace" + +ARCHES=$(cuobjdump --list-elf "$LIB" 2>/dev/null | grep -oE 'sm_[0-9]+a?' | sort -u | tr '\n' ' ') +say "cubin arches in archive: ${ARCHES:-}" +MISSING="" +for want in sm_90a sm_100a; do + case " $ARCHES " in *" $want "*) ;; *) MISSING="$MISSING $want" ;; esac +done +[ -z "$MISSING" ] || fail "D16 unmet: archive is missing$MISSING (has: $ARCHES). \ +The kernels compiled but this binary would PTX-JIT or fail on the missing arch." +say "PASS D16 — Hopper and Blackwell cubins both present" + +# ============================================================================ +# STEP 2 — native full build. PR #98 already built these kernels clean for +# sm_90 in 8m53s, so this is confirmation, not discovery. Skippable. +# ============================================================================ +if [ "${SKIP_NATIVE_BUILD:-0}" = "1" ]; then + say "SKIP step 2 (SKIP_NATIVE_BUILD=1)" +else + say "STEP 2 native build (cuda flash-attn; never cudnn — -62% decode on V4)" + cargo build --release --features "cuda flash-attn" 2>&1 | tail -40 \ + || fail "native build failed" + say "PASS native build" +fi + +# ============================================================================ +# STEP 3 — format tests. CPU-only by nature: these pin the CUDA constant tables +# to the Rust generator. Not a GPU claim, and not offered as one (D14). +# ============================================================================ +say "STEP 3 turboquant format tests" +cargo test --release -p mistralrs-quant turboquant:: 2>&1 | tail -20 \ + || fail "turboquant format tests failed" +say "PASS format tests" + +# ============================================================================ +# STEP 4 — decode no-regression. V4 does NOT use PagedAttention, so this does +# not exercise the TurboQuant kernel at all; it only shows this branch did not +# disturb the shipping path. +# ============================================================================ +UQFF=${UQFF:-/root/models/v4-uqff/qtip2b-0.uqff} +if [ "${SKIP_NATIVE_BUILD:-0}" = "1" ]; then + say "SKIP step 4 (no native binary was built)" +elif [ -f "$UQFF" ]; then + say "STEP 4 V4 decode smoke (no-regression only)" + timeout 900 ./target/release/mistralrs bench \ + -m /root/models/v4-src --uqff "$UQFF" \ + --n-gen 64 --n-prompt 64 --repetitions 2 2>&1 | tail -30 + say "STEP 4 done — compare tok/s against the FACTS.md V4 decode baseline" +else + say "SKIP step 4 — no UQFF at $UQFF" +fi + +say "ALL STEPS COMPLETE" +say "EXIT 0 (pass)" +exit 0 diff --git a/docs/engineering/OPENROUTER_READY.md b/docs/engineering/OPENROUTER_READY.md new file mode 100644 index 000000000..1a1e04887 --- /dev/null +++ b/docs/engineering/OPENROUTER_READY.md @@ -0,0 +1,89 @@ +# Arc, complete — the single gate for selling tokens + +**Merging this PR means Arc is the engine we set out to build, and it can be sold on OpenRouter.** +Not a defect list. The vision, shipped. Nothing reaches `master` for this effort except through here. + +Base `d7742670a` (the 15.0 → **34.48 tok/s** b=1 ladder, on master). + +--- + +## 🎯 THE GATE IS ONE NUMBER: **14,000 tok/s AGGREGATE ON ONE H200** + +| aggregate | tokens/GPU-hr | revenue @ $1/M | vs $4.92 card | verdict | +|---|---|---|---|---| +| ~50 (today, slow build) | 0.18 M | $0.18 | **27× underwater** | — | +| 1,400 | 5.0 M | $5 | break-even | **pointless** | +| **14,000** | **50.4 M** | **$50** | **10× margin** | **the business** | + +**Why 14,000 is an engineering target and not a fantasy:** + +| batch | memory-bound roofline | 14,000 is | +|---|---|---| +| 256 | 16,600 tok/s | **84% of peak** — nobody runs there | +| **512** | **~33,000 tok/s** | **42% of peak — normal for a good engine** | + +Memory traffic per step is **flat** as batch grows — 256 or 512 users read the same 74 GB once. +Every extra concurrent user is nearly free until compute-bound, somewhere around B≈512–1000. +**So 14K is a batching problem, not a kernel-speed problem.** + +🔴 **And batching is exactly what is broken.** Amortisation only pays when memory is the +constraint. **Our memory controller sits at 4%** — we are instruction-bound, so adding users +adds work instead of sharing it. Measured: b=1 is 85× off its bound, **aggregate at B=8 is +227× off** — we are *further* off at batch, which is the signature. + +**b=1 speed is a waypoint, not the goal.** Single-user is what one customer feels; aggregate +is what you sell. + +--- + +## What has to be true + +### 1 · Batching actually amortises ← **the gate** +- [ ] Sweep **B = 1/8/32/64/128/256/512**, per-request **and** aggregate. Current sweep stops at 128 — **too early; the answer lives past 256.** +- [ ] Find the knee. 256 experts top-6 means tokens barely collide until batch is large; the knee is where sharing starts and it is the single most important measurement in the program. +- [ ] Memory-controller utilisation at each point. **4% → 60%+ is the whole job.** +- [ ] KV bytes resident per user, and the OOM ceiling. + +### 2 · The engine runs near the silicon +- [ ] **Kernels at their own bounds.** Real profile on `qtip2b`: tail (sinkhorn + fast_sum + bmul + ucopy) **27.2%**, `fp8_gemv_warp` **16.0% → 13.0%** (1.31× landed, bit-exactness **not yet verified**), `qtip2b_gemv_tuned_kernel` **15.4%**. 🚩 A **4×4** Sinkhorn costing 9.3% of an H200 is structural. +- [ ] **4-stage async pipelining** — measured 54.8% → **92.9%** of peak across depths 1→4, saturating at 4. Issue rate first, then this. +- [ ] **TCFRAG** — 26.31 → 3.56 inst/weight, zero shuffles, cosine 1.0, **B=8 costs what b=1 costs** (the capacity property). Attaches to `qtip2b_gemv_tuned_kernel`. **Permute at LOAD (D22)** — artifacts stay architecture-neutral, one file for every GPU generation. +- [ ] **ArcGraph is INERT by default** — candle sits on the legacy NULL stream, which CUDA forbids capturing, so decode runs fully eagerly. Capture never engages unless env vars are set. +- [ ] **GPU WHILE-node loop** — syncs once per *generation* instead of once per token. Written, **zero callers**. Blocked by an unconditional `return Ok(None)` (`pipeline/normal.rs:2324`) and a U8-dtype bug failing all 28 call sites. **Neither competitor has this.** +- [ ] **`CudaSampler`** — complete, tested bit-exact, **zero callers, no gate**. Just unwired. + +### 3 · Prefill is a product, not an afterthought +- [ ] **TTFT ~24 s ⇒ usable.** ~12 ms per *input* token, ~100× off bound. **Prefill is the opposite disease from decode — the GPU is already 99.5% busy, so graphs and launch reduction buy it nothing.** Its own lane, its own kernels. +- [ ] **Prefill crashes**: b=1 on `qtip2b` exits 139 (SIGSEGV) and prints `pp = 0.000±0.000` into a results table instead of failing. +- [ ] **Chunked prefill** — inert until the MoE gather lands. + +### 4 · The memory system Arc was designed around +- [ ] **Prefix caching is silently disabled** by our own default (TurboQuant is the default cache type). Most real traffic is repeated prefixes. +- [ ] **56× KV over-retention** — root-caused; the fix reportedly never reached V4. +- [ ] **Long context measured**: at 1M ctx the answer is **~3 users or ~55** depending on whether the 128-token window + compress ratios are live. Unmeasured, and it is the capacity claim. +- [ ] **Ragged batching / ArcSched** — admit mixed-length cohorts; the per-seq advance dies on ragged batches today. + +### 5 · The multipliers we designed and haven't switched on +- [ ] **Speculative decoding (MTP)** — measured 1.93 tok/step at b=1, 1.06 at b=128. A near-2× at low batch, off. +- [ ] **Expert parallelism** — does not exist. The natural MoE scaling axis. +- [ ] **ArcTarget multi-arch** — B200/B300 as first-class, one artifact, N load-time tables (D22). + +### 6 · Trust — because our instruments have been the problem +- [ ] **Twelve instruments lied in one session.** bench counted *requested* tokens (**1280 tok/s for a zero-token run**) · `cargo clean -p` silently no-ops on git deps · `nsys` hides graph-internal kernels by default · a stale `.o` satisfied a link · `replay()` returned a clone of the graded buffer · `append_graph` compared a buffer against itself · an XOR control cancelled over an even tile count · an atomic read 0 while the work ran · a sampler slept past the run · bf16 parity vacuous at 8 mantissa bits · a guard died on Windows CRLF · prefill printed 0.000. +- [ ] **The live V4 fixture is all zeros** — identity tests against it prove nothing. +- [ ] **Quality re-confirmed on the exact artifact served.** GSM8K 96.0% exists — on a different rung. +- [ ] **One measurement harness**, and nothing measures any other way. + +### 7 · Commercial and operational +- [ ] **Token accounting exact** — OpenRouter bills on it, and our own bench was miscounting. +- [ ] **Sustained-load soak** — 400 requests passed once; nobody has run hours. +- [ ] **Deploy, monitor, restart-on-crash.** +- [ ] **OpenRouter onboarding + pricing** — requirements unknown; needs a doc read, not a guess. + +--- + +### Standing rules for anything landing here +**Fast path is the DEFAULT** — flags are kill switches, never opt-ins. **Artifacts stay +architecture-neutral; permute at load (D22).** Every arm proves it engaged. Compare **F32 bit +patterns before narrowing**, with a control that must fire — a single-element 1-ULP +perturbation, never an XOR that can cancel. **A scoping result is never a verdict (D21).** diff --git a/mistralrs-bench/Cargo.toml b/mistralrs-bench/Cargo.toml index 8d7fb240f..224d2dbf4 100644 --- a/mistralrs-bench/Cargo.toml +++ b/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 diff --git a/mistralrs-bench/src/main.rs b/mistralrs-bench/src/main.rs index c77b677dd..5f22d45d5 100644 --- a/mistralrs-bench/src/main.rs +++ b/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, 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 = 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) { let backend = match device { Device::Cpu => "CPU", @@ -205,6 +324,11 @@ fn print_usage(model: &str, device: &Device, results: Vec) { (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) { "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); } diff --git a/mistralrs-cli/src/args/mod.rs b/mistralrs-cli/src/args/mod.rs index 675fcc069..4cf461e3c 100644 --- a/mistralrs-cli/src/args/mod.rs +++ b/mistralrs-cli/src/args/mod.rs @@ -427,7 +427,23 @@ pub struct GlobalOptions { #[derive(clap::Args, Clone, Deserialize)] pub struct RuntimeOptions { /// How many requests may run concurrently. - #[arg(long, default_value_t = 32, help_heading = "Serving")] + /// + /// Default 256. The previous default of 32 capped aggregate throughput far + /// below the hardware: at B=32 the physics ceiling on an H200 is ~3,740 + /// tok/s, so the 14,000 tok/s aggregate gate (which needs B >= 214) was + /// unreachable on the default configuration — an admission cap, not a + /// kernel limit. + /// + /// Memory is the real bound, and it scales with this knob: KV/state memory + /// is per-sequence (~222 MiB/seq at ctx-4096 for V4 Flash-class models, so + /// 256 concurrent sequences is ~55 GiB of context memory on top of the + /// weights). PagedAttention deployments are protected by block-budget + /// admission (a sequence that does not fit waits for free blocks); the + /// `DefaultScheduler` (dense-cache models, including V4) admits by COUNT + /// only and has no VRAM-aware check — on small cards, lower this + /// explicitly. Wiring the default through a VRAM-aware admission check for + /// the dense path is tracked as follow-up, not built here. + #[arg(long, default_value_t = 256, help_heading = "Serving")] #[serde(default = "default_max_seqs")] pub max_seqs: usize, @@ -558,7 +574,7 @@ impl Default for GlobalOptions { impl Default for RuntimeOptions { fn default() -> Self { Self { - max_seqs: 32, + max_seqs: 256, no_kv_cache: false, prefix_cache_n: 16, chat_template: None, @@ -580,7 +596,7 @@ fn default_token_source() -> TokenSource { } fn default_max_seqs() -> usize { - 32 + 256 } fn default_prefix_cache_n() -> usize { diff --git a/mistralrs-cli/src/args/paged_attn.rs b/mistralrs-cli/src/args/paged_attn.rs index 9e057abf8..d52afa8b3 100644 --- a/mistralrs-cli/src/args/paged_attn.rs +++ b/mistralrs-cli/src/args/paged_attn.rs @@ -59,15 +59,18 @@ pub struct PagedAttentionOptions { /// this flag off gives you TurboQuant KV and, with it, no prefix caching. /// Pass `auto` to opt out. Every other geometry — MLA layouts and every /// other head_dim — auto-falls back to `auto` with a warning, so those - /// models are unaffected. Setting the flag explicitly turns that fallback - /// into a hard error instead. + /// models are unaffected. Setting the flag explicitly widens the accepted + /// set to every instantiated head_dim (64/128/256/512) and turns that + /// fallback into a hard error instead. /// Provenance: TurboQuant K4/V3 has one end-to-end serving run behind it — /// Qwen3-32B on a B200, 55 tok/s with correct output (2026-04-06). That run /// was b=1, head_dim 128, on Blackwell, and it did not isolate TurboQuant /// from the rest of the decode path. There is no A/B against an /// unquantized cache and no quality evaluation at any preset, so treat - /// quality as unestablished; `turboquant-3` and `turboquant-aggressive` - /// have never run at all. + /// quality as unestablished; 64/256/512 are compiled but have never + /// executed on hardware, which is why the default will not choose them for + /// you, and `turboquant-3` and `turboquant-aggressive` have never run at + /// all. // Hidden: overriding Arc's resolved KV cache type is the clearest way to // silently make a deployment worse — every TurboQuant variant currently // disables prefix caching, and an explicit choice converts the safe diff --git a/mistralrs-cli/src/commands/tune.rs b/mistralrs-cli/src/commands/tune.rs index 028e6cb3e..b5cf70aee 100644 --- a/mistralrs-cli/src/commands/tune.rs +++ b/mistralrs-cli/src/commands/tune.rs @@ -174,7 +174,7 @@ fn emit_toml_config( out.push_str("host = \"0.0.0.0\"\n"); out.push_str("port = 1234\n\n"); out.push_str("[runtime]\n"); - out.push_str("max_seqs = 32\n\n"); + out.push_str("max_seqs = 256\n\n"); out.push_str("[[models]]\n"); out.push_str(&format!("kind = \"{}\"\n", model_kind(model_type))); 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/attention/backends/sinks.rs b/mistralrs-core/src/attention/backends/sinks.rs index 7a41e8fec..cec200002 100644 --- a/mistralrs-core/src/attention/backends/sinks.rs +++ b/mistralrs-core/src/attention/backends/sinks.rs @@ -7,7 +7,9 @@ use crate::{attention::SdpaParams, pipeline::text_models_inputs_processor::Flash /// Dispatches to: /// CUDA -> flash_attn_sinks / flash_attn_sinks_varlen /// Metal -> flash_attn_sinks_metal / flash_attn_sinks_varlen_metal -/// CPU -> unfused matmul + softmax_with_sinks +/// otherwise -> unfused matmul + softmax_with_sinks (device-agnostic candle +/// ops — this is the fall-through for ANY device the fused kernels decline, +/// not a CPU-only path; DeepSeek-V4's head_dim 512 takes it on CUDA too) /// /// Varlen is used when flash_params contains cu_seqlens_k for this device AND /// q has batch > 1 AND the caller supplied no explicit mask. @@ -209,8 +211,16 @@ fn sinks_attn_regular( }); } + // 🔴 SAME DEFECT, OTHER BACKEND. The gate above stops the CUDA arm from + // swallowing masked work; this arm had no such gate, so a masked call on + // Metal (GPT-OSS passes a mask it asserts is `Some`, and every fused head + // dim {64,80,96,128,256} is reachable there) went to + // `flash_attn_sinks_metal` — which takes no mask parameter at all — and the + // mask was dropped on the floor. Unlike the CUDA kernel, the Metal entry + // point cannot even refuse a mask at the interface (it has no argument to + // refuse with), so this dispatch gate is the only line of defence. #[cfg(feature = "metal")] - if q.device().is_metal() && flash_sinks_ok_metal { + if q.device().is_metal() && flash_sinks_ok_metal && mask.is_none() { crate::attention::arcflash::note(crate::attention::arcflash::Path::VendorSinks); return mistralrs_quant::flash_attn_sinks_metal( q, @@ -221,6 +231,19 @@ fn sinks_attn_regular( window_size, ); } + #[cfg(feature = "metal")] + if q.device().is_metal() && flash_sinks_ok_metal && mask.is_some() { + // Loud once: correct-and-slower on purpose, exactly as on CUDA. + static WARNED: std::sync::OnceLock<()> = std::sync::OnceLock::new(); + WARNED.get_or_init(|| { + tracing::warn!( + target: "arcflash", + "fused head_dim={hd} Metal sinks kernel DECLINED: an explicit attention \ + mask is present and the kernel cannot honour one (it has no mask \ + parameter). Falling back to the unfused masked path." + ); + }); + } // ArcFlash/Tile — the head_dim the vendor sinks kernel does not take. // DeepSeek-V4's 512 is the whole reason this arm exists. @@ -299,7 +322,12 @@ fn sinks_attn_varlen( ); } - // CPU fallback: per-sequence loop (to_vec1 is fine on CPU path) + // Unfused fallback: per-sequence loop. Device-agnostic, and reached on + // CUDA/Metal too whenever the fused varlen kernel declines the head_dim — + // where these `to_vec1` calls are device-to-host syncs, once per layer per + // step. Hoisting the host copy of cu_seqlens into `FlashParams` is the fix + // for that; it is plumbing shared with the eager path and out of scope + // here. let cu_q_vec: Vec = cu_seqlens_q.to_vec1()?; let cu_k_vec: Vec = cu_seqlens_k.to_vec1()?; sinks_attn_cpu_varlen( @@ -350,6 +378,47 @@ fn varlen_causal_mask( Tensor::from_vec(data, (1, 1, q_len, kv_len), device)?.to_dtype(dtype) } +/// Memoizing front end for [`varlen_causal_mask`]. +/// +/// The builder is a `Tensor::from_vec` — a host-to-device copy when the +/// tensors live on a GPU (CLAUDE.md pitfall #5) — and `sinks_attn_cpu_varlen` +/// calls it once per **sequence** per layer per step. The mask is a pure +/// function of the integer key, and in a uniform-length batch every sequence +/// of every layer shares one, so a single slot turns per-(sequence × layer) +/// uploads into one per shape change. A ragged batch alternates keys and +/// rebuilds per sequence — exactly the pre-memo cost, never worse. +fn varlen_causal_mask_cached( + q_len: usize, + kv_len: usize, + window_size: usize, + dtype: candle_core::DType, + device: &candle_core::Device, +) -> Result { + type Key = ( + usize, + usize, + usize, + candle_core::DType, + candle_core::DeviceLocation, + ); + thread_local! { + static MEMO: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; + } + let key = (q_len, kv_len, window_size, dtype, device.location()); + MEMO.with(|c| -> Result { + let mut c = c.borrow_mut(); + if let Some((k, t)) = c.as_ref() { + if *k == key { + return Ok(t.clone()); + } + } + let t = varlen_causal_mask(q_len, kv_len, window_size, dtype, device)?; + *c = Some((key, t.clone())); + Ok(t) + }) +} + /// Sinks attention on **ArcFlash/Tile** — the path every DeepSeek-V4 layer /// takes, because head_dim 512 is outside the fused sinks kernel's /// `{64,80,96,112,128,192,256}`. @@ -376,7 +445,9 @@ fn sinks_attn_cpu( crate::attention::arcflash::union_attention(q, k, v, mask, Some(sinks), sdpa_params) } -/// CPU fallback for varlen: per-sequence unfused loop. +/// Unfused fallback for varlen: per-sequence loop. Despite the name this is +/// device-agnostic candle math, not CPU-only — it is the fall-through on any +/// device whose fused varlen kernel declines the head_dim. /// /// Each sequence is masked with `varlen_causal_mask`, matching the fused /// varlen kernels. Passing `None` here (the pre-fix behavior) made every @@ -416,7 +487,7 @@ fn sinks_attn_cpu_varlen( .transpose(0, 1)? .unsqueeze(0)?; - let mask = varlen_causal_mask(q_len, kv_len, window_size, qi.dtype(), device)?; + let mask = varlen_causal_mask_cached(q_len, kv_len, window_size, qi.dtype(), device)?; let oi = sinks_attn_cpu(&qi, &ki, &vi, sinks, Some(&mask), sdpa_params)?; // Pad back to max_q @@ -579,6 +650,90 @@ mod tests { ); Ok(()) } + + /// Additive causal mask whose last `pad` key columns are dead PADDING — + /// the exact class of information the fused kernels cannot re-derive: + /// they rebuild causality and windows internally, but nothing tells them + /// which trailing keys belong to a shorter neighbour's padding. + fn padding_mask(seq: usize, pad: usize, dev: &Device) -> Result { + let mut m = vec![0f32; seq * seq]; + for i in 0..seq { + for j in 0..seq { + if j > i || j >= seq - pad { + m[i * seq + j] = f32::NEG_INFINITY; + } + } + } + Tensor::from_vec(m, (1, 1, seq, seq), dev) + } + + /// A padding mask must survive the dispatch end to end, and dropping it + /// must CHANGE the output — so this test fails on any regression that + /// re-routes masked work to a maskless kernel, rather than passing by + /// luck because the mask happened to be redundant with plain causality. + #[test] + fn padding_mask_is_honored_and_dropping_it_changes_the_output() -> Result<()> { + let dev = Device::Cpu; + let (b, h, seq, d) = (2usize, 2usize, 6usize, 16usize); + let sinks = Tensor::from_vec(vec![0.25f32, 0.5], (1, h, 1, 1), &dev)?; + let sdpa = SdpaParams { + sinks: Some(sinks.clone()), + ..params(d, h, None) + }; + + let q = mk(&[b, h, seq, d], 0.05, &dev)?; + let k = mk(&[b, 1, seq, d], 0.13, &dev)?; + let v = mk(&[b, 1, seq, d], 0.23, &dev)?; + let mask = padding_mask(seq, 2, &dev)?; + + // flash_params present and b > 1: the configuration where a dropped + // mask would additionally re-route the call onto the varlen path. + let flash = varlen_params(&[seq as u32; 2], &dev); + let got = sinks_attn(&q, &k, &v, &sinks, Some(&mask), Some(&flash), &sdpa)?; + let want = sinks_attn_cpu(&q, &k, &v, &sinks, Some(&mask), &sdpa)?; + assert_eq!( + max_abs_diff(&flat(&got), &flat(&want)), + 0.0, + "the dispatch did not honor the padding mask" + ); + + // Teeth: with the mask dropped (the defect), the answer is different + // by a wide margin — the padding columns vote in the softmax. + let dropped = sinks_attn_cpu(&q, &k, &v, &sinks, None, &sdpa)?; + let signal = max_abs_diff(&flat(&want), &flat(&dropped)); + assert!( + signal > 1e-2, + "masked and maskless attention are indistinguishable ({signal}); \ + this test has no teeth" + ); + Ok(()) + } + + /// The memoized varlen mask must be bit-identical to a fresh build across + /// key changes — same shape with a different window is exactly the + /// stale-serving hazard — and must actually HIT on a repeated key: a memo + /// that always misses passes every value check while doing nothing. + #[test] + fn varlen_mask_memo_hits_and_never_serves_stale() -> Result<()> { + let dev = Device::Cpu; + for &(q_len, kv_len, window) in &[(5usize, 8usize, 0usize), (5, 8, 3), (5, 8, 0), (4, 9, 2)] + { + let fresh = varlen_causal_mask(q_len, kv_len, window, DType::F32, &dev)?; + let cached = varlen_causal_mask_cached(q_len, kv_len, window, DType::F32, &dev)?; + assert_eq!( + flat(&fresh), + flat(&cached), + "memo served a wrong mask for ({q_len}, {kv_len}, {window})" + ); + } + // Engagement: a repeated key returns the same tensor handle. Skipping + // the build leaves no trace in the values, so only the handle proves + // the cache did anything. + let a = varlen_causal_mask_cached(7, 7, 2, DType::F32, &dev)?; + let b = varlen_causal_mask_cached(7, 7, 2, DType::F32, &dev)?; + assert_eq!(a.id(), b.id(), "repeated key did not hit the memo"); + Ok(()) + } } /// The fused `head_dim=512` flash-sinks path takes no mask argument, so it @@ -820,3 +975,181 @@ mod fused_sinks_mask_tests { Ok(()) } } + +/// The Metal twin of `fused_sinks_mask_tests`. `flash_attn_sinks_metal` takes +/// no mask parameter at all, so — unlike CUDA, where the kernel itself now +/// refuses a mask — the dispatch gate is the ONLY line of defence for masked +/// work on Metal. head_dim 128 sits inside Metal's fused set, so the fused arm +/// is exactly the one that would engage. +/// +/// cargo test -p mistralrs-core --features metal metal_sinks -- --nocapture +#[cfg(all(test, feature = "metal"))] +mod metal_sinks_mask_tests { + use super::*; + use candle_core::{DType, Device}; + + fn mk(dims: &[usize], seed: f32, dev: &Device) -> Result { + let n: usize = dims.iter().product(); + let data: Vec = (0..n).map(|i| ((i as f32) * seed).sin()).collect(); + Tensor::from_vec(data, dims, dev) + } + + fn to_f32(t: &Tensor) -> Vec { + t.to_dtype(DType::F32) + .unwrap() + .flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + } + + fn max_abs(a: &[f32], b: &[f32]) -> f32 { + assert_eq!(a.len(), b.len()); + a.iter() + .zip(b) + .map(|(x, y)| (x - y).abs()) + .fold(0f32, f32::max) + } + + /// Additive 0 / -inf mask [1, 1, q, k]: causal, and only the first `keep` + /// keys are admissible — the padding-column class the kernel cannot + /// re-derive. + fn band_mask(q_len: usize, kv_len: usize, keep: usize, dev: &Device) -> Result { + let mut m = vec![0f32; q_len * kv_len]; + for i in 0..q_len { + for j in 0..kv_len { + let causal_ok = j <= i + (kv_len - q_len); + if !causal_ok || j >= keep { + m[i * kv_len + j] = f32::NEG_INFINITY; + } + } + } + Tensor::from_vec(m, (1, 1, q_len, kv_len), dev) + } + + const B: usize = 1; + const H: usize = 8; + const KVH: usize = 1; + const QL: usize = 16; + const KL: usize = 48; + const D: usize = 128; // inside Metal's fused set {64, 80, 96, 112, 128, 192, 256} + + fn fixtures(dev: &Device) -> Result<(Tensor, Tensor, Tensor, Tensor, SdpaParams)> { + let q = mk(&[B, H, QL, D], 0.017, dev)?; + let k = mk(&[B, KVH, KL, D], 0.023, dev)?; + let v = mk(&[B, KVH, KL, D], 0.031, dev)?; + let sinks = mk(&[H], 0.041, dev)?; + let p = SdpaParams { + n_kv_groups: H / KVH, + softcap: None, + softmax_scale: 1.0 / (D as f32).sqrt(), + sliding_window: None, + sinks: None, + }; + Ok((q, k, v, sinks, p)) + } + + /// THE BASELINE. The fused Metal kernel, called exactly as the dispatch + /// calls it, cannot depend on a mask it is never given: its output tracks + /// the causal-only reference and misses the masked one by the mask's whole + /// effect. This is the delta the dispatch gate exists to keep out of + /// serving — and the size of the wrong answer if the gate regresses. + #[test] + fn fused_metal_kernel_cannot_honour_a_mask() -> Result<()> { + let dev = Device::new_metal(0).expect("this test requires a Metal device"); + let (q, k, v, sinks, p) = fixtures(&dev)?; + + let mask_causal = band_mask(QL, KL, KL, &dev)?; + let mask_narrow = band_mask(QL, KL, 8, &dev)?; + + let ref_causal = to_f32(&crate::attention::arcflash::union_attention( + &q, + &k, + &v, + Some(&mask_causal), + Some(&sinks), + &p, + )?); + let ref_narrow = to_f32(&crate::attention::arcflash::union_attention( + &q, + &k, + &v, + Some(&mask_narrow), + Some(&sinks), + &p, + )?); + let fused = to_f32(&mistralrs_quant::flash_attn_sinks_metal( + &q, + &k, + &v, + Some(&sinks), + p.softmax_scale, + 0, + )?); + + let d_causal = max_abs(&fused, &ref_causal); + let d_narrow = max_abs(&fused, &ref_narrow); + let d_masks = max_abs(&ref_causal, &ref_narrow); + eprintln!( + "max|fused - ref(causal only)| = {d_causal:.6}\n\ + max|fused - ref(masked)| = {d_narrow:.6}\n\ + max|ref(causal) - ref(masked)| = {d_masks:.6} <- the mask's own effect" + ); + assert!( + d_masks > 1e-3, + "the two masks did not change the reference ({d_masks}); test would be vacuous" + ); + assert!( + d_narrow > d_causal * 10.0, + "expected the maskless Metal kernel to track the causal-only reference \ + (d_causal {d_causal}, d_narrow {d_narrow})" + ); + Ok(()) + } + + /// THE GATE. A masked call through the dispatch must match the masked + /// unfused reference — i.e. it must NOT have gone to the fused kernel the + /// baseline above proves would drop the mask. Reverting the + /// `mask.is_none()` guard on the Metal arm fails this test by + /// approximately the `d_narrow` the baseline prints. + #[test] + fn masked_sinks_attention_on_metal_matches_the_masked_reference() -> Result<()> { + let dev = Device::new_metal(0).expect("this test requires a Metal device"); + let (q, k, v, sinks, p) = fixtures(&dev)?; + let mask = band_mask(QL, KL, 8, &dev)?; + + let got = to_f32(&sinks_attn(&q, &k, &v, &sinks, Some(&mask), None, &p)?); + let want = to_f32(&crate::attention::arcflash::union_attention( + &q, + &k, + &v, + Some(&mask), + Some(&sinks), + &p, + )?); + + // Negative control on the OUTPUT: a shared-input perturbation would + // move both sides together and could never fire. + { + let mut probe = got.clone(); + assert_eq!(max_abs(&probe, &got), 0.0); + probe[7] = f32::from_bits(probe[7].to_bits().wrapping_add(1)); + assert!(max_abs(&probe, &got) > 0.0, "comparator is inert"); + } + + let d = max_abs(&got, &want); + eprintln!("masked Metal dispatch vs masked reference: max abs diff {d:.8}"); + assert!( + d < 1e-3, + "masked sinks attention on Metal does not match the masked reference \ + (max abs diff {d}) -- the Metal arm is dropping the mask again" + ); + + // ...and the maskless fused fast path still works. + assert!( + sinks_attn(&q, &k, &v, &sinks, None, None, &p).is_ok(), + "declining masks must not break the maskless fused path" + ); + Ok(()) + } +} diff --git a/mistralrs-core/src/cuda/qk_norm_rope.rs b/mistralrs-core/src/cuda/qk_norm_rope.rs index 68314ec23..3f15456dc 100644 --- a/mistralrs-core/src/cuda/qk_norm_rope.rs +++ b/mistralrs-core/src/cuda/qk_norm_rope.rs @@ -49,6 +49,64 @@ pub fn fused_enabled() -> bool { *ENABLED.get_or_init(|| !matches!(std::env::var("ARC_QK_FUSED").as_deref(), Ok("0"))) } +/// `ARC_QK_FUSED_COHORT=0` stops the fused kernel from serving a **batched** +/// cohort (`batch > 1`). **Default ON since the s10 measurement** (box +/// `arc-s10-ledger`, H200, candidate `4d03b9e25`, leg 4): B=256 decode +/// aggregate 599.35 vs 575.12 tok/s baseline, seeded canaries byte-identical, +/// ENGAGED line present at B=256, 0/512 request errors. The +4.2% sits within +/// the observed inter-leg variance band (~±5%, leg 4b's no-op delta), so the +/// default stands on "engaged + bit-identical + not slower", not on a claimed +/// speedup. +/// +/// # Why this was opt-in until measured, and [`fused_enabled`] never was +/// +/// The kernel is batch-correct by construction — `qk_norm_rope.cu` flattens +/// `(b, t)` into `blockIdx.y`, recovers `b = bt / seq_len`, and reads the table +/// at `pos_offset + t`, a position independent of `b` — and until #198 the +/// dispatch gate tested `seqlen_offsets.len() != 1`, the LENGTH of the offset +/// vector where the property needed is the DISTINCTNESS of its values. So the +/// fused path was dead at every batch size above one. #198 fixes that gate. +/// +/// Fixing the gate, though, makes a kernel **live on a path it has never run +/// on**, and this repo has paid for that twice in one week: +/// +/// * TCFRAG shipped default-ON with "UNVERIFIED ON HARDWARE — NEVER RUN" in its +/// own header. It held 63 GB and permanently broke a layer through a poisoned +/// `OnceLock`; #209 made it opt-in. +/// * The fused-512 attention path ran for four days silently dropping the +/// attention mask, at 12% agreement with the reference, because nobody ran +/// the path it had quietly become live on. +/// +/// The doctrine was: **unverified means default-off, and "unverified" means +/// unmeasured, not new.** That measurement has now happened (s10 leg 4, above) +/// and the default flipped in the same session, exactly as the old comment +/// promised. `batch == 1` was always unaffected. +/// +/// Read by VALUE, `!= Some("0")`, not by presence: 23 `ARC_*` flags were +/// converted (#212) precisely because `var_os(..).is_some()` made +/// `ARC_FOO=0` mean ON. +pub fn fused_cohort_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + fused_cohort_enabled_from(std::env::var("ARC_QK_FUSED_COHORT").ok().as_deref()) + }) +} + +/// Pure half of [`fused_cohort_enabled`], so the polarity is testable without +/// mutating the process environment. +/// +/// Same reasoning as `mistralrs_quant::env_flag_value` and +/// `qtip::tcfrag2b_enabled_from`: mutating the environment from a test is racy +/// across `cargo test`'s threads and, since the 2024 edition, `unsafe` — and +/// the caller latches its answer in a `OnceLock`, so such a test would prove +/// only whatever the first reader in the process happened to see. Testing the +/// pure function against every input is the only form of this test that means +/// anything. +pub fn fused_cohort_enabled_from(value: Option<&str>) -> bool { + value != Some("0") +} + /// `ARC_QK_VERIFY=1` runs the eager chain alongside the fused kernel at every /// layer and compares the raw output bytes. /// @@ -443,3 +501,71 @@ static VERIFIED: AtomicU64 = AtomicU64::new(0); pub fn verified_elements() -> u64 { VERIFIED.load(Ordering::Relaxed) } + +#[cfg(test)] +mod cohort_gate_tests { + use super::*; + + /// The batched fused path must be OFF unless explicitly asked for. + /// + /// This is the assertion that stops a repeat of TCFRAG — default-ON with + /// "UNVERIFIED ON HARDWARE — NEVER RUN" in its own header, 63 GB held, a + /// layer permanently broken through a poisoned `OnceLock`, retired by #209 + /// — and of the fused-512 attention path, which ran for four days silently + /// dropping the attention mask at 12% agreement with the reference. Both + /// were live on a path nobody had run. + /// + /// These run with no features and no GPU, on three operating systems, on + /// every PR, because the polarity is a pure function. + #[test] + fn the_batched_cohort_is_on_by_default_since_the_s10_measurement() { + assert!( + fused_cohort_enabled_from(None), + "unset must mean ON — measured s10 leg 4 (arc-s10-ledger, 4d03b9e25): \ + engaged at B=256, canaries byte-identical, 599.35 vs 575.12 tok/s" + ); + assert!(fused_cohort_enabled_from(Some("1"))); + } + + /// Read by VALUE, not presence. #212 converted 23 `ARC_*` flags because + /// `var_os(..).is_some()` made `ARC_FOO=0` mean ON — which also silently + /// cancels any A/B whose control arm sets the flag to zero. After the s10 + /// flip the kill switch is the exact string "0" and nothing else. + #[test] + fn only_literal_zero_disables_the_cohort() { + assert!( + !fused_cohort_enabled_from(Some("0")), + "\"0\" must disable — it is the A/B control arm and the kill switch" + ); + for v in ["", "false", "no", "off", "true", "yes", "on", "2", "1 ", "0 "] { + assert!( + fused_cohort_enabled_from(Some(v)), + "{v:?} must leave the measured default ON; only the exact string \"0\" disables" + ); + } + } + + /// [`engaged_count`] is the instrument the A/B harness reads to tell "the + /// kernel ran" apart from "the benchmark got faster" — and until this + /// test, nothing in the tree read it, so it could have silently stopped + /// advancing and every claim built on it would still look green. The + /// counters are process-global and other tests may bump them + /// concurrently, so the assertion is on a lower bound of the delta, which + /// only this thread's calls can guarantee. + #[test] + fn engaged_count_advances_when_the_paths_are_taken() { + let (e0, d0) = engaged_count(); + note_engaged(128, 64, 16, "bf16"); + note_declined("test: exercising the decline counter"); + let (e1, d1) = engaged_count(); + assert!( + e1 >= e0 + 1, + "note_engaged did not advance the engaged counter ({e0} -> {e1}); \ + an A/B reading engaged_count() would call a dead path live" + ); + assert!( + d1 >= d0 + 1, + "note_declined did not advance the declined counter ({d0} -> {d1})" + ); + } +} diff --git a/mistralrs-core/src/cuda/sinkhorn.cu b/mistralrs-core/src/cuda/sinkhorn.cu index 767ac55a1..4deca694c 100644 --- a/mistralrs-core/src/cuda/sinkhorn.cu +++ b/mistralrs-core/src/cuda/sinkhorn.cu @@ -62,6 +62,7 @@ #include #include +#include #if defined(__USE_FAST_MATH__) #error "sinkhorn.cu must be compiled WITHOUT --use_fast_math: fast math rewrites expf to __expf and IEEE division to approximate reciprocals, breaking bit-identity with candle-kernels (which build with plain -O3). See the dedicated no-fast-math builder in mistralrs-core/build.rs." @@ -257,6 +258,138 @@ __global__ void sinkhorn_normalize_f32_kernel( } // namespace +// --------------------------------------------------------------------------- +// WARP-BUTTERFLY PATH (default) +// --------------------------------------------------------------------------- +// The templated kernel above got the row/tree buffers into registers, but it +// still stages the matrix through shared memory and pays TWO __syncthreads() +// per iteration -- 40 barriers and ~80 shared round trips per call -- to move +// each column to the thread that divides by it. With hc <= 16 every thread is +// a lane of ONE warp, so none of that is needed: the column reduction is a +// warp butterfly and the matrix never leaves registers. +// +// MEASURED, H200 / nvcc 12.4 / sm_90, n=1 hc=4 iters=20, 86 calls per token: +// legacy (pre-ladder, runtime hc, 192 B stack) 30.578 us/call +// templated + shared/barriers (the kernel above) 9.935 us/call +// warp butterfly (this kernel) 7.302 us/call 1.36x +// +// Bit-identity holds by construction. candle's level `s` is +// `buf[t] = __fadd_rn(buf[t], buf[t + s])` for t < s; the butterfly gives lane +// t `__fadd_rn(buf[t], buf[t ^ s])`, which IS that expression for t < s and its +// commuted form for t >= s. IEEE-754 addition is commutative (only +// associativity fails), so every lane ends each level holding bit-identical +// values and lane 0 holds precisely candle's result. Lanes with row >= HC feed +// 0.0f, which is exactly candle's `shr[tid] = 0` identity padding. +// +// KILL SWITCH: ARC_NO_SINKHORN_WARP=1 restores the templated shared-memory +// kernel above. Default is this path (house policy: fast-by-default with a kill +// switch, never opt-in). The two arms are distinguishable in a profile by +// kernel name -- `arc_sinkhorn::warp_kernel` vs the anonymous-namespace +// `sinkhorn_normalize_f32_kernel` -- so which arm ran is provable from nsys +// output alone and cannot be faked by a stale object file. + +namespace arc_sinkhorn { + +template +__device__ __forceinline__ void col_sums( + const float (&r)[HC], + float (&cs)[HC], + bool active, + unsigned mask, + float eps +) { +#pragma unroll + for (int c = 0; c < HC; ++c) { + float v = __fadd_rn(0.0f, active ? r[c] : 0.0f); +#pragma unroll + for (int s = P >> 1; s > 0; s >>= 1) { + v = __fadd_rn(v, __shfl_xor_sync(mask, v, s)); + } + cs[c] = __fadd_rn(v, eps); + } +} + +// in/out: [n, HC, HC] row-major F32. One block per matrix, next_pow2(HC) +// threads -- all lanes of a single warp, so no shared memory and no barriers. +template +__global__ void warp_kernel( + const float* __restrict__ in, + float* __restrict__ out, + int n, + int iters, + float eps +) { + const int batch = blockIdx.x; + const int row = threadIdx.x; + // Uniform across the block, so it cannot desynchronise the shuffles below. + if (batch >= n) return; + + const int P = next_pow2_ce(HC); + const bool active = (row < HC); + const unsigned mask = (P >= 32) ? 0xffffffffu : ((1u << P) - 1u); + + float r[HC]; + if (active) { + const float* my_in = in + (size_t)batch * HC * HC + (size_t)row * HC; +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = my_in[j]; + } else { +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = 0.0f; + } + + // ---- 1. stable row softmax, then + eps (entirely within this lane) ---- + const float m = candle_tree_max(r); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = expf(__fsub_rn(r[j], m)); + const float rs = candle_tree_sum(r); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fadd_rn(__fdiv_rn(r[j], rs), eps); + + // ---- 2. initial column normalize: x / (colsum + eps) ---- + float cs[HC]; + col_sums(r, cs, active, mask, eps); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(r[j], cs[j]); + + // ---- 3. (iters - 1) more row->col passes ---- + for (int it = 0; it < iters - 1; ++it) { + const float rsum = __fadd_rn(candle_tree_sum(r), eps); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(r[j], rsum); + col_sums(r, cs, active, mask, eps); +#pragma unroll + for (int j = 0; j < HC; ++j) r[j] = __fdiv_rn(r[j], cs[j]); + } + + if (active) { + float* my_out = out + (size_t)batch * HC * HC + (size_t)row * HC; +#pragma unroll + for (int j = 0; j < HC; ++j) my_out[j] = r[j]; + } +} + +// Read once; getenv is not cheap and this is on the per-layer decode path. +inline bool warp_disabled() { + static const bool disabled = [] { + const char* v = getenv("ARC_NO_SINKHORN_WARP"); + return v != nullptr && v[0] == '1' && v[1] == '\0'; + }(); + return disabled; +} + +} // namespace arc_sinkhorn + +// Defined outside the namespace because it is expanded in `extern "C"` scope +// below; `warp_kernel` has to be qualified there. +#define SINKHORN_WARP_LAUNCH(HC_) \ + case HC_: \ + arc_sinkhorn::warp_kernel \ + <<>>( \ + (const float*)in, (float*)out, n, iters, eps); \ + return; + + extern "C" { void sinkhorn_normalize_f32( @@ -268,6 +401,31 @@ void sinkhorn_normalize_f32( float eps, long long stream ) { + if (!arc_sinkhorn::warp_disabled()) { + switch (hc) { + SINKHORN_WARP_LAUNCH(1) + SINKHORN_WARP_LAUNCH(2) + SINKHORN_WARP_LAUNCH(3) + SINKHORN_WARP_LAUNCH(4) + SINKHORN_WARP_LAUNCH(5) + SINKHORN_WARP_LAUNCH(6) + SINKHORN_WARP_LAUNCH(7) + SINKHORN_WARP_LAUNCH(8) + SINKHORN_WARP_LAUNCH(9) + SINKHORN_WARP_LAUNCH(10) + SINKHORN_WARP_LAUNCH(11) + SINKHORN_WARP_LAUNCH(12) + SINKHORN_WARP_LAUNCH(13) + SINKHORN_WARP_LAUNCH(14) + SINKHORN_WARP_LAUNCH(15) + SINKHORN_WARP_LAUNCH(16) + default: + // Unreachable (the Rust guard rejects hc > SINKHORN_MAX_HC); fall + // through to the templated kernel rather than leave `out` untouched. + break; + } + } + dim3 grid(n, 1, 1); dim3 block(hc, 1, 1); switch (hc) { @@ -295,5 +453,6 @@ void sinkhorn_normalize_f32( } } #undef SINKHORN_LAUNCH +#undef SINKHORN_WARP_LAUNCH } // extern "C" diff --git a/mistralrs-core/src/engine/mod.rs b/mistralrs-core/src/engine/mod.rs index c30053ca4..00495b282 100644 --- a/mistralrs-core/src/engine/mod.rs +++ b/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"); diff --git a/mistralrs-core/src/kv_cache/mod.rs b/mistralrs-core/src/kv_cache/mod.rs index b3e21b693..f21b01612 100644 --- a/mistralrs-core/src/kv_cache/mod.rs +++ b/mistralrs-core/src/kv_cache/mod.rs @@ -13,7 +13,7 @@ mod hybrid_cache; mod rotating_cache; mod single_cache; pub mod turboquant_cache; -mod xs_rolling; +pub(crate) mod xs_rolling; /// Thread-local `ARC_V4_XS_PER_SEQ` override, so a test outside this module can /// exercise both sides of the flag (the production read is a `OnceLock`). #[cfg(test)] @@ -2874,6 +2874,35 @@ mod clone_in_cache_invariant_tests { use crate::sequence::{SeqStepType, SequenceGroup, SequenceRecognizer}; use candle_core::Device; + /// Run `f` with the xs window pin OFF. + /// + /// Three tests below assert the geometry of the *unpinned* retained window + /// — exact tail widths, and that column 0 holds token `base`. Both are + /// properties of the resizing policy specifically: #121's pin holds the + /// buffer at `span_groups * ratio + margin` and leaves real tokens sitting + /// ahead of `base`, so column 0 is no longer `base` and the widths are the + /// pinned constants instead. + /// + /// The pin is opt-in (`ARC_V4_XS_PIN_WINDOW=1`), so as things stand these + /// three would pass without asking. They ask anyway, deliberately: the + /// default is expected to flip to ON once the pin is measured on this tree + /// (see `xs_rolling::xs_pin_window_enabled_from`'s FLIP CONDITION), and a + /// test that depends on a default it does not name is a test that breaks on + /// the day someone flips it. Naming the mode makes these three immune to + /// that, in either direction. + /// + /// They are not testing a dead path either way: the resizing window is a + /// supported serving mode and is one arm of the pin A/B. The pinned + /// geometry is covered separately by `xs_rolling`'s own tests and by + /// `deepseek4`'s pin A/B. + /// + /// Uses #121's `pin_test_override` — thread-local, not an env mutation, + /// because `cargo test` is multi-threaded and `xs_pin_window_enabled` + /// latches its env answer in a `OnceLock`. + fn unpinned(f: impl FnOnce() -> T) -> T { + crate::kv_cache::xs_rolling::pin_test_override::with(false, f) + } + /// A cache slot whose only interesting property is its `current_seq_len`. /// `all_data` stays `None` — `first_mismatched_cache_len` reads lengths, /// never tensors, which is exactly the point: `CACHE_GROW_SIZE = 512` means @@ -3958,8 +3987,23 @@ mod clone_in_cache_invariant_tests { /// compressor history and cannot share one dense buffer. It must be /// refused, by name — never papered over, and never a panic on the engine /// task. + /// + /// ⚠️ Run with the window pin OFF, deliberately. The widths 18 and 22 are + /// what the *resizing* buffer produced, and they are the whole + /// discriminator: pinning the window makes both 144 + /// (`span_groups * ratio + margin`), so the shape mismatch at :499 cannot + /// arise and this fixture would silently stop testing anything. The + /// historical defect is still worth pinning — the pin is a flag, and the + /// refusal is what has to hold when it is off. + /// `the_pin_removes_this_fixtures_discriminator` records the other side. #[test] fn ragged_xs_tail_is_refused_by_name_not_panicked() { + crate::kv_cache::xs_rolling::pin_test_override::with(false, || { + ragged_xs_tail_is_refused_by_name_not_panicked_inner(); + }); + } + + fn ragged_xs_tail_is_refused_by_name_not_panicked_inner() { let mut short = xs_state(128, 1); feed_xs(&mut short, 274); let mut long = xs_state(128, 1); @@ -4015,6 +4059,38 @@ mod clone_in_cache_invariant_tests { ); } + /// The other side of the fixture above: with the window pinned, the two + /// sequences' tails are the SAME width, so wave51-CB's shape mismatch is + /// not merely refused — it cannot be constructed. + /// + /// That is a consequence of the pin worth recording, not a reason to stop + /// refusing: the two caches still hold different history, and the batch is + /// still refused, just one layer earlier and by the token-count invariant + /// (`ensure_uniform_batch_cache_lens`) rather than by a tensor shape. A + /// refusal that depends on two buffers happening to differ in size is a + /// weaker guarantee than one that reads the lengths, and this shows the + /// weaker one is not what is holding. + #[test] + fn the_pin_removes_this_fixtures_discriminator() { + crate::kv_cache::xs_rolling::pin_test_override::with(true, || { + let mut short = xs_state(128, 1); + feed_xs(&mut short, 274); + let mut long = xs_state(128, 1); + feed_xs(&mut long, 278); + let (sw, lw) = ( + short.tail.as_ref().unwrap().dims()[1], + long.tail.as_ref().unwrap().dims()[1], + ); + assert_eq!( + (sw, lw), + (144, 144), + "pinned HCA windows must both be span_groups*ratio+margin = 144" + ); + // …and the lengths still disagree, which is what the refusal reads. + assert_ne!(short.current_seq_len(), long.current_seq_len()); + }); + } + /// The ragged-tail refusal must not depend on the K/V slots noticing /// first: on V4 the K/V halves are `[1, 1, capacity, head_dim]`, whose /// dim-1 is the head count, so they batch happily at *any* pair of lengths. @@ -4073,6 +4149,10 @@ mod clone_in_cache_invariant_tests { /// `shape mismatch on dim 1, 4 <> 132`. #[test] fn xs_base_divergence_at_equal_lengths_is_reconciled_not_refused() { + unpinned(xs_base_divergence_at_equal_lengths_is_reconciled_not_refused_inner); + } + + fn xs_base_divergence_at_equal_lengths_is_reconciled_not_refused_inner() { // Restored from a prefix-cache entry stored at 300 tokens, truncated to // 260 — `base` stays at canonical(300), which is past canonical(260). let mut restored = xs_state(128, 1); @@ -4189,6 +4269,10 @@ mod clone_in_cache_invariant_tests { /// looks at shapes. #[test] fn per_row_xs_bases_survive_a_batch_round_trip_and_are_not_flattened() { + unpinned(per_row_xs_bases_survive_a_batch_round_trip_and_are_not_flattened_inner); + } + + fn per_row_xs_bases_survive_a_batch_round_trip_and_are_not_flattened_inner() { xs_rolling::test_override::with(true, || { let mut restored = xs_state(128, 1); feed_xs(&mut restored, 300); @@ -4356,6 +4440,10 @@ mod clone_in_cache_invariant_tests { /// (token `t` carries the value `t`) and reads the actual numbers back. #[test] fn trimming_the_retained_window_drops_the_oldest_rows_not_the_newest() { + unpinned(trimming_the_retained_window_drops_the_oldest_rows_not_the_newest_inner); + } + + fn trimming_the_retained_window_drops_the_oldest_rows_not_the_newest_inner() { use candle_core::IndexOp; let dev = candle_core::Device::Cpu; let mut state = xs_state(128, 1); diff --git a/mistralrs-core/src/kv_cache/xs_rolling.rs b/mistralrs-core/src/kv_cache/xs_rolling.rs index fa30bed65..5e580547e 100644 --- a/mistralrs-core/src/kv_cache/xs_rolling.rs +++ b/mistralrs-core/src/kv_cache/xs_rolling.rs @@ -100,34 +100,42 @@ //! [`XsRollingCache::advance`] keeps a verbatim uniform fast path and every //! B=1 and every uniform batch runs byte-identical code. //! -//! # What `W` is now allowed to be +//! # The window is a fixed-size buffer, and `advance` was the one thing +//! # breaking that //! -//! Nothing below still requires `W == max_i (tokens_i - base_i)`. -//! [`plan_xs_advance`] and [`XsRollingCache::set_row_lens`] take `W` as *given* -//! and ask only that every row fit inside it; every read offset is -//! `need_start + W - tokens_i`, which tracks `W` exactly. So `W` may be any -//! width at or above the widest row's run — including a **constant** one. -//! `an_oversized_window_reads_exactly_the_same_absolute_tokens` pins that. +//! Read the description of `tail` at the top of this file again: *"the raw rows +//! behind `tokens`, **bounded by `span_groups * ratio + margin` and independent +//! of context length**"*. That is an invariant this type has always claimed — +//! a buffer whose size is a function of the layer's geometry and nothing else. //! -//! It is worth writing down because this buffer does **not** have a constant -//! size today. `tail` is rebuilt by `Tensor::cat` + `narrow` every step at -//! width `tokens - base`, and that width is not stable: `tokens` climbs by one -//! per token while `base` jumps a whole `ratio` at a group boundary, so the -//! allocation cycles through `ratio` consecutive sizes. The ArcGraph chain -//! measured exactly that from the outside — `4096 × {18, 19, 20, 21}` at -//! `hidden = 4096`, `ratio = 4`. (Their separate hypothesis, that this is what -//! stops a CUDA graph replaying, they have since retracted as unproven; the -//! size measurement is a fact about this buffer either way.) +//! It was not true of the allocation. `advance` rebuilt `tail` with +//! `Tensor::cat` + `narrow` on every step at width exactly `tokens - base`, and +//! that width is not stable: `tokens` climbs by one per token while `base` jumps +//! a whole `ratio` at a group boundary, so the buffer cycled through `ratio` +//! consecutive sizes and was reallocated every single decode step, on every one +//! of the 41 compressed layers. The ArcGraph chain measured exactly that shape +//! from outside the engine — `4096 × {18, 19, 20, 21}` at `hidden = 4096`, +//! `ratio = 4`. (Their further hypothesis, that this is what stops a CUDA graph +//! replaying, they have retracted as unproven. The reallocation is a per-token +//! cost on the hot path whether or not anything is capturing, which is the +//! reason it is fixed here.) //! -//! Pinning `W` to the bound the retention rule already guarantees it never -//! exceeds — `span_groups * ratio + margin` — would make the allocation -//! constant per step **without touching any of the semantics above**: a row -//! would simply carry more dead lead columns, exactly as a short row in a -//! ragged batch already does. The places that decide `W` are the retention -//! narrow at the end of [`XsRollingCache::advance_ragged`] and its counterpart -//! in [`XsRollingCache::advance_uniform`], plus the fast-path predicate -//! [`XsRollingCache::can_advance_uniform`], which reads the exact equality as -//! its dispatch condition. Nothing here depends on that change being made. +//! [`xs_pin_window_enabled`] holds the buffer at the documented bound instead. +//! It required no new semantics: [`plan_xs_advance`] and +//! [`XsRollingCache::set_row_lens`] already take `W` as *given* and ask only +//! that every row fit inside it, and every read offset is +//! `need_start + W - tokens_i`, which tracks `W` exactly. A pinned window is +//! therefore just the widest case of the ragged one that already had to work — +//! a row carrying dead lead columns it never promised. +//! `an_oversized_window_reads_exactly_the_same_absolute_tokens` pins the +//! arithmetic and `pinning_the_window_is_numerically_inert` pins the outcome. +//! +//! The two widths, reconciled: the steady-state band is +//! `[capacity - ratio, capacity - 1]`, i.e. `{20..23}` at CSA geometry, and the +//! `{18..21}` measured outside was the pre-saturation **ramp** — `base` is still +//! at 0 while `tokens` climbs, so `W` has not finished growing. +//! `the_window_ramps_then_settles_to_ratio_consecutive_sizes` checks both +//! halves of that, so the bound is right for a reason rather than by luck. use std::sync::OnceLock; @@ -236,6 +244,161 @@ pub fn xs_per_sequence_enabled() -> bool { }) } +/// `ARC_V4_XS_PIN_WINDOW=1` — hold the retained raw window at a **constant** +/// width instead of re-narrowing it to `tokens - base` every step. +/// **Default OFF.** See [`xs_pin_window_enabled_from`] for the polarity, the +/// reason, and the named measurement that flips this default back on. +/// +/// # The argument that once made this default-ON, kept because it is still true +/// +/// This shipped default-ON on #121's branch, on the reasoning below. The +/// reasoning is sound as far as it goes — it is why the pin is *safe* — but it +/// is an argument about correctness, and the default was also a claim about +/// throughput, which is what was never measured on this tree. Correctness +/// argument retained; default retracted. +/// +/// The usual rule on this project is measure-then-default-on, and it exists +/// because FP8 KV shipped defaulted-on and broke every forward for a day. The +/// difference is what the two changes touch. FP8 KV changed the *values*. This +/// changes only the *size of an allocation*: the compressor is handed a slice +/// covering the same absolute tokens either way, because every offset is +/// derived from the row's own token count rather than from the buffer's width +/// (see [`plan_xs_advance`]). `pinning_the_window_is_numerically_inert` pins +/// that as an exact-equality test rather than as an argument. +/// +/// # The bound, which is the reviewable part +/// +/// The retained width is `W = tokens - base`, and `base` is set from the +/// retention rule at `tokens - margin`, then only ever raised: +/// +/// ```text +/// keep_from = ((tokens - margin) / ratio + 1 - span_groups) * ratio +/// base = max(keep_from, previous base), capped at tokens +/// ``` +/// +/// Write `m = tokens - margin` and `q = floor(m / ratio)`. Then +/// `q * ratio > m - ratio`, so +/// +/// ```text +/// keep_from = (q + 1 - span_groups) * ratio > m - span_groups * ratio +/// W = tokens - base <= tokens - keep_from < margin + span_groups * ratio +/// ``` +/// +/// so **`W <= span_groups * ratio + margin - 1`, always**. Pinning to +/// `span_groups * ratio + margin` is therefore provably sufficient and provably +/// never truncates a row's promised history. At V4's CSA geometry +/// (`ratio 4, span_groups 2, margin 16`) that is 24, against a measured width +/// cycling through `{20, 21, 22, 23}` — `ratio` consecutive sizes, which is +/// exactly the `4096 x {18..21}`-shaped reallocation the ArcGraph chain +/// observed from outside the engine. +/// +/// This is a bound with a proof, not a chosen constant. The capacity is a +/// function of the layer's own `ratio`/`span_groups`/`margin`, so it does not +/// grow with context length. +pub fn xs_pin_window_enabled() -> bool { + #[cfg(test)] + match pin_test_override::current() { + Some(on) => return on, + None => {} + } + static ON: OnceLock = OnceLock::new(); + *ON.get_or_init(|| { + let on = xs_pin_window_enabled_from(std::env::var("ARC_V4_XS_PIN_WINDOW").ok().as_deref()); + // Named once per process, so an A/B that means to compare pinned + // against unpinned can ASSERT it got two different builds of behaviour + // rather than two identical arms reporting a 1.000x ratio. A flag whose + // name is wrong is indistinguishable from a flag that costs nothing. + tracing::info!( + target: "xs_rolling", + "xs rolling window is {} (ARC_V4_XS_PIN_WINDOW)", + if on { "PINNED" } else { "RESIZING" } + ); + on + }) +} + +/// Pure half of [`xs_pin_window_enabled`], so the polarity is testable without +/// mutating the process environment. +/// +/// # Why the pin is OPT-IN +/// +/// This shipped default-ON, with `=0` as the control arm, because that made the +/// A/B guard meaningful. It is now default-OFF, `== Some("1")`, and the A/B +/// carries the flag on its TREATMENT arm instead. +/// +/// The width bound above is proved, and the pin is very probably a win. But +/// "probably a win" is not the bar. It changes what the serving path retains — +/// at V4's HCA geometry the window is held at 144 columns where the resizing +/// policy keeps as few as 4 at some residues — and the +23.2% at uniform B=32 +/// that motivates it was measured on this change's own branch, against a tree +/// that no longer exists. That is an unmeasured default, and unmeasured +/// defaults are how this repo lost a week: TCFRAG was default-ON with +/// "UNVERIFIED ON HARDWARE — NEVER RUN" in its own header and held 63 GB +/// (#209), and the fused-512 attention path silently dropped the attention mask +/// for four days at 12% agreement. +/// +/// So: **unverified means default-off, and "unverified" means unmeasured, not +/// new.** +/// +/// # FLIP CONDITION — this default is meant to change +/// +/// Off is a **temporary state with an owner**, not the finish line. Arc's +/// larger problem is not unbuilt work, it is finished, correct, tested work +/// left switched off — a complete GPU sampler with no callers, a fused decode +/// path nobody enabled, 34 tok/s of landed work while master served 17. A flag +/// that ships off with no named experiment behind it is how that happens, so +/// here is the experiment. +/// +/// * **Gate:** `ARC_V4_XS_PIN_WINDOW`, read at +/// `mistralrs-core/src/kv_cache/xs_rolling.rs` in +/// [`xs_pin_window_enabled_from`] (this function). +/// * **The measurement that flips it:** one binary, uniform B=32 decode, same +/// prompt and same seed, two arms — `ARC_V4_XS_PIN_WINDOW=1` against the flag +/// unset. `arc-tools/arcspec_perseq_ladder.sh` already runs exactly this +/// shape and asserts, from the per-process `xs rolling window is …` log line, +/// that the two arms really differed. Pass = the pinned arm is faster on +/// aggregate tok/s, with no change in generated tokens. +/// * **What it was, on its own branch and not on this tree:** +23.2% at uniform +/// B=32. That number was taken on #121 before rebase, against a tree that no +/// longer exists. It is **not a current fact** and must not be quoted as one. +/// * **On pass, the default flips to ON in the same change that records the +/// number.** Leaving it off after the measurement succeeds is a failure +/// state, not a safe one. +/// + +/// Nothing about CUDA-graph capture depends on this default: capture asks for +/// the pin per-cache through [`XsRollingCache::pin_tail_width`], and +/// `pin_is_on` honours that trigger independently of this flag. +/// +/// Read by VALUE and not by presence — #212 converted 23 `ARC_*` flags because +/// `var_os(..).is_some()` made `ARC_FOO=0` mean ON, which silently turns an A/B +/// control arm into a second treatment arm. +pub fn xs_pin_window_enabled_from(value: Option<&str>) -> bool { + value == Some("1") +} + +/// Test-only override for [`xs_pin_window_enabled`], thread-local for the same +/// reason [`test_override`] is. +#[cfg(test)] +pub(crate) mod pin_test_override { + use std::cell::Cell; + + thread_local! { + static STATE: Cell> = const { Cell::new(None) }; + } + + pub(crate) fn current() -> Option { + STATE.with(|s| s.get()) + } + + pub(crate) fn with(on: bool, f: impl FnOnce() -> R) -> R { + let prev = STATE.with(|s| s.replace(Some(on))); + let out = f(); + STATE.with(|s| s.set(prev)); + out + } +} + /// Test-only override for [`xs_per_sequence_enabled`]. /// /// The production read is latched in a `OnceLock` on purpose — a mid-run change @@ -537,8 +700,14 @@ impl XsRollingCache { /// `r + margin + (span_groups - 1) * ratio`, maximised at `r = ratio - 1` /// as `span_groups * ratio + margin - 1`. So pinning here only ever /// retains MORE than the compressor needs, never less. + /// + /// The same quantity as [`Self::window_capacity`], which #121 derived + /// independently for the serving pin. Delegating rather than repeating the + /// expression: two pins that computed "the pinned width" separately could + /// drift, and then capture and serving would disagree about a number whose + /// whole value is that it does not move. pub fn graph_tail_width(&self) -> usize { - self.span_groups * self.ratio + self.margin + self.window_capacity() } /// Hold the retained raw tail at [`Self::graph_tail_width`] columns instead @@ -662,24 +831,103 @@ impl XsRollingCache { )) })?; let have = tail.dim(1)?; - // `want > have` cannot happen (`set_row_lens` refuses such a batch), - // but if it ever did, keeping the whole window and RAISING `base` - // degrades to "refuses a rollback it could have served" — never to - // resuming the compressor from a gap. - let want = (tok - base[i]).min(have); + // What this row actually promised. `promised > have` cannot happen + // (`set_row_lens` refuses such a batch), but if it ever did, keeping + // the whole window and RAISING `base` degrades to "refuses a rollback + // it could have served" — never to resuming the compressor from a gap. + let promised = (tok - base[i]).min(have); + // 🔑 With the window pinned, do NOT re-narrow to `promised` here. + // `clone_out_cache` calls this once per layer per sequence on EVERY + // engine step, so narrowing would reallocate the buffer every step and + // undo the pin exactly on the hot path it exists for. Keeping the + // pinned width costs nothing: the extra columns are older real tokens + // ahead of this row's `base`, and `plan_xs_advance` refuses to read + // below `base` regardless of how wide the buffer is. + // + // Clamped by `tok` so a per-sequence cache never ends up wider than the + // row is long — that state is legal but it would send every subsequent + // B=1 step down `advance_ragged`, which is the slower path. + let width = if self.pin_is_on() { + self.window_capacity().min(tok).min(have).max(promised) + } else { + promised + }; let mut out = self.clone(); out.comp.all_data = Some(comp); out.comp.current_seq_len = tok / self.ratio; - out.tail = Some(if want == have { + out.tail = Some(if width == have { tail } else { - tail.narrow(1, have - want, want)?.contiguous()? + tail.narrow(1, have - width, width)?.contiguous()? }); out.tokens = vec![tok]; - out.base = vec![tok - want]; + // The resume point is this row's own, unchanged by how wide the buffer + // it landed in happens to be. + out.base = vec![tok - promised]; Ok(out) } + /// The width the retained raw window is pinned to — `span_groups * ratio + + /// margin`, the bound derived in `xs_pin_window_enabled` that + /// `tokens - base` provably never reaches. + /// + /// (Not an intra-doc link: `XsRollingCache` is re-exported from the crate + /// root but `xs_pin_window_enabled` is not, so rustdoc reads a public item + /// linking to a private one and warns. Exporting it alongside its exact + /// analogue `xs_per_sequence_enabled` would also close this, but that is an + /// API-surface change and not this PR's business.) + /// + /// A function of the layer's own geometry, so it does not grow with context + /// length: 24 for V4's CSA layers, 144 for HCA. That is the whole point — + /// "the capacity it can never exceed" has to be a small number with a + /// proof, not a context-length ceiling. + pub fn window_capacity(&self) -> usize { + self.span_groups * self.ratio + self.margin + } + + /// The width to retain this step: `needed` when the pin is off (which is + /// bit-identical to what this always did), otherwise the pinned capacity — + /// clamped so it can never exceed what the window physically holds, and + /// never fall below what the rows actually promised. + /// + /// `w_win` matters after a rollback: `set_len` narrows the buffer, so the + /// step after one may have fewer than `capacity` columns to give. Clamping + /// there degrades to "this one step reallocates", never to a truncated + /// history — `needed <= w_win` holds because `tokens - base <= w_phys` is + /// the type's invariant and `w_win = w_phys + t_new`. + fn retained_width(&self, needed: usize, tokens_new: usize, w_win: usize) -> usize { + if !self.pin_is_on() { + return needed; + } + self.window_capacity() + .min(tokens_new) + .min(w_win) + .max(needed) + } + + /// Is the width pin on, by either of its two triggers? + /// + /// There are two, because the pin was arrived at twice for two different + /// reasons and both callers are real: + /// + /// * [`Self::pin_tail_width`] — a per-cache request, made by CUDA-graph + /// capture, for which a moving allocation size is not slow but *invalid* + /// (an allocation that misses the warm pool during capture becomes an + /// unstable graph memory node). See that method for the measured + /// H200 evidence. + /// * [`xs_pin_window_enabled`] — a process-wide serving switch, for which + /// the point is that the reallocation per decode step is simply waste. + /// + /// One width policy answers both: the width each asks for is + /// `span_groups * ratio + margin`, the same expression, so they cannot + /// disagree about what "pinned" means. Keeping both triggers rather than + /// collapsing them keeps capture working when the env switch is off, which + /// is its default. + #[inline] + fn pin_is_on(&self) -> bool { + self.pin_tail || xs_pin_window_enabled() + } + /// Physical width of the retained raw window (0 when there is none). pub fn tail_width(&self) -> Result { match self.tail.as_ref() { @@ -717,11 +965,22 @@ impl XsRollingCache { /// batch with `ARC_V4_XS_PER_SEQ` off satisfies it (one length, one base, /// and `clone_in_cache` refuses any tail-width disagreement outright), so /// the "flag off is byte-identical" claim is unchanged. + /// + /// The width is deliberately **not** part of the test. It was, while the + /// buffer was re-narrowed to `tokens - base` every step; pinning the window + /// ([`xs_pin_window_enabled`]) makes it wider than that on purpose, and the + /// scalar path now measures its offsets from the buffer's physical start, + /// so a wider buffer is no longer an assumption it makes. What it still + /// needs is that the window begins at or after token 0 — `w_phys <= tokens` + /// — because otherwise its `tokens - w_phys` would underflow. That case is + /// a batch assembled around a longer neighbour, and `advance_ragged` is the + /// path written for it. fn can_advance_uniform(&self) -> Result { let (t0, b0) = (self.tokens[0], self.base[0]); Ok(self.tokens.iter().all(|&t| t == t0) && self.base.iter().all(|&b| b == b0) - && self.tail_width()? == t0 - b0) + && b0 <= t0 + && self.tail_width()? <= t0) } pub fn reset(&mut self) { @@ -1149,10 +1408,24 @@ impl XsRollingCache { { let tokens = self.tokens[0]; let base = self.base[0]; + let w_phys = self.tail_width()?; - // `window` covers tokens [win_start, tokens + t_new). + // `window` covers tokens [win_start, tokens + t_new), where `win_start` + // is the buffer's PHYSICAL start, not `base`. + // + // Those were the same number while the buffer was re-narrowed to + // exactly `tokens - base` every step. They are not once it is pinned: a + // pinned window keeps older columns this row no longer promises, so its + // physical start sits at or below `base`. Measuring offsets from the + // physical start is what makes the two cases one code path — and it is + // the same thing the ragged path already does, where a neighbour's + // width has always been able to widen the window. + // + // `tokens - w_phys` cannot underflow here: `can_advance_uniform` + // refuses this path when `w_phys > tokens`, sending that case to + // `advance_ragged`, which handles a window wider than the row is long. let (window, win_start) = match self.tail.as_ref() { - Some(tail) if tokens > base => (Tensor::cat(&[tail, xs_new], 1)?, base), + Some(tail) if w_phys > 0 => (Tensor::cat(&[tail, xs_new], 1)?, tokens - w_phys), _ => (xs_new.clone(), tokens), }; let tokens_new = tokens + t_new; @@ -1167,11 +1440,16 @@ impl XsRollingCache { need_start, (g_done + 1).saturating_sub(self.span_groups) * self.ratio ); - if need_start < win_start { + // Bounded by `base`, not by the buffer's physical start — the same + // tightening `plan_xs_advance` carries. A pinned window physically + // holds tokens below `base`, and they are real, but they are not + // history this row promised to keep; reading them because the + // buffer happens to be wider would make the pin change answers. + if need_start < base { candle_core::bail!( "xs rolling cache: compressor history gap — row {g_done} needs tokens from \ - {need_start} but the retained window starts at {win_start}. This means a \ - truncation was applied without going through `set_len`." + {need_start} but that row only retained from {base}. This means a truncation \ + was applied without going through `set_len`." ); } let off = need_start - win_start; @@ -1209,22 +1487,32 @@ impl XsRollingCache { let rollback_floor = tokens_new.saturating_sub(self.margin); let keep_from = ((rollback_floor / self.ratio + 1).saturating_sub(self.span_groups) * self.ratio) - .max(win_start); - // Under `pin_tail_width` the retention point is the FIXED-width one - // instead, so the tail (and the cat, narrow and contiguous that build - // it) keeps one shape for the whole run — a capture precondition, see - // `pin_tail_width`. `graph_tail_width()` is an upper bound on what the - // unpinned policy would keep, so this only ever retains more, never - // less; `.max(win_start)` covers the transition step, where the window - // in hand may still be narrower than the pinned width. - let pinned_base = (self.pin_tail && tokens_new >= self.graph_tail_width()) - .then(|| (tokens_new - self.graph_tail_width()).max(win_start)); - let new_base = pinned_base.unwrap_or(keep_from).min(tokens_new); - self.tail = Some( - window - .narrow(1, new_base - win_start, tokens_new - new_base)? - .contiguous()?, - ); + .max(base); + let new_base = keep_from.min(tokens_new); + // End-anchored, which is what it always was: with the pin off, + // `w_win - needed == new_base - win_start` identically, so this is the + // same narrow the scalar path has always taken. With it on, the extra + // columns are older real tokens ahead of `new_base` that nothing reads. + // + // This supersedes the `pinned_base` form that reached the integration + // branch first (#181's `pin_tail_width`). Both hold the retained tail + // at the SAME constant width -- `graph_tail_width()` and + // `window_capacity()` are the same expression, `span_groups * ratio + + // margin` -- which is the property capture needs. They differed only in + // where the constancy came from: that one moved `base` earlier so the + // logical span `tokens - base` became constant, this one leaves `base` + // at the retention point the type documents and widens the physical + // buffer instead. This form is the one that composes, because + // `win_start` here is the buffer's PHYSICAL start (`tokens - w_phys`) + // rather than `base`, so a buffer wider than the row promised is + // already a representable state -- and `plan_xs_advance` still refuses + // to read below `base`, so the pin cannot change any answer. Both + // triggers are honoured; see `retained_width`. + let w_win = window.dim(1)?; + let needed = tokens_new - new_base; + let width = self.retained_width(needed, tokens_new, w_win); + debug_assert!(width >= needed && width <= w_win); + self.tail = Some(window.narrow(1, w_win - width, width)?.contiguous()?); for (t, ba) in self.tokens.iter_mut().zip(self.base.iter_mut()) { *t = tokens_new; *ba = new_base; @@ -1320,11 +1608,10 @@ impl XsRollingCache { // row's own token count and `plan_xs_advance` refuses any that would // reach below the row's `base`. let w_win = window.dim(1)?; - self.tail = Some( - window - .narrow(1, w_win - plan.tail_width, plan.tail_width)? - .contiguous()?, - ); + let widest_row = plan.tokens_new.iter().copied().max().unwrap_or(0); + let width = self.retained_width(plan.tail_width, widest_row, w_win); + debug_assert!(width >= plan.tail_width && width <= w_win); + self.tail = Some(window.narrow(1, w_win - width, width)?.contiguous()?); self.tokens = plan.tokens_new; self.base = plan.base_new; self.compressed_rows() @@ -1698,3 +1985,60 @@ mod tests { ); } } + +#[cfg(test)] +mod pin_gate_polarity_tests { + use super::*; + + /// The pin is OFF unless explicitly asked for. + /// + /// It shipped default-ON in #121, which made the A/B's control arm the one + /// carrying `=0`. That is the wrong way round for something unmeasured on + /// this tree: the pin changes what serving retains (144 columns at V4's HCA + /// geometry, against as few as 4 at some residues under the resizing + /// policy), and its +23.2% at uniform B=32 was measured on #121's own + /// branch. Same lesson as #209's TCFRAG retirement. + /// + /// This is NOT a parking assertion. See `xs_pin_window_enabled_from`'s doc + /// comment for the single measurement that flips this default back on. + #[test] + fn the_pin_is_off_unless_explicitly_enabled() { + assert!( + !xs_pin_window_enabled_from(None), + "unset must mean OFF — the pin is unmeasured on the current tree" + ); + assert!(xs_pin_window_enabled_from(Some("1"))); + } + + /// Read by VALUE. #212 converted 23 `ARC_*` flags because + /// `var_os(..).is_some()` made `ARC_FOO=0` mean ON — which turns an A/B's + /// control arm into a second treatment arm and reports a clean 1.000x. + #[test] + fn zero_means_off_and_so_does_every_other_value() { + for v in ["0", "", "false", "no", "off", "true", "yes", "on", "2", " 1"] { + assert!( + !xs_pin_window_enabled_from(Some(v)), + "{v:?} must not enable the pin; only the exact string \"1\" does" + ); + } + } + + /// Capture must not depend on the env default: it asks per-cache through + /// `pin_tail_width()`, and `pin_is_on` honours that trigger on its own. + /// If this regressed, CUDA-graph capture would silently go back to a + /// per-step-varying allocation size — which is not a slow graph but an + /// invalid one. + #[test] + fn capture_still_gets_a_constant_width_with_the_env_flag_off() { + pin_test_override::with(false, || { + let mut c = XsRollingCache::new(4, 2, 16, 64); + assert!(!c.pin_is_on(), "env off and pin_tail unset => not pinned"); + c.pin_tail_width(); + assert!( + c.pin_is_on(), + "capture's per-cache trigger must work with the env flag off" + ); + assert_eq!(c.graph_tail_width(), c.window_capacity()); + }); + } +} diff --git a/mistralrs-core/src/layers.rs b/mistralrs-core/src/layers.rs index 71701e852..818f235a1 100644 --- a/mistralrs-core/src/layers.rs +++ b/mistralrs-core/src/layers.rs @@ -375,12 +375,27 @@ impl QRmsNorm { /// `None` keeps the old loop verbatim, so ragged cohorts behave exactly as they /// do today. This is a dispatch fix, not a new ragged path. /// -/// # Relationship to `DeepSeekV2RotaryEmbedding::uniform_offset` +/// # The V4 rotary, collapsed onto this predicate /// -/// The V4 rotary got this same correction separately (PR #198, still unmerged -/// at the time of writing) as a private associated function with an identical -/// body. When these two lines of work meet, collapse that one onto this one -- -/// there must not be two spellings of the same predicate. +/// The V4 rotary (`DeepSeekV2RotaryEmbedding`) got this same correction +/// separately, in #198, as a private associated function `uniform_offset` with +/// a byte-identical body. #198 is now merged and that second spelling is gone: +/// `DeepSeekV2RotaryEmbedding::{forward, forward_inverse_tail}` call this +/// function directly, so the two cannot drift apart. +/// +/// Two details #198 established that this doc did not previously carry: +/// +/// * The uniformity is not an accident of the current scheduler, it is +/// *enforced* by it. `scheduler::default_scheduler::select_running_bucket` +/// admits a forward pass only when cache lengths are exactly equal +/// ("Sequences may only share a forward pass when their cache lengths are +/// exactly equal"), and the single producer of a ragged dense batch is gated +/// behind `ARC_MTP_PER_SEQ_KV`, which defaults off +/// (`models::deepseek4::ragged_row_q0`). +/// * Measured at B=256 on a real decode profile, the per-sequence loop cost +/// ~100,000 GPU launches per step, 19.9% of step time -- of which the +/// `copy2d` half existed only to undo a split the same function had just +/// made. That number was measured on #198's branch, not on this tree. #[inline] pub(crate) fn uniform_seqlen_offset(seqlen_offsets: &[usize]) -> Option { match seqlen_offsets { @@ -418,14 +433,105 @@ pub mod rope_cohort_stats { #[inline] pub(super) fn record_cohort() { COHORT.fetch_add(1, Ordering::Relaxed); + #[cfg(test)] + bump_local(1, 0); } #[inline] pub(super) fn record_per_sequence() { PER_SEQUENCE.fetch_add(1, Ordering::Relaxed); + #[cfg(test)] + bump_local(0, 1); + } + + // ---- Test-only per-thread mirror of the two counters above. ----------- + // + // The atomics are process-global and, since #198 collapsed its own private + // copy of this module onto this one, they are written by EVERY rotary in + // this file -- `RotaryEmbedding`, `Llama3RotaryEmbedding`, + // `Phi4MMRotaryEmbedding` and `DeepSeekV2RotaryEmbedding` alike. A test + // that brackets one call with `counts()` is therefore not measuring its own + // call: `cargo test` runs the suite on many threads, and any other test + // that drives any rotary lands between the `before` and the `after`. + // + // That is not hypothetical. Two rotary forwards in `deepseek4.rs` + // (`standard_layer_rope_is_unscaled_and_compressed_is_not`, + // `mtp_block_takes_the_standard_unscaled_rope_table`) race the cohort tests + // exactly this way -- caught by merging #198 and #121 together and running + // the full suite, where `--test-threads=1` passes 707 and the parallel run + // fails one. + // + // A mutex was the first fix and it is the wrong one: it obliges every + // present and future test that touches any rotary to know about a lock in + // another module, which nothing enforces and the next such test will + // forget. Counting per-thread instead makes each test's delta its own by + // construction, needs no cooperation from anyone else, and leaves the + // serving instrument -- the global `counts()` -- byte-for-byte unchanged. + + #[cfg(test)] + thread_local! { + static LOCAL: std::cell::Cell<(u64, u64)> = const { std::cell::Cell::new((0, 0)) }; + } + + #[cfg(test)] + #[inline] + fn bump_local(cohort: u64, per_sequence: u64) { + LOCAL.with(|c| { + let (a, b) = c.get(); + c.set((a + cohort, b + per_sequence)); + }); + } + + /// `(cohort, per_sequence)` for THIS thread only — the test-facing + /// counterpart of [`counts`]. + /// + /// Same monotonic-delta contract, same meaning, but immune to whatever the + /// rest of the suite is doing on other threads. Tests assert on deltas + /// around their own call; nothing resets it. + #[cfg(test)] + pub fn local_counts() -> (u64, u64) { + LOCAL.with(|c| c.get()) } } +/// `ARC_ROPE_COHORT=1` routes a **ragged** decode cohort — rows at genuinely +/// different positions — through ONE gathered `rope_i` per projection instead +/// of the per-sequence loop. **Default OFF**; unset and every off-spelling +/// keep the loop byte-for-byte as it is today. +/// +/// Uniform cohorts are NOT behind this flag: when every row shares an offset, +/// [`uniform_seqlen_offset`] already dispatches one batched `rope_i` with a +/// 2-D table, which is bit-identical to the loop and needs no gather. This +/// flag governs only the arm the uniform fix could not reach — offsets that +/// genuinely differ, where the loop was the only spelling. On the measured +/// B=256 profile that loop is ~9 ops × 43 layers × B ≈ 99,000 launches per +/// step (19.9% of step time); the gathered path is ~5 launches per layer. +/// +/// candle's `rope_i` accepts a rank-3 `[B, T, D/2]` cos/sin when `B` matches +/// the input batch (`rope_check_cs`, `candle-nn/src/rotary_emb.rs`; the CUDA +/// wrapper sets `stride_b`, the CPU path indexes `i + b_i * t * d / 2`), so a +/// per-row table is expressible in one launch — the loop was never structural. +/// +/// Why opt-in rather than default-ON: the same doctrine as +/// `cuda::qk_norm_rope::fused_cohort_enabled` — unverified means default-off, +/// and unverified means unmeasured on hardware, not new. The A/B box flips +/// this flag; when the gathered arm has been measured correct and faster, the +/// default flips in the same change that records the number. +/// +/// Read by VALUE (`mistralrs_quant::env_flag_is_set`), never by presence: +/// `ARC_ROPE_COHORT=0` means OFF (#212 converted 23 flags because +/// `var_os(..).is_some()` made `=0` mean ON). Latched once per process — the +/// dispatch runs on every layer of every step and `std::env::var` takes a +/// global lock. The polarity table lives with `env_flag_value`'s own tests; +/// the two ragged arms themselves are compared bit-for-bit in +/// `deepseek_ragged_cohort_tests`, which drives them directly so the latch +/// cannot make the parity claim untestable. +pub(crate) fn rope_cohort_enabled() -> bool { + use std::sync::OnceLock; + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| mistralrs_quant::env_flag_is_set("ARC_ROPE_COHORT")) +} + /// Tile a `[max_seq_len, rot_dim]` rotary table into the `[b * seq_len, rot_dim]` /// layout `mistralrs_quant::rotary::apply_rotary_inplace` requires. /// @@ -1540,8 +1646,52 @@ impl Qwen2_5VLRotaryEmbedding { pub struct DeepSeekV2RotaryEmbedding { sin: Tensor, cos: Tensor, + /// `-sin`, materialised once at construction. + /// + /// [`Self::forward_inverse_tail`] rotates by the conjugate `(cos, -sin)`. + /// It used to spell that as `self.sin.narrow(..)?.neg()?` *inside* the + /// per-sequence loop, which at B=256 × 43 layers launched `uneg_bf16` + /// **11,008 times per decode step** to negate a `[1, 32]` slice of a + /// **constant** table — 0.8% of a measured B=256 step spent re-deriving a + /// value that never changes. + /// + /// Narrowing a pre-negated table is bit-identical to negating a narrowed + /// one: negation is a sign-bit flip, exact in every float format and + /// elementwise, so it commutes with a view. See `neg_sin_commutes_with_narrow`. + neg_sin: Tensor, + /// Memoised per-row table gather for the ragged-cohort path + /// ([`rope_cohort_enabled`]). Behind an `Arc` so the per-device sharing in + /// `deepseek4.rs` (one rotary per device per table kind, cloned into every + /// layer on that device) shares ONE cache: the offsets are host data, so + /// one small H2D per step per device is the floor, and this cache holds it + /// there instead of paying it again on all 43 layers × 3 rope calls + /// (CLAUDE.md pitfall #5 — no `Tensor::from_vec` per layer on the hot + /// path). `None` until the gated path first runs; the flag-off paths never + /// touch it. + cohort_gather: Arc>>, +} + +/// One step's gathered `[B, T, rope_dim / 2]` tables for a ragged cohort, +/// keyed by the exact `(seqlen_offsets, seq_len)` that built them. All three +/// tables are gathered together: `forward` reads (cos, sin), +/// `forward_inverse_tail` reads (cos, neg_sin), and they share the step. +#[derive(Debug)] +struct CohortGather { + offsets: Vec, + seq_len: usize, + /// `[B, T, rope_dim / 2]` each, contiguous, on the tables' device. + cos: Tensor, + sin: Tensor, + neg_sin: Tensor, } +// The engagement counters for this path are the ones declared beside +// `uniform_seqlen_offset` (`rope_cohort_stats`, above). #198 carried a second +// module of the same name, with the same two atomics and a `counts()` of its +// own; two `pub mod rope_cohort_stats` in one file is a hard name collision, +// and two counters for one claim would let a harness read `cohort > 0` off the +// Phi path while this one never engaged. Collapsed onto the shared module, as +// `uniform_seqlen_offset`'s doc comment asks. #[derive(Debug, Clone, Deserialize, Serialize)] #[serde(untagged)] pub enum DeepSeekV2RopeScaling { @@ -1583,6 +1733,183 @@ pub struct DeepSeekV2RopeConfig { } impl DeepSeekV2RotaryEmbedding { + /// The single place the tables are stored, so `neg_sin` cannot be derived + /// in one constructor and forgotten in the other. + fn from_tables(sin: Tensor, cos: Tensor) -> Result { + let neg_sin = sin.neg()?; + Ok(Self { + sin, + cos, + neg_sin, + cohort_gather: Arc::new(std::sync::Mutex::new(None)), + }) + } + + /// The step's `(cos, sin, neg_sin)` — each `[B, T, rope_dim / 2]` — for a + /// ragged cohort, one table row per `(b, t)` at absolute position + /// `seqlen_offsets[b] + t`. + /// + /// The row indices are host data (`seqlen_offsets` is a host slice), so + /// one H2D copy per step is the floor; everything after it — the three + /// `index_select`s included — happens once per step per device and is + /// served from [`Self::cohort_gather`] to the other rope calls of the same + /// step (all layers on this device share one instance of `self`; see the + /// field's doc). `index_select` output is contiguous and `reshape` of a + /// contiguous tensor is a view, so the returned tables satisfy `rope_i`'s + /// contiguity requirement with no extra copy. + fn cohort_tables( + &self, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor, Tensor)> { + let mut cache = self + .cohort_gather + .lock() + .expect("cohort gather cache poisoned"); + if let Some(entry) = cache.as_ref() { + if entry.seq_len == seq_len && entry.offsets == seqlen_offsets { + return Ok((entry.cos.clone(), entry.sin.clone(), entry.neg_sin.clone())); + } + } + // Host-side bounds guard, BEFORE any index reaches `index_select`. The + // loop arm's `narrow(0, offset, seq_len)` hard-errors past the table + // end; `index_select` with an out-of-range index is not reliably + // bounds-checked on the CUDA backend and can read garbage SILENTLY. + // Both arms must fail loud, and identically-loud — the offsets are + // host data, so the check is free. + let table_rows = self.cos.dim(0)?; + if let Some(&worst) = seqlen_offsets.iter().max() { + if worst + seq_len > table_rows { + candle_core::bail!( + "rope cohort gather out of bounds: offset {worst} + seq_len {seq_len} \ + exceeds the {table_rows}-row cos/sin table (offsets {seqlen_offsets:?})" + ); + } + } + let b = seqlen_offsets.len(); + let mut rows: Vec = Vec::with_capacity(b * seq_len); + for &off in seqlen_offsets { + for t in 0..seq_len { + rows.push((off + t) as u32); + } + } + let idx = Tensor::from_vec(rows, b * seq_len, self.cos.device())?; + let half = self.cos.dim(D::Minus1)?; + let cos = self + .cos + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + let sin = self + .sin + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + let neg_sin = self + .neg_sin + .index_select(&idx, 0)? + .reshape((b, seq_len, half))?; + *cache = Some(CohortGather { + offsets: seqlen_offsets.to_vec(), + seq_len, + cos: cos.clone(), + sin: sin.clone(), + neg_sin: neg_sin.clone(), + }); + Ok((cos, sin, neg_sin)) + } + + /// The gated ragged arm of [`Self::forward`]: ONE `rope_i` per projection + /// over the whole cohort, with the per-row rotation expressed through a + /// rank-3 `[B, T, D/2]` cos/sin (see [`rope_cohort_enabled`] for why that + /// is expressible at all). Bit-identical to [`Self::forward_ragged_loop`]: + /// the gather copies table rows verbatim, and each output element is the + /// same two-multiply-one-add of the same inputs — batching changes which + /// launch computes an element, never the arithmetic. Asserted bit-for-bit + /// in `deepseek_ragged_cohort_tests`. + fn forward_ragged_cohort( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor)> { + let (cos, sin, _neg_sin) = self.cohort_tables(seqlen_offsets, seq_len)?; + let q_embed = candle_nn::rotary_emb::rope_i(&q.contiguous()?, &cos, &sin)?; + let k_embed = candle_nn::rotary_emb::rope_i(&k.contiguous()?, &cos, &sin)?; + Ok((q_embed, k_embed)) + } + + /// The per-sequence loop, extracted verbatim so both [`Self::forward`] + /// arms name the thing they dispatch between and the parity tests can + /// drive each directly whatever `ARC_ROPE_COHORT` latched to. + fn forward_ragged_loop( + &self, + q: &Tensor, + k: &Tensor, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result<(Tensor, Tensor)> { + let mut q_embeds = Vec::new(); + let mut k_embeds = Vec::new(); + for (i, offset) in seqlen_offsets.iter().enumerate() { + let cos = self.cos.narrow(0, *offset, seq_len)?; + let sin = self.sin.narrow(0, *offset, seq_len)?; + let q_embed = + candle_nn::rotary_emb::rope_i(&q.i(i)?.unsqueeze(0)?.contiguous()?, &cos, &sin)?; + let k_embed = + candle_nn::rotary_emb::rope_i(&k.i(i)?.unsqueeze(0)?.contiguous()?, &cos, &sin)?; + q_embeds.push(q_embed); + k_embeds.push(k_embed); + } + Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + } + + /// Gated ragged arm of [`Self::forward_inverse_tail`]: one `rope_i` over + /// the whole cohort, conjugate rotation via the pre-negated table gathered + /// per row. Returns only the rotated tail; the caller re-attaches the NoPE + /// dims (shared with the loop arm, so the `cat` is spelled once). + fn inverse_tail_ragged_cohort( + &self, + x: &Tensor, + rope_dim: usize, + nope: usize, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result { + let (cos, _sin, neg_sin) = self.cohort_tables(seqlen_offsets, seq_len)?; + let x_pe = x.narrow(3, nope, rope_dim)?.contiguous()?; + candle_nn::rotary_emb::rope_i(&x_pe, &cos, &neg_sin) + } + + /// Per-sequence loop arm of [`Self::forward_inverse_tail`], extracted + /// verbatim (minus the trailing NoPE `cat`, which the caller owns). + fn inverse_tail_ragged_loop( + &self, + x: &Tensor, + rope_dim: usize, + nope: usize, + seqlen_offsets: &[usize], + seq_len: usize, + ) -> Result { + let mut outs = Vec::new(); + for (i, offset) in seqlen_offsets.iter().enumerate() { + let cos = self.cos.narrow(0, *offset, seq_len)?; + let sin = self.neg_sin.narrow(0, *offset, seq_len)?; + let x_pe = x + .i(i)? + .unsqueeze(0)? + .narrow(3, nope, rope_dim)? + .contiguous()?; + outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); + } + Tensor::cat(&outs, 0) + } + + // The cohort predicate for this rotary is the free function + // `uniform_seqlen_offset`, not a method here. #198 introduced + // `Self::uniform_offset` with a byte-identical body; per that function's + // doc comment ("there must not be two spellings of the same predicate") + // the duplicate is collapsed away, and the reasoning it carried has moved + // onto the survivor. fn new_unscaled(cfg: &DeepSeekV2RopeConfig, dtype: DType, dev: &Device) -> Result { let max_seq_len = cfg.max_position_embeddings; let dim = cfg.qk_rope_head_dim; @@ -1601,7 +1928,7 @@ impl DeepSeekV2RotaryEmbedding { let sin = freqs.sin()?.to_dtype(dtype)?; let cos = freqs.cos()?.to_dtype(dtype)?; - Ok(Self { sin, cos }) + Self::from_tables(sin, cos) } fn yarn_find_correction_dim( @@ -1693,7 +2020,7 @@ impl DeepSeekV2RotaryEmbedding { let sin = (freqs.sin()? * mscale as f64)?.to_dtype(dtype)?; let cos = (freqs.cos()? * mscale as f64)?.to_dtype(dtype)?; - Ok(Self { sin, cos }) + Self::from_tables(sin, cos) } pub fn new(cfg: &DeepSeekV2RopeConfig, dtype: DType, dev: &Device) -> Result { @@ -1746,32 +2073,32 @@ impl DeepSeekV2RotaryEmbedding { ) -> Result<(Tensor, Tensor)> { let (_b_sz, _h, seq_len, _n_embd) = q.dims4()?; - if seqlen_offsets.len() == 1 { - let cos = self.cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = self.sin.narrow(0, seqlen_offsets[0], seq_len)?; + // One `rope_i` over the whole cohort whenever every row rotates at the + // same position — which is every request today. See + // [`uniform_seqlen_offset`] for why this used to read `len() == 1`, and + // what that cost. + // + // Bit-identical to the loop below, not merely close: with a 2-D + // `[T, D/2]` cos/sin the kernel's `stride_b` is 0 + // (`candle-nn/src/rotary_emb.rs`), so every batch row reads the same + // cos/sin row, and each output element is an independent + // two-multiply-one-add of its own inputs. Batching changes which + // launch computes an element, never the arithmetic that produces it. + if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.sin.narrow(0, offset, seq_len)?; let q_embed = candle_nn::rotary_emb::rope_i(&q.contiguous()?, &cos, &sin)?; let k_embed = candle_nn::rotary_emb::rope_i(&k.contiguous()?, &cos, &sin)?; Ok((q_embed, k_embed)) + } else if rope_cohort_enabled() { + // Ragged rows, one gathered rope_i per projection — opt-in, see + // [`rope_cohort_enabled`]. + rope_cohort_stats::record_cohort(); + self.forward_ragged_cohort(q, k, seqlen_offsets, seq_len) } else { - let mut q_embeds = Vec::new(); - let mut k_embeds = Vec::new(); - for (i, offset) in seqlen_offsets.iter().enumerate() { - let cos = self.cos.narrow(0, *offset, seq_len)?; - let sin = self.sin.narrow(0, *offset, seq_len)?; - let q_embed = candle_nn::rotary_emb::rope_i( - &q.i(i)?.unsqueeze(0)?.contiguous()?, - &cos, - &sin, - )?; - let k_embed = candle_nn::rotary_emb::rope_i( - &k.i(i)?.unsqueeze(0)?.contiguous()?, - &cos, - &sin, - )?; - q_embeds.push(q_embed); - k_embeds.push(k_embed); - } - Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + rope_cohort_stats::record_per_sequence(); + self.forward_ragged_loop(q, k, seqlen_offsets, seq_len) } } @@ -1789,24 +2116,23 @@ impl DeepSeekV2RotaryEmbedding { let (_b, _h, seq_len, head_dim) = x.dims4()?; let nope = head_dim - rope_dim; let x_nope = x.narrow(3, 0, nope)?; - let rotated = if seqlen_offsets.len() == 1 { - let cos = self.cos.narrow(0, seqlen_offsets[0], seq_len)?; - let sin = self.sin.narrow(0, seqlen_offsets[0], seq_len)?.neg()?; + // Same cohort dispatch as `forward`, and the conjugate `sin` now comes + // from the pre-negated table (see the `neg_sin` field) rather than a + // `neg` launch per sequence per layer. + let rotated = if let Some(offset) = uniform_seqlen_offset(seqlen_offsets) { + rope_cohort_stats::record_cohort(); + let cos = self.cos.narrow(0, offset, seq_len)?; + let sin = self.neg_sin.narrow(0, offset, seq_len)?; let x_pe = x.narrow(3, nope, rope_dim)?.contiguous()?; candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)? + } else if rope_cohort_enabled() { + // Ragged rows, one gathered rope_i — opt-in, see + // [`rope_cohort_enabled`]. + rope_cohort_stats::record_cohort(); + self.inverse_tail_ragged_cohort(x, rope_dim, nope, seqlen_offsets, seq_len)? } else { - let mut outs = Vec::new(); - for (i, offset) in seqlen_offsets.iter().enumerate() { - let cos = self.cos.narrow(0, *offset, seq_len)?; - let sin = self.sin.narrow(0, *offset, seq_len)?.neg()?; - let x_pe = x - .i(i)? - .unsqueeze(0)? - .narrow(3, nope, rope_dim)? - .contiguous()?; - outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); - } - Tensor::cat(&outs, 0)? + rope_cohort_stats::record_per_sequence(); + self.inverse_tail_ragged_loop(x, rope_dim, nope, seqlen_offsets, seq_len)? }; Tensor::cat(&[&x_nope, &rotated], 3)?.contiguous() } @@ -4110,9 +4436,9 @@ mod rope_cohort_tests { let k = xs(&[B, H, T, D_HEAD], 0.019); let offsets = vec![5usize; B]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); // A green result must prove work happened: the batched arm has to have // been the one that ran. @@ -4143,9 +4469,9 @@ mod rope_cohort_tests { let k = xs(&[3, H, T, D_HEAD], 0.019); let offsets = vec![0usize, 4, 9]; - let before = rope_cohort_stats::counts(); + let before = rope_cohort_stats::local_counts(); let (qb, kb) = re.forward(&q, &k, &offsets).unwrap(); - let after = rope_cohort_stats::counts(); + let after = rope_cohort_stats::local_counts(); assert!( after.1 > before.1, "ragged cohort must NOT take the batched path" @@ -4243,3 +4569,507 @@ mod rope_cohort_tests { } } } + +/// Bit-exactness and engagement tests for the DeepSeek V2/V4 RoPE cohort path. +/// +/// The change under test replaced `seqlen_offsets.len() == 1` — a test of the +/// *length* of the offset vector — with a test of the *distinctness of its +/// values*. At B=256 the length is 256, so the batched path was unreachable at +/// every batch size above one, while the values were uniform anyway; the +/// per-sequence loop was therefore performing B bit-identical recomputations +/// and concatenating them back together (~100,000 GPU launches per decode step +/// on a measured profile, 19.9% of step time). +/// +/// Each test below holds the old loop as a verbatim oracle and asserts the new +/// dispatch reproduces it **bit for bit**, not within a tolerance — a +/// tolerance-based assertion here would be vacuous, since the claim is that no +/// arithmetic changed at all. +#[cfg(test)] +mod deepseek_rope_cohort_tests { + use super::*; + + const MAX_POS: usize = 32; + const ROPE_DIM: usize = 8; + const HEAD_DIM: usize = 12; + const N_HEADS: usize = 3; + const BATCH: usize = 4; + + // Engagement is asserted through `rope_cohort_stats::local_counts()`, the + // per-thread mirror of the global counters. #198 originally serialised + // these tests on a private mutex, which worked only while this module also + // carried its own private copy of `rope_cohort_stats`; once the two counter + // modules were collapsed the atomics became shared and a private lock + // stopped serialising anything that mattered. A shared lock was the next + // attempt and it is still wrong -- it obliges unrelated tests in other + // modules to know about it. See `rope_cohort_stats` for why per-thread + // counting is the fix. + + fn rope() -> DeepSeekV2RotaryEmbedding { + let cfg = DeepSeekV2RopeConfig { + rope_scaling: None, + max_position_embeddings: MAX_POS, + rope_theta: 10000.0, + qk_rope_head_dim: ROPE_DIM, + }; + DeepSeekV2RotaryEmbedding::new(&cfg, DType::F32, &Device::Cpu).unwrap() + } + + /// Deterministic, non-degenerate input. Both paths consume the same tensor, + /// so randomness would be sound too, but a fixed pattern makes a failure + /// reproducible. + fn xs(b: usize, h: usize, t: usize, d: usize) -> Tensor { + let n = b * h * t * d; + let v: Vec = (0..n) + .map(|i| ((i as f32) * 0.37).sin() * 1.7 + 0.11) + .collect(); + Tensor::from_vec(v, (b, h, t, d), &Device::Cpu).unwrap() + } + + /// Raw IEEE-754 bit patterns. Float equality would accept `-0.0 == 0.0` and + /// silently pass a sign error in the conjugate `sin`; bit patterns do not. + fn bits(t: &Tensor) -> Vec { + t.flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + .into_iter() + .map(f32::to_bits) + .collect() + } + + /// The pre-fix `forward`, verbatim, as the oracle. Always loops, whatever + /// the offsets are. + fn reference_forward( + re: &DeepSeekV2RotaryEmbedding, + q: &Tensor, + k: &Tensor, + offsets: &[usize], + ) -> Result<(Tensor, Tensor)> { + let (_b, _h, seq_len, _n) = q.dims4()?; + let mut q_embeds = Vec::new(); + let mut k_embeds = Vec::new(); + for (i, offset) in offsets.iter().enumerate() { + let cos = re.cos.narrow(0, *offset, seq_len)?; + let sin = re.sin.narrow(0, *offset, seq_len)?; + q_embeds.push(candle_nn::rotary_emb::rope_i( + &q.i(i)?.unsqueeze(0)?.contiguous()?, + &cos, + &sin, + )?); + k_embeds.push(candle_nn::rotary_emb::rope_i( + &k.i(i)?.unsqueeze(0)?.contiguous()?, + &cos, + &sin, + )?); + } + Ok((Tensor::cat(&q_embeds, 0)?, Tensor::cat(&k_embeds, 0)?)) + } + + /// The pre-fix `forward_inverse_tail`, verbatim — including the old + /// `sin.narrow(..).neg()` spelling, so this oracle also pins the `neg_sin` + /// precomputation. + fn reference_inverse_tail( + re: &DeepSeekV2RotaryEmbedding, + x: &Tensor, + rope_dim: usize, + offsets: &[usize], + ) -> Result { + let (_b, _h, seq_len, head_dim) = x.dims4()?; + let nope = head_dim - rope_dim; + let x_nope = x.narrow(3, 0, nope)?; + let mut outs = Vec::new(); + for (i, offset) in offsets.iter().enumerate() { + let cos = re.cos.narrow(0, *offset, seq_len)?; + let sin = re.sin.narrow(0, *offset, seq_len)?.neg()?; + let x_pe = x + .i(i)? + .unsqueeze(0)? + .narrow(3, nope, rope_dim)? + .contiguous()?; + outs.push(candle_nn::rotary_emb::rope_i(&x_pe, &cos, &sin)?); + } + let rotated = Tensor::cat(&outs, 0)?; + Tensor::cat(&[&x_nope, &rotated], 3)?.contiguous() + } + + // The predicate itself is tested once, in `rope_cohort_tests` + // (`uniform_offset_tests_values_not_length`), against the same six cases. + // #198 carried a byte-identical copy of that test aimed at + // `DeepSeekV2RotaryEmbedding::uniform_offset`; with the two spellings + // collapsed onto `uniform_seqlen_offset` there is one predicate, so a + // second copy would assert nothing the first does not. What is NOT + // duplicated, and stays below, is the V4-specific behaviour: `neg_sin`, + // `forward_inverse_tail`, and the bit-exactness of this rotary's own + // cohort dispatch. + #[test] + fn neg_sin_commutes_with_narrow() { + let re = rope(); + for offset in [0usize, 1, 17] { + let precomputed = re.neg_sin.narrow(0, offset, 3).unwrap(); + let on_the_fly = re.sin.narrow(0, offset, 3).unwrap().neg().unwrap(); + assert_eq!( + bits(&precomputed), + bits(&on_the_fly), + "pre-negating the table must be bit-identical at offset {offset}" + ); + } + } + + #[test] + fn cohort_forward_is_bit_identical_to_the_per_sequence_loop() { + let re = rope(); + // seq_len 1 is decode; seq_len 2 covers uniform-offset chunked prefill. + for seq_len in [1usize, 2] { + let q = xs(BATCH, N_HEADS, seq_len, ROPE_DIM); + let k = xs(BATCH, 1, seq_len, ROPE_DIM); + let offsets = vec![5usize; BATCH]; + + let before = rope_cohort_stats::local_counts(); + let (q_new, k_new) = re.forward(&q, &k, &offsets).unwrap(); + let after = rope_cohort_stats::local_counts(); + let (q_ref, k_ref) = reference_forward(&re, &q, &k, &offsets).unwrap(); + + assert_eq!(q_new.dims(), q_ref.dims()); + assert_eq!(bits(&q_new), bits(&q_ref), "Q differs at seq_len={seq_len}"); + assert_eq!(bits(&k_new), bits(&k_ref), "K differs at seq_len={seq_len}"); + // Prove the batched branch is the one that ran: a fast path that + // silently declined would pass the equality above trivially. + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (1, 0), + "expected exactly one cohort call and no per-sequence loop" + ); + } + } + + #[test] + fn ragged_forward_still_takes_the_loop_and_still_matches() { + let re = rope(); + let q = xs(BATCH, N_HEADS, 1, ROPE_DIM); + let k = xs(BATCH, 1, 1, ROPE_DIM); + let offsets = vec![5usize, 6, 5, 9]; + + let before = rope_cohort_stats::local_counts(); + let (q_new, k_new) = re.forward(&q, &k, &offsets).unwrap(); + let after = rope_cohort_stats::local_counts(); + let (q_ref, k_ref) = reference_forward(&re, &q, &k, &offsets).unwrap(); + + assert_eq!(bits(&q_new), bits(&q_ref)); + assert_eq!(bits(&k_new), bits(&k_ref)); + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (0, 1), + "ragged rows must still take the per-sequence loop" + ); + } + + #[test] + fn cohort_inverse_tail_is_bit_identical_to_the_per_sequence_loop() { + let re = rope(); + for seq_len in [1usize, 2] { + let x = xs(BATCH, N_HEADS, seq_len, HEAD_DIM); + let offsets = vec![5usize; BATCH]; + + let before = rope_cohort_stats::local_counts(); + let new = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + let after = rope_cohort_stats::local_counts(); + let reference = reference_inverse_tail(&re, &x, ROPE_DIM, &offsets).unwrap(); + + assert_eq!(new.dims(), reference.dims()); + assert_eq!( + bits(&new), + bits(&reference), + "inverse tail differs at seq_len={seq_len}" + ); + assert_eq!((after.0 - before.0, after.1 - before.1), (1, 0)); + } + } + + #[test] + fn ragged_inverse_tail_still_takes_the_loop_and_still_matches() { + let re = rope(); + let x = xs(BATCH, N_HEADS, 1, HEAD_DIM); + let offsets = vec![5usize, 6, 5, 9]; + + let before = rope_cohort_stats::local_counts(); + let new = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + let after = rope_cohort_stats::local_counts(); + let reference = reference_inverse_tail(&re, &x, ROPE_DIM, &offsets).unwrap(); + + assert_eq!(bits(&new), bits(&reference)); + assert_eq!((after.0 - before.0, after.1 - before.1), (0, 1)); + } + + /// The comparator must be able to fail. Every bit-equality assertion above + /// is worthless if `bits` had degenerated into something that always + /// matches, so poison one element by a single ULP and require it flagged. + #[test] + fn the_bit_comparator_is_live() { + // Drives a forward, so it must serialise with the counter assertions. + let re = rope(); + let x = xs(BATCH, N_HEADS, 1, HEAD_DIM); + let offsets = vec![5usize; BATCH]; + let good = re.forward_inverse_tail(&x, ROPE_DIM, &offsets).unwrap(); + + let mut poisoned = bits(&good); + poisoned[0] += 1; // one ULP + assert_ne!( + bits(&good), + poisoned, + "the comparator accepted a one-ULP difference; every other \ + assertion in this module would be vacuous" + ); + } +} + +/// Parity for the **gated** ragged-cohort arms ([`rope_cohort_enabled`], +/// `ARC_ROPE_COHORT`): one gathered rank-3 `rope_i` against the per-sequence +/// loop it replaces, bit for bit. +/// +/// The arms are driven **directly**, not through the public dispatch: the flag +/// latches once per process (`OnceLock`), so a test that reached the gathered +/// arm by mutating the environment would prove only whatever the first reader +/// happened to see — the same reasoning as `env_flag_value`'s own tests. The +/// dispatch's default (flag unset ⇒ ragged rows keep the loop) is pinned by +/// `deepseek_rope_cohort_tests::ragged_forward_still_takes_the_loop_and_still_matches`, +/// which runs the public `forward` in exactly that environment. +#[cfg(test)] +mod deepseek_ragged_cohort_tests { + use super::*; + + const MAX_POS: usize = 32; + const ROPE_DIM: usize = 8; + const HEAD_DIM: usize = 12; + const N_HEADS: usize = 3; + + fn rope() -> DeepSeekV2RotaryEmbedding { + let cfg = DeepSeekV2RopeConfig { + rope_scaling: None, + max_position_embeddings: MAX_POS, + rope_theta: 10000.0, + qk_rope_head_dim: ROPE_DIM, + }; + DeepSeekV2RotaryEmbedding::new(&cfg, DType::F32, &Device::Cpu).unwrap() + } + + fn xs(b: usize, h: usize, t: usize, d: usize) -> Tensor { + let n = b * h * t * d; + let v: Vec = (0..n) + .map(|i| ((i as f32) * 0.29).sin() * 1.3 + 0.07) + .collect(); + Tensor::from_vec(v, (b, h, t, d), &Device::Cpu).unwrap() + } + + /// Raw IEEE-754 bit patterns — the claim is that no arithmetic changed at + /// all, so a tolerance would be vacuous and `-0.0 == 0.0` would slip a + /// sign error through. + fn bits(t: &Tensor) -> Vec { + t.flatten_all() + .unwrap() + .to_vec1::() + .unwrap() + .into_iter() + .map(f32::to_bits) + .collect() + } + + /// Mixed offsets, B ∈ {1, 3}, decode (T=1) and uniform-length chunk (T=2): + /// the gathered arm must reproduce the loop bit for bit. + #[test] + fn ragged_cohort_forward_matches_the_loop_bitwise() { + let re = rope(); + for (offsets, t) in [ + (vec![0usize, 4, 9], 1usize), + (vec![0usize, 4, 9], 2), + (vec![9usize, 4, 0], 1), // descending — order must be per-row + (vec![7usize], 1), // B = 1 through the same arms + ] { + let b = offsets.len(); + let q = xs(b, N_HEADS, t, ROPE_DIM); + let k = xs(b, 1, t, ROPE_DIM); + + let (q_c, k_c) = re.forward_ragged_cohort(&q, &k, &offsets, t).unwrap(); + let (q_l, k_l) = re.forward_ragged_loop(&q, &k, &offsets, t).unwrap(); + + assert_eq!(q_c.dims(), q_l.dims()); + assert_eq!( + bits(&q_c), + bits(&q_l), + "gathered Q differs from the loop for offsets {offsets:?}, T={t}" + ); + assert_eq!( + bits(&k_c), + bits(&k_l), + "gathered K differs from the loop for offsets {offsets:?}, T={t}" + ); + + // Non-vacuity: rows at different offsets must actually rotate + // differently, or the equality above would pass under a + // broadcast-one-row bug on BOTH sides. + if b > 1 { + let d01 = (q_c.i(0).unwrap() - q_c.i(1).unwrap()) + .unwrap() + .abs() + .unwrap() + .max_all() + .unwrap() + .to_scalar::() + .unwrap(); + assert!(d01 > 0.0, "rows at different offsets rotated identically"); + } + } + } + + /// Same parity for the conjugate (inverse-tail) rotation, and for the + /// fully assembled `[NoPE | rotated]` output. + #[test] + fn ragged_cohort_inverse_tail_matches_the_loop_bitwise() { + let re = rope(); + let nope = HEAD_DIM - ROPE_DIM; + for (offsets, t) in [ + (vec![0usize, 4, 9], 1usize), + (vec![0usize, 4, 9], 2), + (vec![7usize], 1), + ] { + let b = offsets.len(); + let x = xs(b, N_HEADS, t, HEAD_DIM); + + let tail_c = re + .inverse_tail_ragged_cohort(&x, ROPE_DIM, nope, &offsets, t) + .unwrap(); + let tail_l = re + .inverse_tail_ragged_loop(&x, ROPE_DIM, nope, &offsets, t) + .unwrap(); + assert_eq!(tail_c.dims(), tail_l.dims()); + assert_eq!( + bits(&tail_c), + bits(&tail_l), + "gathered inverse tail differs from the loop for offsets {offsets:?}, T={t}" + ); + + // The assembled output the model actually consumes. + let x_nope = x.narrow(3, 0, nope).unwrap(); + let full_c = Tensor::cat(&[&x_nope, &tail_c], 3) + .unwrap() + .contiguous() + .unwrap(); + let full_l = Tensor::cat(&[&x_nope, &tail_l], 3) + .unwrap() + .contiguous() + .unwrap(); + assert_eq!(bits(&full_c), bits(&full_l)); + } + } + + /// The gather memo must serve repeats AND rebuild when the cohort moves — + /// a stale entry would rotate every layer of the next step at the previous + /// step's positions, which is exactly the class of wrong-but-plausible + /// output that never crashes. + #[test] + fn cohort_gather_rebuilds_when_the_cohort_moves() { + let re = rope(); + let t = 1usize; + + for offsets in [vec![0usize, 4, 9], vec![1usize, 5, 10], vec![2usize, 3, 4]] { + // Twice with the same key: the second call is the memo hit. + for pass in 0..2 { + let (cos, sin, neg_sin) = re.cohort_tables(&offsets, t).unwrap(); + assert_eq!(cos.dims(), &[offsets.len(), t, ROPE_DIM / 2]); + for (b_i, &off) in offsets.iter().enumerate() { + let want_cos = re.cos.narrow(0, off, t).unwrap(); + let want_sin = re.sin.narrow(0, off, t).unwrap(); + let want_neg = re.neg_sin.narrow(0, off, t).unwrap(); + assert_eq!( + bits(&cos.i(b_i).unwrap()), + bits(&want_cos), + "cos row {b_i} wrong for offsets {offsets:?} (pass {pass})" + ); + assert_eq!(bits(&sin.i(b_i).unwrap()), bits(&want_sin)); + assert_eq!(bits(&neg_sin.i(b_i).unwrap()), bits(&want_neg)); + } + } + } + } + + /// `ARC_ROPE_COHORT` is read by VALUE through `env_flag_is_set`, whose + /// polarity table (`=0`/unset ⇒ OFF, only yes-spellings ⇒ ON) is pinned in + /// `mistralrs_quant::env_flag`. What this pins HERE is the dispatch + /// default: with the flag unset — the state of every CI environment this + /// runs in — ragged rows must take the per-sequence loop, so the gathered + /// arm cannot go live by accident of a latch. + #[test] + fn ragged_dispatch_defaults_to_the_loop() { + assert!( + std::env::var_os("ARC_ROPE_COHORT").is_none(), + "test environment unexpectedly sets ARC_ROPE_COHORT; this test \ + pins the UNSET default and cannot do so with the flag present" + ); + assert!(!rope_cohort_enabled(), "unset must mean OFF"); + + let re = rope(); + let q = xs(3, N_HEADS, 1, ROPE_DIM); + let k = xs(3, 1, 1, ROPE_DIM); + let before = rope_cohort_stats::local_counts(); + re.forward(&q, &k, &[0usize, 4, 9]).unwrap(); + let after = rope_cohort_stats::local_counts(); + assert_eq!( + (after.0 - before.0, after.1 - before.1), + (0, 1), + "with ARC_ROPE_COHORT unset a ragged cohort must take the loop" + ); + } + + /// An offset past the table end must ERROR on the cohort arm exactly as it + /// does on the loop arm — and it must be THIS module's host-side guard + /// that fires, not the backend. The loop's `narrow` fails loud everywhere, + /// but `index_select` with an out-of-range index is not reliably + /// bounds-checked on CUDA and can read garbage silently; the arm we A/B on + /// hardware must not be the one with the quiet failure mode (D18). The + /// message substring is pinned so a backend error passing this test by + /// coincidence is impossible. + #[test] + fn out_of_bounds_offset_fails_loud_on_both_arms() { + let re = rope(); + let t = 1usize; + // MAX_POS = 32 rows: offset 31 is the last valid row at T=1; 32 is + // past the end. The valid rows in the same cohort prove the guard + // reports the MAX, not merely offsets[0]. + let offsets = vec![5usize, 32]; + let q = xs(2, N_HEADS, t, ROPE_DIM); + let k = xs(2, 1, t, ROPE_DIM); + + // Loop arm: narrow(0, 32, 1) on a 32-row table hard-errors. + assert!( + re.forward_ragged_loop(&q, &k, &offsets, t).is_err(), + "the loop arm accepted an offset past the table end" + ); + + // Cohort arm: must also error, and with the guard's own message. + let err = re + .forward_ragged_cohort(&q, &k, &offsets, t) + .expect_err("the cohort arm accepted an offset past the table end"); + let msg = err.to_string(); + assert!( + msg.contains("rope cohort gather out of bounds"), + "the cohort arm rejected the offset, but not via the host-side \ + guard — got: {msg}" + ); + assert!( + msg.contains("offset 32") && msg.contains("32-row"), + "the guard's message must name the offending offset and the table \ + size — got: {msg}" + ); + + // Same guard on the inverse-tail arm (shared `cohort_tables`). + let x = xs(2, N_HEADS, t, HEAD_DIM); + let err = re + .inverse_tail_ragged_cohort(&x, ROPE_DIM, HEAD_DIM - ROPE_DIM, &offsets, t) + .expect_err("the inverse-tail cohort arm accepted an offset past the table end"); + assert!(err.to_string().contains("rope cohort gather out of bounds")); + + // The boundary itself is not off by one: the last valid row works. + let ok = vec![5usize, 31]; + re.forward_ragged_cohort(&q, &k, &ok, t) + .expect("offset 31 + seq_len 1 on a 32-row table is in bounds"); + } +} diff --git a/mistralrs-core/src/layers_masker.rs b/mistralrs-core/src/layers_masker.rs index 4f31d12b7..8c83a5148 100644 --- a/mistralrs-core/src/layers_masker.rs +++ b/mistralrs-core/src/layers_masker.rs @@ -100,6 +100,63 @@ impl PastKvLenCache for Vec> { } } +/// The two invariants both left-padded-mask builders require, factored so the +/// device build and its host oracle cannot drift on what they refuse. +fn validate_ragged_live(b_sz: usize, past_kv_len: usize, live: &[usize]) -> Result<()> { + if live.len() != b_sz { + candle_core::bail!( + "ragged mask: {} live lengths for a batch of {b_sz}", + live.len() + ); + } + for &l in live { + if l > past_kv_len { + candle_core::bail!( + "ragged mask: live length {l} exceeds the padded width {past_kv_len}" + ); + } + } + Ok(()) +} + +thread_local! { + /// Per-device cache of an `arange(0, N)` F32 tensor, grown geometrically + /// and narrowed per call. + /// + /// This is what keeps [`CausalMasker::make_left_padded_causal_mask`] off + /// CLAUDE.md pitfall #5: a fresh `Tensor::arange` per decode step is a + /// host-side build plus an H2D sync in the hot loop, and the index rows it + /// needs only ever grow by one column per step. Thread-local because the + /// engine drives every forward from one thread, and a `cargo test` binary + /// runs each test on its own thread — a process-global here would be a + /// lock on the hot path for no aliasing that actually occurs. + static ARANGE_CACHE: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; +} + +/// `[0.0, 1.0, ..., n-1.0]` on `device`, served from [`ARANGE_CACHE`]. +/// +/// Allocates only when `n` first exceeds the cached capacity for that device, +/// growing to the next power of two (>= 1024) so a cohort whose context climbs +/// token by token re-allocates O(log n) times over its whole life. +fn cached_arange_f32(n: usize, device: &Device) -> Result { + ARANGE_CACHE.with(|c| { + let mut c = c.borrow_mut(); + let loc = device.location(); + if let Some((_, t)) = c.iter().find(|(l, _)| *l == loc) { + if t.dim(0)? >= n { + return t.narrow(0, 0, n); + } + } + let cap = n.next_power_of_two().max(1024); + let grown = Tensor::arange(0f32, cap as f32, device)?; + let out = grown.narrow(0, 0, n)?; + c.retain(|(l, _)| *l != loc); + c.push((loc, grown)); + Ok(out) + }) +} + impl CausalMasker { fn make_mask(&self, tgt_len: usize, past_kv_len: usize, device: &Device) -> Result { let offset = tgt_len + past_kv_len; @@ -218,6 +275,24 @@ impl CausalMasker { /// `[past - live[b], past + i]`. Everything outside that is killed. Note /// the row is never fully masked — column `past + i` is the query's own /// position — so softmax cannot produce a NaN row. + /// + /// # Built on the DEVICE, not the host + /// + /// This runs on **every decode step** while a cohort is ragged (the + /// `ragged_lead_pad` channel stays set until batch membership changes), and + /// the first implementation was a scalar triple loop over + /// `B * t_q * (past + t_q)` host `f32`s followed by a `Tensor::from_vec` + /// H2D copy — 4.2 MB/token at B=256, ctx-4096, i.e. CLAUDE.md pitfall #5 + /// landing on the hot path the moment ragged decode opens. Measured as the + /// thing that would make enabling ragged read as a regression. + /// + /// Now the only per-step H2D is the `[B]` lead vector; the row/column index + /// tensors come from a per-device cache ([`cached_arange_f32`]) and the + /// mask itself is a handful of broadcast comparisons on the device. + /// + /// [`Self::make_left_padded_causal_mask_host`] keeps the loop verbatim as + /// the test oracle; `the_device_mask_matches_the_host_oracle_element_for_element` + /// pins the two together. pub fn make_left_padded_causal_mask( &self, b_sz: usize, @@ -228,29 +303,70 @@ impl CausalMasker { dtype: DType, device: &Device, ) -> Result { - if live.len() != b_sz { - candle_core::bail!( - "ragged mask: {} live lengths for a batch of {b_sz}", - live.len() - ); - } + validate_ragged_live(b_sz, past_kv_len, live)?; + let k_len = past_kv_len + tgt_len; + + // [1, 1, 1, k]: key column index j. [1, 1, t, 1]: query row index i. + // Cached per device — no `arange` allocation in the hot loop. + let js = cached_arange_f32(k_len, device)?.reshape((1, 1, 1, k_len))?; + let is_ = cached_arange_f32(tgt_len, device)?.reshape((1, 1, tgt_len, 1))?; + + // [B, 1, 1, 1]: each row's dead prefix. The one per-step H2D copy, of + // exactly B elements. (f32 is exact for integers up to 2^24, far past + // any context length this engine serves.) + let lead: Vec = live + .iter() + .map(|l| (past_kv_len - l) as f32) + .collect::>(); + let lead = Tensor::from_vec(lead, (b_sz, 1, 1, 1), device)?; + + // last[i] = past + i — the query's own column, the causal upper edge. + let last = is_.affine(1.0, past_kv_len as f64)?; + + // A sliding window is expressed in ABSOLUTE positions, and this buffer + // is shifted by `lead`, so the window's low edge has to be shifted with + // it: the window admits `j > last - w`, applied on top of the dead + // prefix, never instead of it. (The host oracle's `saturating_sub` is + // subsumed by the `max` with `lead >= 0`.) + let start = match sliding_window { + Some(w) => lead.broadcast_maximum(&last.affine(1.0, 1.0 - w as f64)?)?, + None => lead, + }; + + // ok[b, 0, i, j] = (j >= start[b, i]) && (j <= last[i]). + let ok = js + .broadcast_ge(&start)? + .broadcast_mul(&js.broadcast_le(&last)?)? + .broadcast_as((b_sz, 1, tgt_len, k_len))?; + let zero = Tensor::zeros((b_sz, 1, tgt_len, k_len), DType::F32, device)?; + let neg_inf = Tensor::full(f32::NEG_INFINITY, (b_sz, 1, tgt_len, k_len), device)?; + ok.where_cond(&zero, &neg_inf)?.to_dtype(dtype) + } + + /// The host-loop reference implementation of + /// [`Self::make_left_padded_causal_mask`], kept **verbatim** as the test + /// oracle: three scalar loops whose every branch is legible, against which + /// the device build is pinned element-for-element. + /// + /// Not for the hot path — this is the `O(B * t_q * k)` host build plus a + /// full-mask H2D copy per call that the device build exists to remove. + pub fn make_left_padded_causal_mask_host( + &self, + b_sz: usize, + tgt_len: usize, + past_kv_len: usize, + live: &[usize], + sliding_window: Option, + dtype: DType, + device: &Device, + ) -> Result { + validate_ragged_live(b_sz, past_kv_len, live)?; let k_len = past_kv_len + tgt_len; let mut data: Vec = Vec::with_capacity(b_sz * tgt_len * k_len); for &l in live { - if l > past_kv_len { - candle_core::bail!( - "ragged mask: live length {l} exceeds the padded width {past_kv_len}" - ); - } let lead = past_kv_len - l; for i in 0..tgt_len { let last = past_kv_len + i; - // A sliding window is expressed in ABSOLUTE positions, and this - // buffer is shifted by `lead`, so the window's low edge has to - // be shifted with it: key `j` sits at absolute position - // `j - lead`, the query at `l + i`, so the window admits - // `j > last - w`. Applied on top of the dead prefix, never - // instead of it. let start = match sliding_window { Some(w) => lead.max((last + 1).saturating_sub(w)), None => lead, @@ -790,6 +906,102 @@ mod ragged_mask_tests { crate::kv_cache::set_ragged_lead_pad(None); } + /// 🔑 The device-built mask IS the host triple loop, element for element. + /// + /// The host loop ([`CausalMasker::make_left_padded_causal_mask_host`]) is + /// kept verbatim as the oracle: every branch of it is legible scalar code. + /// The production build replaces it with broadcast comparisons and one + /// `[B]` H2D copy, and this test is the only thing standing between that + /// rewrite and a silently wrong mask — a wrong element here is not an + /// error anywhere downstream, it is softmax weight on a key the row must + /// not see. + /// + /// Cases cover: mixed lengths with lead 0 and live 0 in one batch, + /// multi-row queries, no window, a window tighter than the live run, a + /// window wider than everything, and the degenerate past=0. + /// + /// Mutation check (run 2026-08-21): drop the `lead` term from the device + /// build (use `0.0`) and this fails on case 0 with 4 mismatching elements + /// (the dead prefix of the `live=2` row reads `0.0` where the oracle says + /// `-inf`). + #[test] + fn the_device_mask_matches_the_host_oracle_element_for_element() { + let device = Device::Cpu; + let cases: &[(usize, usize, &[usize], Option)] = &[ + (1, 6, &[6, 2], None), + (2, 6, &[6, 2, 0], None), + (3, 4, &[4, 1], None), + (1, 5, &[5, 1], Some(2)), + (2, 6, &[6, 3], Some(4)), + (1, 5, &[3, 5], Some(32768)), + (1, 0, &[0], None), + (4, 0, &[0, 0], None), + ]; + for (case, &(tgt, past, live, window)) in cases.iter().enumerate() { + let host = CausalMasker + .make_left_padded_causal_mask_host( + live.len(), + tgt, + past, + live, + window, + DType::F32, + &device, + ) + .unwrap(); + let dev = CausalMasker + .make_left_padded_causal_mask( + live.len(), + tgt, + past, + live, + window, + DType::F32, + &device, + ) + .unwrap(); + assert_eq!(host.dims(), dev.dims(), "case {case}: shape"); + let h = host.flatten_all().unwrap().to_vec1::().unwrap(); + let d = dev.flatten_all().unwrap().to_vec1::().unwrap(); + let mismatches: Vec = h + .iter() + .zip(&d) + .enumerate() + .filter(|(_, (a, b))| a != b) + .map(|(i, _)| i) + .collect(); + assert!( + mismatches.is_empty(), + "case {case} (tgt={tgt}, past={past}, live={live:?}, window={window:?}): \ + {} of {} elements differ, first at flat index {} (host {}, device {})", + mismatches.len(), + h.len(), + mismatches[0], + h[mismatches[0]], + d[mismatches[0]], + ); + // Neither builder may emit a NaN — 0 * -inf is the classic way to. + assert!(d.iter().all(|x| !x.is_nan()), "case {case}: device NaN"); + } + } + + /// Both builders refuse the same malformed inputs — a live length past the + /// padded width, and a live vector of the wrong batch size. The validator + /// is shared, but shared code can be un-shared by a refactor; this pins the + /// contract itself. + #[test] + fn the_device_mask_and_the_oracle_refuse_the_same_inputs() { + let device = Device::Cpu; + for (b, live) in [(2usize, &[7usize, 2][..]), (3, &[1, 1][..])] { + assert!(CausalMasker + .make_left_padded_causal_mask(b, 1, 6, live, None, DType::F32, &device) + .is_err()); + assert!(CausalMasker + .make_left_padded_causal_mask_host(b, 1, 6, live, None, DType::F32, &device) + .is_err()); + } + } + /// `make_causal_mask_matrix` routes to the ragged builder BEFORE both of /// its early returns — the `tgt_len == 1` shortcut and the flash-attn /// placeholder. A decode step is `tgt_len == 1` and still has a dead prefix diff --git a/mistralrs-core/src/lib.rs b/mistralrs-core/src/lib.rs index 26b8b14d9..2603392c2 100644 --- a/mistralrs-core/src/lib.rs +++ b/mistralrs-core/src/lib.rs @@ -174,7 +174,9 @@ pub use response::*; pub use sampler::{ CustomLogitsProcessor, DrySamplingParams, SamplingParams, StopTokens, TopLogprob, }; -pub use scheduler::{DefaultSchedulerMethod, SchedulerConfig}; +pub use scheduler::{ + sched_bucket_marker, sched_buckets, DefaultSchedulerMethod, SchedBuckets, SchedulerConfig, +}; pub use search::{SearchCallback, SearchFunctionParameters, SearchResult}; use serde::Serialize; pub use speech_models::{utils as speech_utils, SpeechGenerationConfig, SpeechLoaderType}; @@ -556,6 +558,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 +572,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 + ); + } + } } } diff --git a/mistralrs-core/src/models/deepseek4.rs b/mistralrs-core/src/models/deepseek4.rs index 0aa8909c3..64361e466 100644 --- a/mistralrs-core/src/models/deepseek4.rs +++ b/mistralrs-core/src/models/deepseek4.rs @@ -1287,9 +1287,21 @@ impl Attention { if cos.dtype() != candle_core::DType::BF16 || sin.dtype() != candle_core::DType::BF16 { decline!("RoPE tables are not BF16"); } - if seqlen_offsets.len() != 1 { + // The kernel needs ONE `pos_offset`, not a batch of one row. It reads + // the table at `pos_offset + t` where `t = blockIdx.y % seq_len` + // (`qk_norm_rope.cu`), a position that does not depend on `b`, and its + // grid is `(n_heads + 1, batch * seq_len)` with buffers sized + // `batch * ...` — so it is already correct at any batch size whose rows + // share an offset. This used to read `seqlen_offsets.len() != 1`, which + // tested the length of the vector rather than the distinctness of its + // values and therefore declined at EVERY batch size above one, leaving + // the fused path dead in exactly the serving regime it was written for. + // Rows that genuinely differ still decline, and still take the eager + // chain. Prove the swap with `ARC_QK_VERIFY=1`, which bit-compares + // against the eager output at every layer. + let Some(uniform_offset) = crate::layers::uniform_seqlen_offset(seqlen_offsets) else { decline!("per-sequence position offsets"); - } + }; if self.num_kv_heads != 1 { decline!("more than one KV head"); } @@ -1318,11 +1330,21 @@ impl Attention { if qh != self.num_attention_heads * head_dim { decline!("q_proj output width is not n_heads * head_dim"); } + // The batched cohort is OPT-IN. Fixing the length-vs-values gate above + // makes this kernel live on a path it has never run on, and this repo + // has paid for exactly that twice this week (TCFRAG, and the fused-512 + // attention path that dropped the mask for four days). `batch == 1` is + // untouched — already live, already exercised. See + // `qk::fused_cohort_enabled` for the doctrine and for how to turn the + // cohort path on once it has been measured on hardware. + if qb > 1 && !qk::fused_cohort_enabled() { + decline!("batched cohort disabled by ARC_QK_FUSED_COHORT=0 (default is ON since s10)"); + } match kv_normed.dims3() { Ok((kb, kt, kh)) if kb == qb && kt == qt && kh == head_dim => {} _ => decline!("kv_norm output shape does not match [B, T, head_dim]"), } - let pos_offset = seqlen_offsets[0]; + let pos_offset = uniform_offset; match cos.dims2() { Ok((rows, cols)) if cols == rope_dim / 2 && rows >= pos_offset + qt => {} _ => decline!("RoPE table does not cover [pos_offset, pos_offset + T)"), @@ -1969,8 +1991,8 @@ impl Attention { graph_attn } None => { - // FP8 code storage is opt-in (`ARC_V4_FP8_KV=1`); unset stores - // the dense BF16 K the model has always stored. + // FP8 code storage is default-ON since s10 (`ARC_V4_FP8_KV=0` + // restores the dense BF16 K layout). let cached = { let _s = arc_profiler::device_span("kv_cache_append"); append_kv_mqa( @@ -3137,16 +3159,21 @@ fn append_graph_kv_mqa( Ok(k_full) } -/// FP8 storage for the cached K. **Opt-in: off unless `ARC_V4_FP8_KV=1`.** +/// FP8 storage for the cached K. **Default ON since the s10 GPU A/B; +/// `ARC_V4_FP8_KV=0` restores the BF16 layout.** /// /// wave43-BU shipped this defaulted ON without ever running it on a GPU — its /// own notes said *"needs a GPU: first CUDA exercise of a U8 `SingleCache`"* — /// and the first V4 forward on a commit that contained it died, every request, /// with `dtype mismatch in slice-set, lhs: BF16, rhs: U8` (wave48-BY). The /// layout is sound and bit-exact (see [`super::dsv4_kv_fp8`]); what was never -/// exercised is the CUDA leg of it. Until a GPU A/B has shown -/// `ARC_V4_FP8_KV=1` and unset producing token-identical greedy output, the -/// default is the layout that has actually served (wave49-BZ). +/// exercised was the CUDA leg. That A/B has now run: s10 leg 10 (box +/// `arc-s10-ledger`, H200, candidate `4d03b9e25`) — seeded canaries +/// byte-identical to the BF16 arm, b=1 41.49 vs 40.76 tok/s (not slower), and +/// engagement PROVEN by nsys kernel capture (`arc_kv_fp8_quantize_kernel` +/// 16,297 launches, `arc_kv_fp8_dequantize_kernel` 32,594 in one b=1 run — +/// the U8 CUDA leg genuinely executed). This is the difference from wave43: +/// the same default, now standing on a measurement instead of a hope. /// /// Also forced off under `ARC_V4_CAPTURE_PROBE`: the CUDA-graph decode arm /// writes through `mistralrs_quant::kvwrite::write_kv_inplace`, which is @@ -3206,7 +3233,10 @@ fn v4_fp8_kv_enabled() -> bool { /// about the default, which is exactly the kind of test DOCTRINE D12 counts as /// worse than none. fn fp8_kv_enabled_from(var: Option<&str>, capture_probe: bool) -> bool { - var == Some("1") && !capture_probe + // Default ON since s10 (measured; see the doc above). The capture-probe + // veto is unconditional: write_kv_inplace has no U8 arm, so the graph + // decode path and the U8 cache remain mutually exclusive. + var != Some("0") && !capture_probe } /// RUN-161 decode profiler. Enabled by `ARC_TIME_DECODE=1`. Accumulates @@ -5902,18 +5932,28 @@ mod kv_footprint_tests { /// purpose: the latter is a `OnceLock`, so a test that set the variable /// would resolve it once and prove nothing about the default. #[test] - fn v4_fp8_kv_is_opt_in() { - assert!(!fp8_kv_enabled_from(None, false), "unset must be OFF"); - assert!(!fp8_kv_enabled_from(Some("0"), false)); - assert!(!fp8_kv_enabled_from(Some(""), false)); - assert!(!fp8_kv_enabled_from(Some("true"), false)); - assert!(!fp8_kv_enabled_from(Some("on"), false)); - assert!(!fp8_kv_enabled_from(Some("2"), false)); + fn v4_fp8_kv_is_default_on_with_a_literal_zero_kill_switch() { + assert!( + fp8_kv_enabled_from(None, false), + "unset must be ON — measured s10 leg 10 (arc-s10-ledger, 4d03b9e25): \ + canaries byte-identical, engagement proven by nsys kernel capture" + ); + assert!(!fp8_kv_enabled_from(Some("0"), false), "\"0\" is the kill switch"); + assert!(fp8_kv_enabled_from(Some(""), false)); + assert!(fp8_kv_enabled_from(Some("true"), false)); + assert!(fp8_kv_enabled_from(Some("on"), false)); + assert!(fp8_kv_enabled_from(Some("2"), false)); assert!(fp8_kv_enabled_from(Some("1"), false), "=1 must be ON"); assert!( !fp8_kv_enabled_from(Some("1"), true), "ARC_V4_CAPTURE_PROBE must veto it: write_kv_inplace has no U8 arm" ); + assert!( + !fp8_kv_enabled_from(None, true), + "the veto must hold for the DEFAULT too, not just for explicit =1 — \ + the capture probe cannot write a U8 cache regardless of how FP8 KV \ + was selected" + ); } /// The per-row ragged mask is **opt-in**, and off it hands @@ -5957,6 +5997,9 @@ mod kv_footprint_tests { #[cfg(test)] mod tests { use super::*; + // Only the tests drive the pin override; importing it at module scope made + // it an unused import on every non-test build. + use crate::kv_cache::xs_rolling; /// The **verbatim** `config.json` published at /// @@ -8149,6 +8192,156 @@ mod tests { ragged_batch_matches_b1(128, &[129, 143, 1024, 1150], 3) } + /// 🔑 Settles the one discrepancy between two independent models of this + /// buffer, and does it without a GPU. + /// + /// The ArcGraph chain measured the reallocation cycling through + /// `4096 x {18, 19, 20, 21}`. Running the retention rule forward predicts + /// `{20, 21, 22, 23}`. Both are `ratio`-consecutive, both contain 21, and + /// they are offset by exactly 2 — so one of the two models is wrong about + /// *why*, even though the bound holds either way. + /// + /// The pre-committed reconciliation was: theirs is the pre-saturation RAMP + /// (their run generated ~21 tokens, so `base` had not finished advancing + /// and `W = tokens - base` was still climbing), mine is the steady state. + /// This checks that claim directly by recording the width from token 1, + /// and it is the difference between a bound that is right for the right + /// reason and one that happens to be right — which is the trap the next + /// context length springs. + #[test] + fn the_window_ramps_then_settles_to_ratio_consecutive_sizes() -> Result<()> { + xs_rolling::pin_test_override::with(false, || { + let device = Device::Cpu; + let (hidden, head_dim, ratio) = (32usize, 16usize, 4usize); + let compressor = rolling_test_compressor(ratio, hidden, head_dim, &device)?; + let stream = per_row_streams(1, 130, hidden, &device)?.remove(0); + let mut s = XsRollingCache::new(ratio, compressor.coff, head_dim, 4096); + + let mut widths = Vec::new(); + for i in 0..128 { + s.advance(&stream.narrow(1, i, 1)?, |w| compressor.forward_from_xs(w))?; + widths.push(s.tail_width()?); + } + let cap = ratio * compressor.coff + crate::kv_cache::XS_TAIL_MARGIN_TOKENS; + + // The ramp: early widths are BELOW the steady-state band, because + // `base` is still pinned at 0 while `tokens` climbs. + let ramp: std::collections::BTreeSet<_> = widths[..20].iter().copied().collect(); + let steady: std::collections::BTreeSet<_> = widths[64..].iter().copied().collect(); + assert!( + ramp.iter().min() < steady.iter().min(), + "the early widths must be a ramp, not the steady band: ramp {ramp:?} \ + steady {steady:?}" + ); + // The steady state is exactly `ratio` consecutive sizes… + assert_eq!( + steady.len(), + ratio, + "steady-state width must cycle through exactly `ratio` sizes, got {steady:?}" + ); + let lo = *steady.iter().next().unwrap(); + assert_eq!( + steady.iter().copied().collect::>(), + (lo..lo + ratio).collect::>(), + "the steady band must be consecutive, got {steady:?}" + ); + // …and it is the top of the range the bound allows: [cap-ratio, cap). + assert_eq!( + (lo, lo + ratio - 1), + (cap - ratio, cap - 1), + "steady band should be [cap-ratio, cap-1] = [{}, {}]", + cap - ratio, + cap - 1 + ); + // The bound, over every step including the ramp. + assert!( + widths.iter().all(|&w| w < cap), + "some width reached the pinned capacity {cap}, so pinning there could truncate: \ + max was {:?}", + widths.iter().max() + ); + Ok(()) + }) + } + + /// 🔑 Pinning the retained window must change the ALLOCATION and nothing + /// else. This is the evidence that the pin is SAFE — the two settings are + /// run against the same stream for 40 steps and required to agree EXACTLY, + /// on the compressed rows and on both time bases. + /// + /// It was once also offered as the evidence for defaulting + /// `ARC_V4_XS_PIN_WINDOW` on without a throughput number. It is not that: + /// it proves the pin changes no answer, which is necessary and not + /// sufficient. The flag is opt-in until the throughput number exists on + /// this tree — see `xs_rolling::xs_pin_window_enabled_from`'s FLIP + /// CONDITION for the experiment that settles it. + /// + /// The argument for why they must — every offset is derived from the row's + /// own token count, so a wider buffer shifts `off` by exactly the widening + /// and the compressor sees the same absolute tokens — is the kind of + /// argument that has been wrong on this chain before. So it is checked. + #[test] + fn pinning_the_window_is_numerically_inert() -> Result<()> { + let device = Device::Cpu; + let (hidden, head_dim, ratio) = (32usize, 16usize, 4usize); + let compressor = rolling_test_compressor(ratio, hidden, head_dim, &device)?; + let stream = per_row_streams(1, 220, hidden, &device)?.remove(0); + + let run = |pinned: bool| -> Result<(Tensor, Vec, Vec, Vec)> { + xs_rolling::pin_test_override::with(pinned, || { + let mut s = XsRollingCache::new(ratio, compressor.coff, head_dim, 4096); + s.advance(&stream.narrow(1, 0, 37)?, |w| compressor.forward_from_xs(w))?; + let mut widths = Vec::new(); + for i in 0..40 { + s.advance(&stream.narrow(1, 37 + i, 1)?, |w| { + compressor.forward_from_xs(w) + })?; + widths.push(s.tail_width()?); + } + let comp = s.compressed_rows()?.expect("past one full group"); + let (tok, base) = s.row_lens(); + Ok((comp, tok.to_vec(), base.to_vec(), widths)) + }) + }; + + let (comp_off, tok_off, base_off, widths_off) = run(false)?; + let (comp_on, tok_on, base_on, widths_on) = run(true)?; + + assert_eq!( + max_abs_diff(&comp_off, &comp_on)?, + 0.0, + "pinning the window changed the compressed rows — it must only change the allocation" + ); + assert_eq!(tok_off, tok_on, "token counts diverged"); + assert_eq!(base_off, base_on, "resume points diverged"); + + // …and the fixture must actually exercise the difference, or the + // equality above proves nothing (five fixtures on this chain have + // survived a mutation by not distinguishing the two answers). + let distinct_off: std::collections::BTreeSet<_> = widths_off.iter().collect(); + assert!( + distinct_off.len() > 1, + "unpinned widths did not vary ({widths_off:?}), so this proves nothing about pinning" + ); + assert_eq!( + distinct_off.len(), + ratio, + "unpinned width should cycle through `ratio` consecutive sizes, got {distinct_off:?}" + ); + let cap = ratio * compressor.coff + crate::kv_cache::XS_TAIL_MARGIN_TOKENS; + assert!( + widths_on.iter().all(|&w| w == cap), + "pinned widths must be the constant {cap}, got {widths_on:?}" + ); + // The bound the pin rests on: the unpinned width never reaches it. + assert!( + widths_off.iter().all(|&w| w < cap), + "the pinned capacity {cap} must exceed every width the retention rule produces, \ + got {widths_off:?}" + ); + Ok(()) + } + /// The control the ragged tests need: a UNIFORM batch takes the untouched /// scalar path and is also token-identical. If this ever failed, the /// "flag off is byte-identical" claim would be false and the ragged @@ -8297,8 +8490,73 @@ mod tests { /// the per-sequence invariant `window width == tokens - base`, taking that /// row's share from the END of the shared window. Taking it from the front /// would hand the row somebody else's older tokens under its own `base`. + /// ⚠️ Run with the window pin OFF, deliberately. Its premise is a row + /// *narrower* than the shared window, which is what a resizing buffer + /// produces; pinned, every row is the same width and the premise — and the + /// re-anchoring it checks — cannot be constructed. The re-anchoring still + /// has to be right when the pin is off, and + /// `splitting_a_pinned_row_keeps_the_buffer_and_the_resume_point` covers + /// the pinned side, where the requirement is the opposite one. #[test] fn splitting_a_batched_row_restores_the_per_sequence_window() -> Result<()> { + xs_rolling::pin_test_override::with(false, splitting_a_batched_row_inner) + } + + /// 🔑 The pinned counterpart, and the reason it matters is cost, not + /// correctness: `clone_out_cache` calls `split_row` once per layer per + /// sequence on EVERY engine step. A split that re-narrowed to + /// `tokens - base` would reallocate the buffer every step and undo the pin + /// precisely on the hot path it exists for. So the split must hand back the + /// SAME tensor, and must not move the row's resume point while doing it. + #[test] + fn splitting_a_pinned_row_keeps_the_buffer_and_the_resume_point() -> Result<()> { + xs_rolling::pin_test_override::with(true, || { + let device = Device::Cpu; + let (hidden, head_dim, ratio) = (32usize, 16usize, 4usize); + let compressor = rolling_test_compressor(ratio, hidden, head_dim, &device)?; + let lens = [37usize, 40]; + let streams = per_row_streams(2, 64, hidden, &device)?; + let mut refs: Vec = Vec::new(); + for (i, &l) in lens.iter().enumerate() { + let mut s = XsRollingCache::new(ratio, compressor.coff, head_dim, 4096); + s.advance(&streams[i].narrow(1, 0, l)?, |w| { + compressor.forward_from_xs(w) + })?; + refs.push(s); + } + let cap = ratio * compressor.coff + crate::kv_cache::XS_TAIL_MARGIN_TOKENS; + for r in &refs { + assert_eq!( + r.tail.as_ref().unwrap().dim(1)?, + cap, + "a pinned per-sequence window must already be the capacity" + ); + } + let batched = batch_xs(&refs)?; + let shared = batched.tail.as_ref().unwrap().dim(1)?; + assert_eq!(shared, cap, "pinned rows batch without any padding at all"); + + let comps = batched.comp.all_data.as_ref().unwrap().chunk(2, 0)?; + let tails = batched.tail.as_ref().unwrap().chunk(2, 0)?; + for i in 0..2 { + let out = batched.split_row(i, comps[i].clone(), tails[i].clone())?; + assert_eq!( + out.tail.as_ref().unwrap().dim(1)?, + cap, + "row {i}: the split re-narrowed a pinned buffer, which would reallocate it \ + on every engine step" + ); + assert_eq!( + (out.row_lens().0[0], out.row_lens().1[0]), + (refs[i].row_lens().0[0], refs[i].row_lens().1[0]), + "row {i}: keeping the buffer must not move the resume point" + ); + } + Ok(()) + }) + } + + fn splitting_a_batched_row_inner() -> Result<()> { let device = Device::Cpu; let (hidden, head_dim, ratio) = (32usize, 16usize, 4usize); let compressor = rolling_test_compressor(ratio, hidden, head_dim, &device)?; diff --git a/mistralrs-core/src/models/dsv4_attention.rs b/mistralrs-core/src/models/dsv4_attention.rs index eeee5a75c..f916bb1da 100644 --- a/mistralrs-core/src/models/dsv4_attention.rs +++ b/mistralrs-core/src/models/dsv4_attention.rs @@ -576,6 +576,82 @@ pub(crate) fn graph_compressed_mask( Tensor::cat(&[&raw, &comp], 3) } +/// Integer key naming everything the uniform-batch union mask depends on. +/// +/// The additive `[1, 1, t_q, t_k + t_c]` mask [`dsv4_attention`] builds on the +/// uniform-batch, non-graph path is a pure function of these fields — the +/// `arange` calls, casts, compares and the 0/-inf conversion take nothing else as +/// input — so two calls with equal keys produce bit-identical tensors, and a +/// cached one is an identity, not an approximation. +#[derive(Clone, PartialEq, Eq, Debug)] +struct UnionMaskKey { + q0: usize, + t_q: usize, + raw_base: usize, + t_k: usize, + t_c: usize, + ratio: usize, + window: usize, + dtype: DType, + dev: candle_core::DeviceLocation, +} + +/// A handful of slots, FIFO. One step touches at most three distinct keys — +/// the Standard / CSA / HCA layer classes agree on everything but `t_c` and +/// `ratio` — so 4 covers a step with a spare, and earlier steps' entries +/// (whose `q0` has advanced) age out on their own. +const UNION_MASK_MEMO_SLOTS: usize = 4; +thread_local! { + static UNION_MASK_MEMO: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; +} + +/// Memoized union-mask lookup: the cached additive mask for `key`, or `build` +/// run once and cached. +/// +/// WHY: the union mask was rebuilt from scratch once per LAYER per step. On +/// the prefill path that is three `Tensor::arange` host-to-device uploads plus +/// their casts, compares and the 0/-inf conversion (CLAUDE.md pitfall #5 — +/// each upload is a CPU-to-GPU sync), all recomputing a pure function of +/// `key`. Every layer of one compress-ratio class shares a key within a step, +/// so this builds each distinct mask once per step instead of once per layer. +/// +/// Thread-local, like `GRAPH_BLOCK_IDS` above: the engine drives all layers of +/// a step from one thread, which is the reuse that matters. +fn union_mask_memo(key: &UnionMaskKey, build: impl FnOnce() -> Result) -> Result { + UNION_MASK_MEMO.with(|c| -> Result { + { + let cache = c.borrow(); + if let Some((_, t)) = cache.iter().find(|(k, _)| k == key) { + return Ok(t.clone()); + } + } + let t = build()?; + let mut cache = c.borrow_mut(); + if cache.len() >= UNION_MASK_MEMO_SLOTS { + cache.remove(0); + } + cache.push((key.clone(), t.clone())); + Ok(t) + }) +} + +/// Boolean `[t_q, n_keys]` validity → additive `[1, 1, t_q, n_keys]` mask +/// (0 on valid, -inf on invalid) in the compute dtype. +fn additive_union_mask( + valid: &Tensor, + t_q: usize, + n_keys: usize, + dtype: DType, + dev: &Device, +) -> Result { + let zeros = Tensor::zeros((t_q, n_keys), dtype, dev)?; + let neg_inf = neg_inf_full((t_q, n_keys), dtype, dev)?; + valid + .where_cond(&zeros, &neg_inf)? + .reshape((1, 1, t_q, n_keys)) +} + /// Per-call configuration for V4 hybrid attention. #[derive(Debug, Clone, Copy)] pub struct Dsv4AttentionConfig<'a> { @@ -879,83 +955,100 @@ pub fn dsv4_attention( // `arange` copies from a HOST pointer, and host-pointer copies are the // class that makes a decode step expensive to capture as a CUDA graph. let decode_row = t_q == 1; - // `qp` outlives this block only to drive the compressed branch's threshold, - // and only on the paths that still evaluate that threshold in F32. - let qp = if decode_row { - None - } else { - Some( - Tensor::arange(q0 as u32, (q0 + t_q) as u32, dev)? - .to_dtype(DType::F32)? - .reshape((t_q, 1))?, - ) + let ratio = cfg.compress_ratio.ratio(); + let t_c = match compressed_kv { + Some(comp) => comp.dim(2)?, + None => 0, }; - let raw_valid = if decode_row { - Tensor::ones((t_q, t_k), DType::U8, dev)? - } else { - let kp = Tensor::arange(raw_base as u32, (raw_base + t_k) as u32, dev)? - .to_dtype(DType::F32)? - .reshape((1, t_k))?; - let qp = qp - .as_ref() - .expect("qp is built on exactly the `t_q > 1` branch this arm is"); - let causal = kp.broadcast_le(qp)?; - let lower = (qp - window as f64)?; - let in_window = kp.broadcast_gt(&lower)?; - (causal * in_window)? // [t_q, t_k] (u8) + let graph_mode = cfg.graph_positions.is_some(); + + // ---- Union validity [t_q, t_k + t_c] (U8), built LAZILY --------------- + // On the uniform-batch, non-graph path the additive mask derived from this + // is memoized across layers on its integer key (`union_mask_memo`), and a + // hit skips this entire build — the prefill path's `arange` uploads, their + // casts and every compare included. A miss runs it character-for-character + // as before. + let build_valid = || -> Result { + // `qp` drives the raw window and, on the paths that still evaluate the + // compressed branch's threshold in F32, that threshold too. + let qp = if decode_row { + None + } else { + Some( + Tensor::arange(q0 as u32, (q0 + t_q) as u32, dev)? + .to_dtype(DType::F32)? + .reshape((t_q, 1))?, + ) + }; + let raw_valid = if decode_row { + Tensor::ones((t_q, t_k), DType::U8, dev)? + } else { + let kp = Tensor::arange(raw_base as u32, (raw_base + t_k) as u32, dev)? + .to_dtype(DType::F32)? + .reshape((1, t_k))?; + let qp = qp + .as_ref() + .expect("qp is built on exactly the `t_q > 1` branch this arm is"); + let causal = kp.broadcast_le(qp)?; + let lower = (qp - window as f64)?; + let in_window = kp.broadcast_gt(&lower)?; + (causal * in_window)? // [t_q, t_k] (u8) + }; + if compressed_kv.is_none() { + return Ok(raw_valid); + } + // Compressed branch mask: query r attends compressed block b iff + // b < (q0+r+1)/ratio (causal over fully-completed ratio-blocks). + let comp_valid = if graph_mode { + // Fixed-capacity graph decode: `qp` is derived from the buffer + // *capacity*, so this threshold would freeze at a constant (see + // `Dsv4AttentionConfig::graph_positions`). Leave the local half + // permissive and let the device-derived mask folded in below + // carry the real pattern — a `0` here is the additive identity, + // so the two compose without double-counting. + Tensor::ones((t_q, t_c), DType::U8, dev)? + } else { + match &qp { + // Decode (`t_q == 1`). `q0` is a host value and + // `CompressRatio::ratio` is 1 / 4 / 128 — always a power of + // two — so `floor((q0 + 1) / ratio)` is exact integer + // arithmetic on the host and `bp` can stay in the U32 the + // `arange` already produced. Bit-identical to the F32 chain + // below: `q0 + 1`, the quotient and every block index are + // integers, a power-of-two divide is exact in binary + // floating point, and F32 represents every integer below + // 2^24 = 16 777 216 — against V4's 163 840-token position + // ceiling. Same predicate, same operands, four fewer + // launches (a cast, an add, a divide and a floor). + None => { + let threshold = + Tensor::new(&[((q0 + 1) / ratio) as u32], dev)?.reshape((1, 1))?; + // The block indices come from the cached `(ratio, + // device)` table at ratio 1, not `Tensor::arange`. + // `compress_positions(t_c, 1, dev)` IS `[0, 1, …, t_c-1]` + // in U32 — bit-identical to the `arange` it replaces — + // but it is a zero-copy `narrow` view of a table built + // once during warmup, so the per-layer host->device copy + // goes away. See `layers::compress_positions`. + crate::layers::compress_positions(t_c, 1, dev)? + .reshape((1, t_c))? + .broadcast_lt(&threshold)? // [1, t_c] (u8) + } + Some(qp) => { + let bp = Tensor::arange(0u32, t_c as u32, dev)? + .to_dtype(DType::F32)? + .reshape((1, t_c))?; + let threshold = ((qp + 1.0)? / ratio as f64)?.floor()?; // [t_q, 1] + bp.broadcast_lt(&threshold)? // [t_q, t_c] (u8) + } + } + }; + Tensor::cat(&[&raw_valid, &comp_valid], 1) }; - // ---- Build the union key set + validity mask. ------------------------- - let (k_cat, v_cat, valid) = match compressed_kv { + // ---- Build the union key set ------------------------------------------ + let (k_cat, v_cat) = match compressed_kv { Some(comp) => { - let t_c = comp.dim(2)?; - // Compressed branch mask: query r attends compressed block b iff - // b < (q0+r+1)/ratio (causal over fully-completed ratio-blocks). - let ratio = cfg.compress_ratio.ratio(); - let comp_valid = match cfg.graph_positions { - // Fixed-capacity graph decode: `qp` is derived from the buffer - // *capacity*, so this threshold would freeze at a constant (see - // `Dsv4AttentionConfig::graph_positions`). Leave the local half - // permissive and let the device-derived mask folded in below - // carry the real pattern — a `0` here is the additive identity, - // so the two compose without double-counting. - Some(_) => Tensor::ones((t_q, t_c), DType::U8, dev)?, - None => match &qp { - // Decode (`t_q == 1`). `q0` is a host value and - // `CompressRatio::ratio` is 1 / 4 / 128 — always a power of - // two — so `floor((q0 + 1) / ratio)` is exact integer - // arithmetic on the host and `bp` can stay in the U32 the - // `arange` already produced. Bit-identical to the F32 chain - // below: `q0 + 1`, the quotient and every block index are - // integers, a power-of-two divide is exact in binary - // floating point, and F32 represents every integer below - // 2^24 = 16 777 216 — against V4's 163 840-token position - // ceiling. Same predicate, same operands, four fewer - // launches (a cast, an add, a divide and a floor). - None => { - let threshold = - Tensor::new(&[((q0 + 1) / ratio) as u32], dev)?.reshape((1, 1))?; - // The block indices come from the cached `(ratio, - // device)` table at ratio 1, not `Tensor::arange`. - // `compress_positions(t_c, 1, dev)` IS `[0, 1, …, t_c-1]` - // in U32 — bit-identical to the `arange` it replaces — - // but it is a zero-copy `narrow` view of a table built - // once during warmup, so the per-layer host->device copy - // goes away. See `layers::compress_positions`. - crate::layers::compress_positions(t_c, 1, dev)? - .reshape((1, t_c))? - .broadcast_lt(&threshold)? // [1, t_c] (u8) - } - Some(qp) => { - let bp = Tensor::arange(0u32, t_c as u32, dev)? - .to_dtype(DType::F32)? - .reshape((1, t_c))?; - let threshold = ((qp + 1.0)? / ratio as f64)?.floor()?; // [t_q, 1] - bp.broadcast_lt(&threshold)? // [t_q, t_c] (u8) - } - }, - }; - let valid = Tensor::cat(&[&raw_valid, &comp_valid], 1)?; let k_cat = Tensor::cat(&[k, comp], 2)?.contiguous()?; // The second half of the same duplication -- see `v_aliases_k` // above. `v_cat` was an unconditional @@ -969,33 +1062,51 @@ pub fn dsv4_attention( fused_v_stats::record_materialized(); Tensor::cat(&[v, comp], 2)?.contiguous()? }; - (k_cat, v_cat, valid) + (k_cat, v_cat) } - None => (k.clone(), v.clone(), raw_valid), + None => (k.clone(), v.clone()), }; record_launch_shape("k_cat", k_cat.dims()); record_launch_shape("v_cat", v_cat.dims()); // Boolean validity → additive mask (0 / -inf), broadcast over batch+heads. - let n_keys = valid.dim(1)?; + let n_keys = t_k + t_c; let mask = match row_q0 { - // Uniform batch: verbatim the pre-change code. - None => { - let zeros = Tensor::zeros((t_q, n_keys), q.dtype(), dev)?; - let neg_inf = neg_inf_full((t_q, n_keys), q.dtype(), dev)?; - valid - .where_cond(&zeros, &neg_inf)? - .reshape((1, 1, t_q, n_keys))? - } + // Uniform batch, no graph positions: a pure function of the integer + // key, memoized across layers. The build a miss runs — + // `additive_union_mask(&build_valid()?, ..)` — is verbatim the + // pre-memo code. + None if !graph_mode => union_mask_memo( + &UnionMaskKey { + q0, + t_q, + raw_base, + t_k, + t_c, + ratio, + window, + dtype: q.dtype(), + dev: dev.location(), + }, + || additive_union_mask(&build_valid()?, t_q, n_keys, q.dtype(), dev), + )?, + // Uniform batch under fixed-capacity graph decode: NOT memoized, on + // purpose. The constant capacity would make the key constant across + // steps, and these ops must be issued inside the capture region rather + // than hoisted into a warmup-owned tensor whose eviction would leave a + // recorded graph reading freed memory. The build is already + // arange-free here — both halves are fills. + None => additive_union_mask(&build_valid()?, t_q, n_keys, q.dtype(), dev)?, // Ragged cohort: `[b, 1, t_q, n_keys]`, because the compressed - // threshold and the dead prefix are both per-row. + // threshold and the dead prefix are both per-row. Not memoized: the + // per-row leads are part of the content but not of the integer key. Some(rows) => ragged_union_mask( rows, q0, - &valid, + &build_valid()?, raw_base, t_k, - cfg.compress_ratio.ratio(), + ratio, q.dtype(), dev, )?, @@ -3719,3 +3830,140 @@ mod fused_v_alias_tests { Ok(()) } } + +/// The union-mask memo: the mask is a pure function of an integer key and was +/// rebuilt once per layer per step (three `arange` uploads on the prefill +/// path). These tests pin the two failure modes a memo can add: doing nothing +/// (always missing) and serving a stale entry for a changed key. +#[cfg(test)] +#[allow(clippy::cast_precision_loss)] // usize -> f32 on test-sized fixtures, the file's test idiom +mod union_mask_memo_tests { + use super::*; + use candle_core::{DType, Device, Result, Tensor}; + use std::collections::HashMap; + + /// The memo itself: a hit returns the cached handle WITHOUT running the + /// builder, and a changed key runs it. Handle identity is the only + /// observable — a memo that always misses returns correct values while + /// doing nothing. + #[test] + fn memo_hits_without_building_and_misses_on_a_changed_key() -> Result<()> { + let dev = Device::Cpu; + let key_a = UnionMaskKey { + q0: 34, + t_q: 1, + raw_base: 27, + t_k: 8, + t_c: 12, + ratio: 4, + window: 8, + dtype: DType::F32, + dev: dev.location(), + }; + let key_b = UnionMaskKey { + q0: 35, + raw_base: 28, + ..key_a.clone() + }; + + let t1 = Tensor::zeros((1, 1, 1, 20), DType::F32, &dev)?; + let got = union_mask_memo(&key_a, || Ok(t1.clone()))?; + assert_eq!(got.id(), t1.id(), "miss must return the built tensor"); + + let hit = union_mask_memo(&key_a, || panic!("builder ran on a hit"))?; + assert_eq!(hit.id(), t1.id(), "hit must return the cached handle"); + + let t2 = Tensor::ones((1, 1, 1, 20), DType::F32, &dev)?; + let miss = union_mask_memo(&key_b, || Ok(t2.clone()))?; + assert_eq!( + miss.id(), + t2.id(), + "an advanced position was served the previous position's mask" + ); + // ...and key A is still cached alongside. + let hit_a = union_mask_memo(&key_a, || panic!("builder ran on a hit"))?; + assert_eq!(hit_a.id(), t1.id()); + Ok(()) + } + + fn mk(b: usize, h: usize, t: usize, d: usize, seed: f32, dev: &Device) -> Result { + let data: Vec = (0..(b * h * t * d)) + .map(|i| ((i as f32) * seed).sin()) + .collect(); + Tensor::from_vec(data, (b, h, t, d), dev) + } + + /// One deterministic decode step at absolute context `ctx`, returned as + /// output bits. `ctx` 35 vs 36 with ratio 4 and window 8 gives two calls + /// whose MASKS have the same shape but different content (the compressed + /// threshold `floor(ctx / 4)` crosses 8 → 9), which is exactly the + /// stale-serving hazard. + fn run(ctx: usize) -> Result> { + let dev = Device::Cpu; + let (h, d, t_c) = (2usize, 8usize, 12usize); + let q = mk(1, h, 1, d, 0.03, &dev)?; + let k = mk(1, 1, ctx, d, 0.11, &dev)?; + let comp = mk(1, 1, t_c, d, 0.05, &dev)?; + let sinks: Vec = (0..h).map(|i| 0.5 + 0.25 * i as f32).collect(); + let sdpa = SdpaParams { + n_kv_groups: h, + softcap: None, + softmax_scale: 1.0 / (d as f32).sqrt(), + sliding_window: None, + sinks: Some(Tensor::from_vec(sinks, (1, h, 1, 1), &dev)?), + }; + let flash = FlashParams { + max_q: 0, + max_k: 0, + cumulative_seqlens_q: HashMap::new(), + cumulative_seqlens_k: HashMap::new(), + causal: false, + }; + let cfg = Dsv4AttentionConfig { + compress_ratio: CompressRatio::Csa, + sliding_window: 8, + raw_prefix: 0, + row_q0: None, + graph_positions: None, + }; + let out = dsv4_attention(&q, &k, &k, Some(&comp), None, &flash, &sdpa, cfg)?; + Ok(out + .to_dtype(DType::F32)? + .flatten_all()? + .to_vec1::()? + .iter() + .map(|x| x.to_bits()) + .collect()) + } + + /// End to end, the memo must be invisible: sequential calls on ONE thread + /// (where the memo carries state between them) must be bit-identical to + /// the same calls on FRESH threads (where the thread-local memo is empty + /// and every mask is built from scratch). If the memo ever served ctx=35's + /// mask to ctx=36 — same shape, different admitted block set — the + /// sequential run would diverge from the fresh reference. + #[test] + fn memoized_masks_are_bit_identical_to_fresh_builds() -> Result<()> { + let fresh_35 = std::thread::spawn(|| run(35)).join().unwrap()?; + let fresh_36 = std::thread::spawn(|| run(36)).join().unwrap()?; + + // Same thread: 35 (miss), 36 (miss, memo holds 35), 35 again (hit). + let seq_35 = run(35)?; + let seq_36 = run(36)?; + let seq_35_again = run(35)?; + + assert_eq!( + seq_35, fresh_35, + "first build disagrees with a fresh thread" + ); + assert_eq!( + seq_36, fresh_36, + "ctx=36 after ctx=35 disagrees with a fresh thread — the memo served a stale mask" + ); + assert_eq!( + seq_35_again, fresh_35, + "a memo HIT changed the answer — the cached mask is not the built one" + ); + Ok(()) + } +} diff --git a/mistralrs-core/src/paged_attention/cache_engine.rs b/mistralrs-core/src/paged_attention/cache_engine.rs index 5cc88040a..b34adb1b2 100644 --- a/mistralrs-core/src/paged_attention/cache_engine.rs +++ b/mistralrs-core/src/paged_attention/cache_engine.rs @@ -8,10 +8,13 @@ use serde::{Deserialize, Serialize}; use super::config::{KvCacheLayout, ModelConfigLike}; -/// Head dimension the TurboQuant CUDA kernels are specialized for. The packed -/// K4/V3 block layout and every `turbo_*` kernel assume exactly this head -/// size; the kernels exit early (`if (hs != 128) return;`) for anything else. -pub const TURBOQUANT_HEAD_DIM: usize = 128; +/// Head dimensions the TurboQuant CUDA kernels are instantiated for. +/// +/// Re-exported from `mistralrs_quant::turboquant::cuda_tables`, which is also +/// where the unit test lives that pins this list to the `case` arms of the +/// kernel's own dispatch. Keeping one list on both sides of the FFI is what +/// makes an unsupported width a refusal instead of an untouched output buffer. +pub use mistralrs_quant::turboquant::TURBOQUANT_CUDA_HEAD_DIMS; #[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Default)] #[cfg_attr(feature = "pyo3_macros", pyo3::pyclass(eq, eq_int))] @@ -100,20 +103,53 @@ impl PagedCacheType { )) } + /// Head dims the **ambient default** is allowed to route onto TurboQuant. + /// + /// Deliberately narrower than [`TURBOQUANT_CUDA_HEAD_DIMS`]. Compiling a + /// kernel instantiation proves it *builds*; it does not prove it produces + /// correct KV on hardware, and none of 64/256/512 has ever executed. Since + /// [`PagedCacheType::TurboQuant`] is `#[default]`, widening the gate alone + /// would have silently moved every standard-layout model at those widths + /// onto unproven kernels — and silently dropped prefix caching with it + /// ([`Self::supports_prefix_cache`] is false for every TurboQuant variant). + /// Two regressions, neither visible at the call site. + /// + /// This repo has paid for that exact shape once already: FP8 KV shipped + /// default-on and unmeasured (wave43-BU) and every V4 request died. + /// + /// So: **compiled** widens what you may ask for; **measured** widens what + /// you get without asking. Move a width here when its hardware gate has + /// passed, not when its kernel compiles. + const TURBOQUANT_DEFAULT_HEAD_DIMS: [usize; 1] = [128]; + /// Whether the TurboQuant kernels support this model's KV geometry. - fn turboquant_supports_model(config: &dyn ModelConfigLike) -> bool { + /// + /// K and V are compressed independently, so they only have to be widths the + /// kernels are instantiated for — they no longer have to be *the same* + /// width, and no longer have to be 128. + /// + /// `forced` selects which set applies: an explicit request may use any + /// instantiated width, the ambient default only a measured one. + fn turboquant_supports_model(config: &dyn ModelConfigLike, forced: bool) -> bool { + let accepts = |head_dim: usize| { + if forced { + mistralrs_quant::turboquant::cuda_supports_head_dim(head_dim) + } else { + Self::TURBOQUANT_DEFAULT_HEAD_DIMS.contains(&head_dim) + } + }; matches!(config.kv_cache_layout(), KvCacheLayout::Standard) - && config.k_head_dim() == TURBOQUANT_HEAD_DIM - && config.v_head_dim() == TURBOQUANT_HEAD_DIM + && accepts(config.k_head_dim()) + && accepts(config.v_head_dim()) } /// Resolve this cache type against the target model's KV geometry. /// - /// The TurboQuant kernels only support the standard KV layout with - /// `head_dim == 128`: - /// * any other head size makes the `turbo_paged_attention.cu` kernels exit - /// early with their (uninitialized) F32 output buffer untouched — - /// silent garbage; + /// The TurboQuant kernels support the standard KV layout at any head + /// dimension in [`TURBOQUANT_CUDA_HEAD_DIMS`]: + /// * a head size with no instantiation falls through the kernel's dispatch + /// switch, leaving the (uninitialized) output buffer untouched — silent + /// garbage, which is why this gate runs before any allocation; /// * MLA-layout models (DeepSeek V2/V3, GLM4-MoE-lite) write through /// `concat_and_cache_mla`, which bails on a packed U8 cache at runtime. /// @@ -128,7 +164,11 @@ impl PagedCacheType { config: &dyn ModelConfigLike, explicitly_requested: bool, ) -> anyhow::Result { - if !self.is_turboquant() || Self::turboquant_supports_model(config) { + // `TurboQuant3`/`TurboQuantAggressive` are never a default anywhere, so + // selecting one is always an explicit act. Computed before the support + // check because it decides *which* set of head dims applies. + let forced = explicitly_requested || !matches!(self, PagedCacheType::TurboQuant); + if !self.is_turboquant() || Self::turboquant_supports_model(config, forced) { return Ok(self); } let reason = match config.kv_cache_layout() { @@ -136,13 +176,31 @@ impl PagedCacheType { "the model uses an MLA KV cache layout, which TurboQuant does not support" .to_string() } + // Split deliberately: "no kernel exists" and "a kernel exists but has + // never run on hardware, so the default will not choose it for you" + // are different facts, and collapsing them would tell a user to stop + // trying when the real answer is "opt in explicitly". + KvCacheLayout::Standard + if mistralrs_quant::turboquant::cuda_supports_head_dim(config.k_head_dim()) + && mistralrs_quant::turboquant::cuda_supports_head_dim(config.v_head_dim()) => + { + format!( + "the model has head_dim k={}/v={}: the TurboQuant kernels are instantiated \ + for it, but only head_dim in {:?} has been measured on hardware, so the \ + default will not route onto it. Pass `--pa-cache-type turboquant` to opt in \ + deliberately", + config.k_head_dim(), + config.v_head_dim(), + Self::TURBOQUANT_DEFAULT_HEAD_DIMS, + ) + } KvCacheLayout::Standard => format!( - "the model has head_dim k={}/v={}, but the TurboQuant kernels only support head_dim={TURBOQUANT_HEAD_DIM}", + "the model has head_dim k={}/v={}, but the TurboQuant kernels are instantiated \ + for head_dim in {TURBOQUANT_CUDA_HEAD_DIMS:?}", config.k_head_dim(), config.v_head_dim(), ), }; - let forced = explicitly_requested || !matches!(self, PagedCacheType::TurboQuant); if forced { anyhow::bail!( "PagedAttention cache type {self:?} was explicitly requested, but {reason}. \ @@ -569,23 +627,132 @@ mod cache_type_tests { PagedCacheType::TurboQuantAggressive, ]; - /// Supported geometry (standard layout, head_dim 128): every TurboQuant - /// preset resolves to itself, explicit or not. + /// Supported geometry, **explicitly requested**: every TurboQuant preset + /// resolves to itself at any head dim the kernels are instantiated for. + /// + /// Driven off `TURBOQUANT_CUDA_HEAD_DIMS` rather than a literal, so adding + /// a kernel instantiation extends this test automatically instead of + /// leaving the new width silently uncovered. #[test] - fn turboquant_stays_on_supported_geometry() { - let cfg = meta(128, 128, KvCacheLayout::Standard); - for t in TURBO_TYPES { - for explicit in [false, true] { - assert_eq!(t.resolve_for_model(&cfg, explicit).unwrap(), t); + fn explicit_turboquant_stays_on_any_instantiated_geometry() { + for head_dim in TURBOQUANT_CUDA_HEAD_DIMS { + let cfg = meta(head_dim, head_dim, KvCacheLayout::Standard); + for t in TURBO_TYPES { + assert_eq!( + t.resolve_for_model(&cfg, true).unwrap(), + t, + "head_dim={head_dim} is instantiated and an explicit request must be kept" + ); } } } - /// The default TurboQuant type falls back to Auto (instead of producing - /// silent garbage from the head_dim-128-only kernels) for other head dims. + /// **The regression this PR would otherwise have shipped.** + /// + /// `PagedCacheType::TurboQuant` is `#[default]`. Widening the acceptance + /// gate from `head_dim == 128` to every instantiated width would have moved + /// standard-layout models at 64/256/512 onto kernels that have never + /// executed, and silently dropped prefix caching with them — without anyone + /// asking for TurboQuant at all. + /// + /// The ambient default must therefore keep falling back to `Auto` at every + /// instantiated-but-unmeasured width, while the same width stays available + /// on explicit request (covered above). + #[test] + fn default_turboquant_does_not_route_onto_unmeasured_head_dims() { + // Non-vacuity guard. This loop skips every width that IS a default, so + // widening `TURBOQUANT_DEFAULT_HEAD_DIMS` to the full instantiated set + // would empty the loop and make the test pass while asserting nothing — + // which is precisely the silent-success shape this test exists to catch. + // Caught by mutating the const during review; keep this assertion. + let unmeasured: Vec = TURBOQUANT_CUDA_HEAD_DIMS + .into_iter() + .filter(|hd| !PagedCacheType::TURBOQUANT_DEFAULT_HEAD_DIMS.contains(hd)) + .collect(); + assert!( + !unmeasured.is_empty(), + "every instantiated head dim is now a default width, so this test asserts \ + nothing. If a hardware gate genuinely measured them all, delete this test \ + deliberately — do not let it pass vacuously." + ); + + for head_dim in unmeasured { + let cfg = meta(head_dim, head_dim, KvCacheLayout::Standard); + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, false) + .unwrap(), + PagedCacheType::Auto, + "head_dim={head_dim} is instantiated but unmeasured; the DEFAULT must not \ + route onto it" + ); + } + } + + /// Every measured width must still be taken by the ambient default — + /// the narrowing must not become "TurboQuant is off by accident". + #[test] + fn default_turboquant_still_taken_on_measured_head_dims() { + for head_dim in PagedCacheType::TURBOQUANT_DEFAULT_HEAD_DIMS { + assert!( + TURBOQUANT_CUDA_HEAD_DIMS.contains(&head_dim), + "head_dim={head_dim} is a default width but has no kernel instantiation" + ); + let cfg = meta(head_dim, head_dim, KvCacheLayout::Standard); + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, false) + .unwrap(), + PagedCacheType::TurboQuant, + "head_dim={head_dim} is measured and must still be the default" + ); + } + } + + /// DeepSeek-V4's width specifically. It is the reason the 128-only kernel + /// limit was lifted, so it gets its own named assertion — but on the + /// *explicit* path, because 512 has not been measured. + #[test] + fn turboquant_accepts_v4_head_dim_512_when_asked_for() { + assert!( + TURBOQUANT_CUDA_HEAD_DIMS.contains(&512), + "512 must stay in the instantiated set" + ); + let cfg = meta(512, 512, KvCacheLayout::Standard); + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, true) + .unwrap(), + PagedCacheType::TurboQuant + ); + // ...and is NOT taken by the ambient default. + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, false) + .unwrap(), + PagedCacheType::Auto + ); + } + + /// K and V are compressed independently, so mixed widths are fine as long + /// as each is instantiated — on the explicit path. + #[test] + fn turboquant_accepts_mixed_but_instantiated_head_dims() { + let cfg = meta(512, 128, KvCacheLayout::Standard); + assert_eq!( + PagedCacheType::TurboQuant + .resolve_for_model(&cfg, true) + .unwrap(), + PagedCacheType::TurboQuant + ); + } + + /// The default TurboQuant type falls back to Auto (instead of falling + /// through the kernel's dispatch switch and leaving the output buffer + /// untouched) for head dims with no instantiation. #[test] fn default_turboquant_falls_back_on_unsupported_head_dim() { - for head_dim in [64, 96, 192, 256] { + for head_dim in [96, 192, 320, 1024] { let cfg = meta(head_dim, head_dim, KvCacheLayout::Standard); assert_eq!( PagedCacheType::TurboQuant @@ -595,8 +762,8 @@ mod cache_type_tests { "head_dim={head_dim} must fall back" ); } - // Mixed k/v head dims must also fall back. - let cfg = meta(128, 64, KvCacheLayout::Standard); + // A mixed pair still falls back when either half is uninstantiated. + let cfg = meta(128, 192, KvCacheLayout::Standard); assert_eq!( PagedCacheType::TurboQuant .resolve_for_model(&cfg, false) @@ -629,7 +796,7 @@ mod cache_type_tests { /// models rather than silently switching caches. #[test] fn explicit_turboquant_errors_on_unsupported_model() { - let cfg = meta(64, 64, KvCacheLayout::Standard); + let cfg = meta(192, 192, KvCacheLayout::Standard); for t in TURBO_TYPES { let err = t.resolve_for_model(&cfg, true).unwrap_err(); assert!( diff --git a/mistralrs-core/src/paged_attention/mod.rs b/mistralrs-core/src/paged_attention/mod.rs index 803d84ab4..f13c1a1ee 100644 --- a/mistralrs-core/src/paged_attention/mod.rs +++ b/mistralrs-core/src/paged_attention/mod.rs @@ -131,7 +131,15 @@ mod cache_config_tests { use super::*; use candle_core::Device; - fn small_head_dim_model() -> ModelConfigMetadata { + /// A head dim the TurboQuant kernels have no instantiation for. 192 is a + /// real head dim (it is what the MLA models report) and is deliberately + /// *not* in `TURBOQUANT_CUDA_HEAD_DIMS`, so it exercises the fallback + /// without pretending an unrealistic geometry. + fn uninstantiated_head_dim_model() -> ModelConfigMetadata { + assert!( + !mistralrs_quant::turboquant::cuda_supports_head_dim(192), + "192 gained a kernel instantiation; pick another width for this fixture" + ); ModelConfigMetadata { max_seq_len: 4096, num_layers: 2, @@ -139,15 +147,16 @@ mod cache_config_tests { num_kv_heads: 2, num_attn_heads: 4, sliding_window: None, - k_head_dim: 64, - v_head_dim: 64, + k_head_dim: 192, + v_head_dim: 192, kv_cache_layout: KvCacheLayout::Standard, } } - /// A defaulted TurboQuant cache type on a head_dim-64 model must fall back - /// to `Auto` at cache-config construction time — the returned config (used - /// for both sizing and `CacheEngine` allocation) carries the resolved type. + /// A defaulted TurboQuant cache type on a model whose head dim has no + /// kernel instantiation must fall back to `Auto` at cache-config + /// construction time — the returned config (used for both sizing and + /// `CacheEngine` allocation) carries the resolved type. #[test] fn calculate_cache_config_falls_back_for_unsupported_turboquant() -> anyhow::Result<()> { let device = Device::Cpu; @@ -157,7 +166,7 @@ mod cache_config_tests { DType::F16, PagedCacheType::TurboQuant, false, - &small_head_dim_model(), + &uninstantiated_head_dim_model(), &device, &[Some(device.clone())], true, @@ -179,7 +188,7 @@ mod cache_config_tests { DType::F16, PagedCacheType::TurboQuant, true, - &small_head_dim_model(), + &uninstantiated_head_dim_model(), &device, &[Some(device.clone())], true, @@ -307,9 +316,11 @@ pub fn calculate_cache_config( if !SUPPORTED_BLOCK_SIZE.contains(&block_size) { anyhow::bail!("Block size must be in {SUPPORTED_BLOCK_SIZE:?}, got {block_size}"); } - // TurboQuant only supports standard-layout models with head_dim == 128; - // resolve to a supported cache type (or error, if the user explicitly - // forced TurboQuant) before any sizing or allocation happens. + // TurboQuant needs a standard KV layout, and a head dim its kernels are + // instantiated for — narrowed further to the *measured* widths when the + // type is the ambient default rather than an explicit request. Resolve to a + // supported cache type (or error, if the user explicitly forced TurboQuant) + // before any sizing or allocation happens. let cache_type = cache_type.resolve_for_model(config, cache_type_explicit)?; let dtype = cache_type.to_dtype(dtype); let dtype_size = dtype.size_in_bytes(); diff --git a/mistralrs-core/src/pipeline/mod.rs b/mistralrs-core/src/pipeline/mod.rs index 37110ed39..7fa8c878f 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. /// @@ -545,6 +551,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, _)| 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)?; + 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, @@ -982,6 +1070,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 +1111,54 @@ 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; + + // 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 @@ -1204,6 +1350,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) } @@ -1800,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). @@ -1925,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. /// diff --git a/mistralrs-core/src/pipeline/mtp_pipeline.rs b/mistralrs-core/src/pipeline/mtp_pipeline.rs index b026a8df6..169331801 100644 --- a/mistralrs-core/src/pipeline/mtp_pipeline.rs +++ b/mistralrs-core/src/pipeline/mtp_pipeline.rs @@ -1522,6 +1522,16 @@ impl AcceptanceTelemetry { for (b, per_b) in self.snapshot_by_batch() { tracing::info!(target: "mtp_speculative", "{}", per_b.marker(&format!("b={b}"))); } + // 🔑 On the SAME fence, deliberately. `tok_per_step` and aggregate + // throughput are both read against how wide the steps actually were, + // and both schedulers run one cache-length bucket per step — so a batch + // can be serialised for reasons that have nothing to do with MTP. A + // harness that differences these counters across a wall-clock boundary + // gets the scheduler's window and MTP's window identical only if they + // are emitted together. See `scheduler::bucket_telemetry`. + if let Some(sched) = crate::scheduler::sched_bucket_marker() { + tracing::info!(target: "mtp_speculative", "{sched}"); + } tracing::info!(target: "mtp_speculative", "{}", snap.report_line()); // The per-position breakdown, in prose, right where the aggregate rate // is. This is the line the acceptance-gap diagnosis needs: the scalar diff --git a/mistralrs-core/src/pipeline/normal.rs b/mistralrs-core/src/pipeline/normal.rs index b93ae48df..eaaaddcb9 100644 --- a/mistralrs-core/src/pipeline/normal.rs +++ b/mistralrs-core/src/pipeline/normal.rs @@ -96,6 +96,13 @@ pub struct NormalPipeline { /// load time; graph capture is deferred until first decode call. #[cfg(feature = "cuda")] autonomous_runner: Option, + /// ArcInfer/ArcGraph device decode loop — replays the captured graph, + /// samples and scatters the token on device, and never synchronizes. Built + /// lazily on the first admitted decode step and then kept, because every + /// buffer it points at must stay at the address the graph baked. + /// Opt-in via `ARC_GRAPH_DEVICE_LOOP`; `None` until then. + #[cfg(feature = "cuda")] + device_decode_loop: Option, } /// What a forward pass should do with candle's caching allocator before it @@ -160,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, @@ -1986,6 +1989,10 @@ impl Loader for NormalLoader { // lazy init (and graph capture) once a real decode batch arrives. #[cfg(feature = "cuda")] autonomous_runner: None, + // Needs the captured graph's exec handle and output tensor, so it + // cannot exist before capture. Built on the first admitted step. + #[cfg(feature = "cuda")] + device_decode_loop: None, }))) } @@ -2108,6 +2115,205 @@ impl MetadataMixin for NormalPipeline { } } +#[cfg(feature = "cuda")] +/// Drive one burst of the ArcInfer/ArcGraph device decode loop. +/// +/// Returns `Some(logits)` when the step is covered — either because a burst +/// just ran or because a previous burst already produced this step's token +/// — and `None` to fall back. +/// +/// The tokens go to `arc_cuda_graph`'s pending queue and are drained one at +/// a time by `pipeline/sampling.rs`, so the engine's per-token `Sequence` +/// bookkeeping (stop strings, streaming, EOS) is unchanged. +/// +/// ⚠️ On the `pending > 0` path this returns the captured graph's output +/// tensor **without launching anything**, so those logits are the last +/// burst's, not this step's. That is sound only because the sampling hook +/// short-circuits before reading them and, at batch 1, +/// `pipeline/mod.rs:538` does no host copy. It is not a general-purpose +/// return value, which is why this whole path refuses anything that would +/// actually read the logits (`return_logprobs`, logits processors, +/// penalties, grammars). +fn device_decode_burst( + device: &Device, + runner: &mut arc_cuda_graph::CudaGraphRunner, + loop_slot: &mut Option, + bs: usize, + seq_len: usize, + current_position: u32, + input_ids_pinned: &Tensor, +) -> Option { + use arc_cuda_graph as acg; + + // Stand aside for the replay-verification lane (PR #205). This branch sits + // EARLIER in the chain than the `runner.replay(bs)` arm, so without this it + // would take every step and `ARC_GRAPH_VERIFY_ALWAYS=1` would silently + // verify nothing — the operator would get a clean diagnostic that never + // ran. A verification switch that is quietly bypassed is worse than one + // that is off. + if std::env::var("ARC_GRAPH_VERIFY_ALWAYS").as_deref() == Ok("1") { + return None; + } + + let (_exec, out) = runner.replay_handles(bs)?; + + // A burst already produced this step's token: no launch, no sync, no + // 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 { + 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 + // last dimension rather than guessed: `models/deepseek4.rs` builds that + // mask as `[B, 1, 1, cfg_full.sliding_window]`, and `cfg_full` is a + // model fact the pipeline cannot otherwise reach. + let capacity = match crate::layers::graph_mode_mask() { + Some(m) => *m.dims().last().unwrap_or(&0), + None => { + // The mask is what keeps the decode arm from reading unwritten + // KV slots; without it we cannot bound the burst. + acg::stand_down(); + return None; + } + }; + let vocab = *out.dims().last().unwrap_or(&0); + if capacity == 0 || vocab == 0 { + acg::stand_down(); + return None; + } + + let cfg = acg::DeviceLoopConfig { + position_limit: capacity as u32, + ..acg::DeviceLoopConfig::new(vocab, -1, capacity as u32) + }; + + let facts = acg::AdmissionFacts { + batch_size: bs, + seq_len, + graph_ready: runner.replay_output_trusted(), + // Established by `sampling.rs` from the Sequence on the previous + // step; false until one has confirmed it. + greedy: acg::device_loop_eligible(), + return_logprobs: false, + host_side_logits_work: false, + current_position, + }; + let steps = match acg::admit(&facts, &cfg, acg::device_loop_enabled()) { + Ok(s) => s, + Err(reason) => { + tracing::debug!("ArcGraph device loop declined this step: {reason}"); + acg::stand_down(); + return None; + } + }; + + if loop_slot.is_none() { + let positions = crate::layers::graph_mode_positions(); + match acg::DeviceDecodeLoop::new( + device, + _exec, + runner.stream(), + &out, + input_ids_pinned, + positions.as_ref(), + cfg, + ) { + Ok(l) => { + tracing::info!( + "ArcGraph device decode loop ENGAGED: burst={} vocab={vocab} \ + kv_window={capacity}. {}", + cfg.burst, + l.status_line() + ); + *loop_slot = Some(l); + } + Err(e) => { + acg::kill_device_loop(&format!("could not build the device loop: {e}")); + return None; + } + } + } + let dl = loop_slot.as_mut()?; + + match dl.run(steps) { + Ok(outcome) => { + // `first()` rather than `[0]`: a burst that returned no rows would + // be a bug, but it must not panic the engine thread mid-request. + let row = match outcome.tokens.first() { + Some(r) => r, + None => { + acg::kill_device_loop("burst returned no token rows"); + *loop_slot = None; + return None; + } + }; + acg::push_pending_tokens(row); + // ── ALIASING CONTRACT, ENFORCED ───────────────────────────── + // The tensor we are about to return aliases the graph-owned + // storage that the NEXT `cuGraphLaunch` overwrites. It is only + // safe to hand back because the sampling hook short-circuits + // before anything reads it — and that hook fires if and only if a + // token is pending. So the invariant is not "we reasoned it + // through", it is checked here: no pending token, no aliasing + // tensor. `push_pending_tokens` drops negative ids, so a row that + // was entirely garbage lands here rather than silently arriving at + // a `sample_sequence` that would then read clobbered logits and + // return a plausible wrong token. + if acg::pending_token_count() == 0 { + acg::kill_device_loop( + "burst produced no usable token, so the aliasing logits tensor must not be \ + returned", + ); + *loop_slot = None; + return None; + } + runner.record_device_loop_steps(outcome.steps_launched as u64); + if outcome.hit_eos { + tracing::debug!( + "ArcGraph device loop: EOS at token {} of a {}-step burst; the remaining \ + steps were already enqueued and are discarded", + row.len(), + outcome.steps_launched + ); + } + Some(out) + } + Err(e) => { + // A burst failure means the graph input buffer may hold an + // uncommitted or stale id, so this cannot be retried. + acg::kill_device_loop(&format!("burst failed: {e}")); + *loop_slot = None; + None + } + } +} + #[async_trait::async_trait] impl Pipeline for NormalPipeline { fn forward_inputs( @@ -2202,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 @@ -2256,6 +2460,11 @@ impl Pipeline for NormalPipeline { } } } + // Bound before the chain: `self.device()` returns by value, + // and taking a reference to it inside an expression that + // also takes `&mut self.device_decode_loop` would borrow + // `self` both ways at once. + let dev_for_loop = self.device(); let captured: Option = if probe && seq_len == 1 && self.cuda_graph_runner.is_some() @@ -2275,6 +2484,12 @@ impl Pipeline for NormalPipeline { // per-forward alloc count (eager warmups only // reach peak-live; capture needs every alloc // distinct). Output is this step's logits (eager). + // Snapshot the allocator's size profile so we can see + // whether this pass teaches it anything new. + let demand_before = match self.device() { + candle_core::Device::Cuda(cd) => cd.capture_alloc_demand(), + _ => Vec::new(), + }; if let candle_core::Device::Cuda(cd) = self.device() { cd.set_capture_mode(true); } @@ -2312,6 +2527,40 @@ impl Pipeline for NormalPipeline { if let candle_core::Device::Cuda(cd) = self.device() { cd.set_capture_mode(false); } + // Keep warming while the profile is still growing. + // One pass is only enough if every decode step + // allocates the same sizes; V4's compressor tail + // cycles through `ratio` widths, so pass k presents a + // size pass k-1 never did and the profile keeps + // growing until the whole cycle has been walked. A + // pass that adds nothing is the signal to stop — and + // it costs one eager decode step to learn that. + if let candle_core::Device::Cuda(cd) = self.device() { + let demand_after = cd.capture_alloc_demand(); + if demand_after != demand_before { + let granted = runner.grant_extra_deferred_pass(); + tracing::info!( + "ARC prewarm: warm pass grew the alloc profile to {} \ + distinct sizes — {} (extra budget {}, {} pass(es) owed)", + demand_after.len(), + if granted { + "granting another pass" + } else { + "extra budget exhausted; the capture-miss assert \ + decides from here" + }, + runner.deferred_extra_budget(), + runner.deferred_passes_remaining(), + ); + } else { + tracing::info!( + "ARC prewarm: alloc profile CONVERGED at {} distinct \ + sizes ({} pass(es) still owed)", + demand_after.len(), + runner.deferred_passes_remaining(), + ); + } + } match out { Ok(o) => { tracing::info!( @@ -2377,6 +2626,55 @@ impl Pipeline for NormalPipeline { (PROBE — not a fix; the buffer still allocates per step)" ); } + // ── EXACT PRE-WARM ──────────────────────────── + // The warm passes leave the free pool holding + // exactly what the steps they ran on asked for. The + // step about to be CAPTURED is a different step: any + // size whose demand is higher here, or that only + // appears here, is served by `cuMemAllocAsync` from + // the graph's private pool and becomes an unstable + // graph memory node — CUDA_ERROR_ILLEGAL_ADDRESS on + // the first cuGraphLaunch, and a poisoned context + // (reported as glibc heap corruption) after it. + // + // So top every observed size up to its observed + // demand plus slack. Slack is what covers a captured + // step that allocates a little more of some size + // than any warm step did; it is buffers, not bytes, + // and the pool is ~315 MiB, so it is cheap. + // + // Must run with capture mode OFF: a buffer released + // while capturing parks in `deferred` and cannot be + // served, so seeding inside the window warms nothing. + if let candle_core::Device::Cuda(cd) = self.device() { + let slack = std::env::var("ARC_GRAPH_PREWARM_SLACK") + .ok() + .and_then(|v| v.parse::().ok()) + .unwrap_or(4); + let demand = cd.capture_alloc_demand(); + let want: usize = demand.iter().map(|(_, n)| n).sum(); + match cd.prewarm_alloc_cache(slack) { + Ok((sizes, buffers)) => tracing::info!( + "ARC prewarm: profile = {} distinct sizes / {want} \ + allocations per captured forward; topped up {sizes} \ + size(s) with {buffers} buffer(s) at slack={slack}", + demand.len(), + ), + Err(e) => { + tracing::error!( + "ARC prewarm: could not top up the alloc cache: {e}. \ + Capturing now would record unstable memory nodes, so \ + capture is disabled for this run." + ); + runner.disable(); + } + } + // Scope the miss ledger to THIS capture. Without + // the reset it would still carry the misses the + // warm passes made on purpose, and those are + // harmless — no graph is being recorded then. + cd.reset_capture_misses(); + } // CAPTURE: frees are deferred so every allocation // is a stable cache hit (no within-capture // aliasing, no unstable graph memory nodes). @@ -2403,127 +2701,179 @@ impl Pipeline for NormalPipeline { &flash_meta, ) { Ok(output) => { - tracing::info!( - "ARC capture: V4 forward RECORDED (bs={bs}); instantiating + launching" + // ── THE GATE ────────────────── + // Every allocation that missed the + // warm pool during this capture was + // served from the graph's private + // pool and recorded as a memory node + // whose address the first launch has + // to back. One is enough to fault. + // Refusing here is strictly better + // than instantiating: the fault is + // asynchronous and poisons the CUDA + // context, taking the process with + // it some allocations later. + let misses = match self.device() { + candle_core::Device::Cuda(cd) => { + cd.capture_misses() + } + _ => Vec::new(), + }; + if !misses.is_empty() { + let total: usize = + misses.iter().map(|(_, n)| n).sum(); + tracing::error!( + "ARC capture: REFUSING to instantiate — \ + {total} allocation(s) missed the warm pool \ + during capture, at (bytes, count) {misses:?}. \ + Each is an unstable graph memory node. The \ + pre-warm profile did not cover this step, so \ + these sizes are step-dependent beyond the \ + warm cycle: the fix is a shape-constant \ + buffer at those sizes (or allocator \ + size-class bucketing), not more warmup. \ + Falling back to eager." ); - match runner.end_capture_and_cache(bs, output, gp, op) { - Ok(out) => { - tracing::info!( + // Drop the captured output BEFORE + // cancelling: its storage may live + // in the private pool, and + // cancel_capture drains the alloc + // cache before destroying that + // pool. Dropping after the drain + // would file a dangling pointer + // into the free list — the exact + // use-after-free the drain exists + // to prevent. + drop(output); + runner.cancel_capture(gp, op); + runner.disable(); + None + } else { + tracing::info!( + "ARC capture: V4 forward RECORDED (bs={bs}), zero capture-time allocator misses; instantiating + launching" + ); + match runner + .end_capture_and_cache(bs, output, gp, op) + { + Ok(out) => { + tracing::info!( "ARC capture: graph CAPTURED + launched OK" ); - if arc_layer_bisect() { - let bufs = - crate::layers::arc_layer_trace_take(); - tracing::error!( - "ARC BISECT: holding {} captured layer buffers", - bufs.as_ref() - .map(|b| b.len()) - .unwrap_or(0) - ); - ARC_GRAPH_LAYER_BUFS - .with(|c| *c.borrow_mut() = bufs); - } - // ── SAME-STEP PROBE ───────────── - // Separates the two ways a graph can - // be wrong, which the replay-step - // comparison cannot tell apart. - // - // Here the position buffer, the token - // ids and the KV ring are EXACTLY what - // capture recorded against — nothing - // has advanced. So: - // Δ ≈ 0 ⇒ the graph is wired right - // and the replay-step Δ is - // stale per-step state. - // Δ large ⇒ a plain wiring fault - // (wrong buffer / wrong - // node); position state is - // innocent. - // `error!` on purpose: this is the - // measurement the run exists for and - // must never be filtered out. - if std::env::var("ARC_GRAPH_SAMESTEP") - .as_deref() - == Ok("1") - { - if let candle_core::Device::Cuda(cd) = - self.device() - { - cd.set_capture_mode(false); - } - match self.model.forward( - &input_ids, - &seqlen_offsets, - context_lens.clone(), - position_ids.clone(), - paged_attn_meta - .as_ref() - .map(|(a, b)| (a.clone(), b)), - &flash_meta, - ) { - Ok(eager_same) => { - match max_abs_diff( - &out, - &eager_same, - ) { - Ok(d) => tracing::error!( - "ARC SAMESTEP: first-launch vs eager, SAME step: max|Δ|={d:.3e}" - ), - Err(e) => tracing::error!( - "ARC SAMESTEP: first-launch compare failed: {e}" - ), - } - // Keep this step's eager - // logits: every later replay - // is measured against them - // too, which is what tells - // "frozen" from "partially - // advancing". - ARC_CAPTURE_STEP_EAGER.with( - |c| { - *c.borrow_mut() = - Some( - eager_same - .clone(), - ); - }, + if arc_layer_bisect() { + let bufs = + crate::layers::arc_layer_trace_take( ); - // Second launch of the same - // graph, still same step. - // Distinguishes "the graph - // is wrong" from "re-launch - // specifically is wrong". - match runner.replay(bs) { - Ok(again) => { - match max_abs_diff( - &again, - &eager_same, - ) { - Ok(d) => tracing::error!( - "ARC SAMESTEP: re-launch vs eager, SAME step: max|Δ|={d:.3e}" - ), - Err(e) => tracing::error!( - "ARC SAMESTEP: re-launch compare failed: {e}" - ), + tracing::error!( + "ARC BISECT: holding {} captured layer buffers", + bufs.as_ref() + .map(|b| b.len()) + .unwrap_or(0) + ); + ARC_GRAPH_LAYER_BUFS + .with(|c| *c.borrow_mut() = bufs); + } + // ── SAME-STEP PROBE ───────────── + // Separates the two ways a graph can + // be wrong, which the replay-step + // comparison cannot tell apart. + // + // Here the position buffer, the token + // ids and the KV ring are EXACTLY what + // capture recorded against — nothing + // has advanced. So: + // Δ ≈ 0 ⇒ the graph is wired right + // and the replay-step Δ is + // stale per-step state. + // Δ large ⇒ a plain wiring fault + // (wrong buffer / wrong + // node); position state is + // innocent. + // `error!` on purpose: this is the + // measurement the run exists for and + // must never be filtered out. + if std::env::var("ARC_GRAPH_SAMESTEP") + .as_deref() + == Ok("1") + { + if let candle_core::Device::Cuda(cd) = + self.device() + { + cd.set_capture_mode(false); + } + match self.model.forward( + &input_ids, + &seqlen_offsets, + context_lens.clone(), + position_ids.clone(), + paged_attn_meta + .as_ref() + .map(|(a, b)| (a.clone(), b)), + &flash_meta, + ) { + Ok(eager_same) => { + match max_abs_diff( + &out, + &eager_same, + ) { + Ok(d) => tracing::error!( + "ARC SAMESTEP: first-launch vs eager, SAME step: max|Δ|={d:.3e}" + ), + Err(e) => tracing::error!( + "ARC SAMESTEP: first-launch compare failed: {e}" + ), + } + // Keep this step's eager + // logits: every later replay + // is measured against them + // too, which is what tells + // "frozen" from "partially + // advancing". + ARC_CAPTURE_STEP_EAGER.with( + |c| { + *c.borrow_mut() = + Some( + eager_same + .clone(), + ); + }, + ); + // Second launch of the same + // graph, still same step. + // Distinguishes "the graph + // is wrong" from "re-launch + // specifically is wrong". + match runner.replay(bs) { + Ok(again) => { + match max_abs_diff( + &again, + &eager_same, + ) { + Ok(d) => tracing::error!( + "ARC SAMESTEP: re-launch vs eager, SAME step: max|Δ|={d:.3e}" + ), + Err(e) => tracing::error!( + "ARC SAMESTEP: re-launch compare failed: {e}" + ), + } } + Err(e) => tracing::error!( + "ARC SAMESTEP: second launch failed: {e}" + ), } - Err(e) => tracing::error!( - "ARC SAMESTEP: second launch failed: {e}" - ), } + Err(e) => tracing::error!( + "ARC SAMESTEP: eager forward failed: {e}" + ), } - Err(e) => tracing::error!( - "ARC SAMESTEP: eager forward failed: {e}" - ), } + Some(out) } - Some(out) - } - Err(e) => { - tracing::warn!( + Err(e) => { + tracing::warn!( "ARC capture: instantiate/launch failed: {e}; eager" ); - None + None + } } } } @@ -2555,6 +2905,32 @@ impl Pipeline for NormalPipeline { cd.set_capture_mode(false); } cap_result + } else if runner.has_graph(bs) + && arc_cuda_graph::device_loop_enabled() + && !arc_cuda_graph::device_loop_killed() + && device_decode_burst( + &dev_for_loop, + &mut runner, + &mut self.device_decode_loop, + bs, + seq_len, + seqlen_offsets.first().copied().unwrap_or(0) as u32, + &input_ids, + ) + .is_some() + { + // ArcInfer/ArcGraph device decode loop: the burst + // already ran (or a previous one covered this step), + // the token was sampled and scattered on device, and + // `pipeline/sampling.rs` will take it from the ring. + // No sync, no host sample, no D2H. + // + // The logits handed back are the captured graph's + // output tensor. They are NOT read: the sampling + // hook short-circuits first, and at batch 1 + // `pipeline/mod.rs:538` does no host copy. Admission + // refuses everything that would read them. + runner.replay_handles(bs).map(|(_, out)| out) } else if runner.has_graph(bs) { let t = std::time::Instant::now(); let replayed = runner.replay(bs); @@ -3021,6 +3397,16 @@ impl Pipeline for NormalPipeline { top_p, frequency_penalty, presence_penalty, + // mistral.rs uses <=0 to mean "disabled"; the fused sampler + // spells that -1 (`sampling_cpu.rs:16`). + top_k: { + let tk = first_sampler.top_k(); + if tk <= 0 { + -1 + } else { + tk as i32 + } + }, greedy, }; @@ -3412,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 b413e380d..3e99d7212 100644 --- a/mistralrs-core/src/pipeline/sampling.rs +++ b/mistralrs-core/src/pipeline/sampling.rs @@ -527,6 +527,106 @@ pub async fn sample_and_add_toks( Ok(()) } +/// Is this sequence one the ArcGraph device decode loop may drive, and does it +/// already have a token waiting? +/// +/// Two jobs, deliberately in one place. It **publishes** eligibility, because +/// `Pipeline::forward_inputs` decides whether to run a burst but is handed only +/// `inputs` and `return_raw_logits` — it never sees a `Sequence`, so it cannot +/// ask whether sampling is greedy. And it **takes** the token a burst already +/// produced. +/// +/// Greedy only. The device sampler runs Splitmix64 per row against the host's +/// Isaac64, so under any stochastic configuration the two draw different — +/// individually valid — tokens, which would silently break seeded +/// reproducibility. Under greedy they agree exactly, which is also what makes +/// the path verifiable against the eager one. +/// +/// The returned `Logprobs` is not a lossy stand-in: `logprob: 0.0` is exactly +/// what the greedy host path returns (`sampler.rs:1490`), and `bytes` / +/// `top_logprobs` are read only when `return_logprobs` is set, which this +/// refuses. +// `&mut` only because `Sequence::sampler()` takes `&mut self` (`sequence.rs:885`) +// to hand back a clone of an `Arc`. Nothing here mutates the sequence. +// +// 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, + sample_speculative: bool, +) -> 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(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(); + 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) + } + } +} + /// Async sample optionally adding to trie. #[allow(clippy::too_many_arguments)] pub async fn sample_sequence( @@ -538,6 +638,17 @@ pub async fn sample_sequence( sample_speculative: bool, multiple_sequences: bool, ) -> Result { + // ── ArcInfer/ArcGraph device decode loop ──────────────────────────────── + // If a burst already sampled this token ON DEVICE, take it and skip the + // host argmax entirely. That D2H (`sampler.rs:1479`) is the last + // synchronization in the decode step, so this is the half of the change + // that makes removing `replay()`'s `cudaStreamSynchronize` worth anything. + if let Some(pre_sampled) = + device_loop_pre_sampled_token(seq, return_logprobs, sample_speculative)? + { + return Ok(pre_sampled); + } + // Both of these open and close inside a single poll, before the first // `.await`, so they nest correctly under `sample.join_all` even when B // futures are being polled in turn on this thread. @@ -926,3 +1037,323 @@ 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. `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( + 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(), + id, + 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(1); + 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(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 + // 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(3); + seq.set_state(state); + assert_eq!( + arc_cuda_graph::pending_token_count(), + 0, + "transition to {state:?} left parked tokens for the next sequence" + ); + } + } + + /// 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/src/sampler.rs b/mistralrs-core/src/sampler.rs index b282841c6..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,41 @@ 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 + /// nothing that has to touch the logits on the host? + /// + /// This is the exact precondition for ArcGraph's device decode loop to + /// substitute its on-device argmax for the host one + /// (`pipeline/sampling.rs`). It deliberately mirrors the conditions inside + /// `sample` rather than approximating them: `logits_bias` and `top_nsigma` + /// are applied *before* the `trivial` gate, so a sampler carrying either + /// still needs the host even though the gate itself ignores them. + /// + /// Greedy is also the only mode where device and host agree *exactly* — the + /// device sampler runs Splitmix64 per row against the host's Isaac64, so + /// any stochastic mode would draw a different, equally valid token and + /// silently break seeded reproducibility. + pub fn is_greedy_trivial(&self) -> bool { + self.temperature.is_none() + && self.logits_processors.is_empty() + && !self.has_penalties() + && self.logits_bias.is_none() + && self.top_nsigma.is_none() } /// True when this sampler is exactly `argmax` over the model's **raw** @@ -539,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 @@ -1164,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 { @@ -1209,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); @@ -2871,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/src/scheduler/bucket_telemetry.rs b/mistralrs-core/src/scheduler/bucket_telemetry.rs new file mode 100644 index 000000000..2b0985a4f --- /dev/null +++ b/mistralrs-core/src/scheduler/bucket_telemetry.rs @@ -0,0 +1,241 @@ +//! Parent system: ArcInfer / ArcSched +//! +//! Process-wide counters for the one thing that decides how wide a decode step +//! actually is: **bucketing**. +//! +//! Both schedulers partition the running set by an exact cache length and run +//! **one bucket per step**, preempting the rest +//! (`default_scheduler.rs:bucket_and_waitlist_seqs_waiting`, +//! `paged_attention/scheduler.rs`). So a throughput number measured at "B=128" +//! is not necessarily a number about 128-wide steps — if the batch shattered +//! into 40 buckets, the engine ran 3-wide steps and the label lies. +//! +//! That is not a hypothetical confound for the per-sequence-advance A/B: the +//! whole point of ragged admission ([`super::RaggedAdmission`]) is to collapse +//! those buckets into one, so an aggregate throughput delta between the ON and +//! OFF arms could come from the KV mechanism, from the bucketing, or from both. +//! Without these counters the two are indistinguishable, and "aggregate did not +//! move" would be unattributable. +//! +//! The marker is emitted on the same log fence as `MTP[agg]`, so a harness that +//! differences cumulative counters across a wall-clock boundary gets the +//! scheduler's numbers for exactly the window it got MTP's. + +use std::sync::atomic::{AtomicUsize, Ordering}; + +/// Cumulative bucketing counters. Monotone; difference two snapshots to get a +/// window. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct SchedBuckets { + /// Bucketing calls with at least one running sequence. + pub calls: usize, + /// Sum of the number of buckets formed. + pub buckets: usize, + /// Sum of the running-set size *before* selection — the batch the engine + /// was offered. + pub offered: usize, + /// Sum of the winning bucket's size — the batch the engine actually ran. + pub chosen: usize, + /// Calls that formed more than one bucket, i.e. that preempted somebody. + pub shattered: usize, +} + +impl SchedBuckets { + /// `self - earlier`, for differencing across a fence. Saturating, so a + /// counter reset between snapshots degrades to zero rather than to a + /// gigantic wrapped number that would read as a spectacular result. + #[must_use] + pub fn since(&self, earlier: &Self) -> Self { + Self { + calls: self.calls.saturating_sub(earlier.calls), + buckets: self.buckets.saturating_sub(earlier.buckets), + offered: self.offered.saturating_sub(earlier.offered), + chosen: self.chosen.saturating_sub(earlier.chosen), + shattered: self.shattered.saturating_sub(earlier.shattered), + } + } + + #[allow(clippy::cast_precision_loss)] + fn per_call(&self, total: usize) -> Option { + (self.calls > 0).then(|| total as f64 / self.calls as f64) + } + + /// Mean buckets formed per scheduling step. 1.0 means the batch was never + /// split; 40.0 at a nominal B=128 means the engine ran ~3-wide steps. + #[must_use] + pub fn buckets_per_step(&self) -> Option { + self.per_call(self.buckets) + } + + /// Mean size of the bucket that actually ran — the **real** batch width, + /// as opposed to the number of sequences in flight. + #[must_use] + pub fn running_bucket_size(&self) -> Option { + self.per_call(self.chosen) + } + + /// Mean number of running sequences offered to the scheduler per step. + #[must_use] + pub fn offered_per_step(&self) -> Option { + self.per_call(self.offered) + } + + /// Mean sequences preempted per step (`offered - chosen`). + #[must_use] + pub fn preempted_per_step(&self) -> Option { + self.per_call(self.offered.saturating_sub(self.chosen)) + } + + /// Fraction of steps that split the batch at all. + #[allow(clippy::cast_precision_loss)] + #[must_use] + pub fn shattered_frac(&self) -> Option { + (self.calls > 0).then(|| self.shattered as f64 / self.calls as f64) + } + + /// The machine-greppable one-liner, in the project's marker convention + /// (`MTP[...]`, `SPEED[...]`). Every raw count is on the line so the ratios + /// are auditable without trusting the formatter. + #[must_use] + pub fn marker(&self, scope: &str) -> String { + let fmt = |v: Option| v.map_or_else(|| "n/a".to_string(), |x| format!("{x:.4}")); + format!( + "SCHED[{scope}] calls={} buckets={} offered={} chosen={} shattered={} \ + buckets_per_step={} running_bucket_size={} offered_per_step={} \ + preempted_per_step={} shattered_frac={}", + self.calls, + self.buckets, + self.offered, + self.chosen, + self.shattered, + fmt(self.buckets_per_step()), + fmt(self.running_bucket_size()), + fmt(self.offered_per_step()), + fmt(self.preempted_per_step()), + fmt(self.shattered_frac()), + ) + } +} + +static CALLS: AtomicUsize = AtomicUsize::new(0); +static BUCKETS: AtomicUsize = AtomicUsize::new(0); +static OFFERED: AtomicUsize = AtomicUsize::new(0); +static CHOSEN: AtomicUsize = AtomicUsize::new(0); +static SHATTERED: AtomicUsize = AtomicUsize::new(0); + +/// Record one bucketing decision. Called once per scheduling step that had +/// anything to schedule. +pub fn record_bucketing(buckets: usize, offered: usize, chosen: usize) { + if offered == 0 { + return; + } + CALLS.fetch_add(1, Ordering::Relaxed); + BUCKETS.fetch_add(buckets, Ordering::Relaxed); + OFFERED.fetch_add(offered, Ordering::Relaxed); + CHOSEN.fetch_add(chosen, Ordering::Relaxed); + if buckets > 1 { + SHATTERED.fetch_add(1, Ordering::Relaxed); + } +} + +/// Snapshot of the process-wide counters. +#[must_use] +pub fn sched_buckets() -> SchedBuckets { + SchedBuckets { + calls: CALLS.load(Ordering::Relaxed), + buckets: BUCKETS.load(Ordering::Relaxed), + offered: OFFERED.load(Ordering::Relaxed), + chosen: CHOSEN.load(Ordering::Relaxed), + shattered: SHATTERED.load(Ordering::Relaxed), + } +} + +/// The aggregate `SCHED[agg] …` line, or `None` when nothing was ever +/// scheduled — the honest answer to "how wide were the steps" when none ran is +/// *nothing*, not `0`. +#[must_use] +pub fn sched_bucket_marker() -> Option { + let snap = sched_buckets(); + (snap.calls > 0).then(|| snap.marker("agg")) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The ratios are what get quoted, so they are what must be pinned. A batch + /// that shatters is the case the counters exist to make visible. + #[test] + fn a_shattered_batch_reports_the_width_that_actually_ran() { + // Three steps: 128 offered, split 40 ways, 3 of them ran. + let mut snap = SchedBuckets::default(); + for _ in 0..3 { + snap.calls += 1; + snap.buckets += 40; + snap.offered += 128; + snap.chosen += 3; + snap.shattered += 1; + } + assert_eq!(snap.buckets_per_step(), Some(40.0)); + assert_eq!(snap.offered_per_step(), Some(128.0)); + // 🔑 The number a throughput result has to be read against: the engine + // ran 3-wide steps while the harness called it B=128. + assert_eq!(snap.running_bucket_size(), Some(3.0)); + assert_eq!(snap.preempted_per_step(), Some(125.0)); + assert_eq!(snap.shattered_frac(), Some(1.0)); + let m = snap.marker("agg"); + assert!(m.contains("running_bucket_size=3.0000"), "{m}"); + assert!(m.contains("buckets_per_step=40.0000"), "{m}"); + } + + /// An unshattered batch must be visibly different, or the counter cannot + /// distinguish "the scheduler serialised us" from "it did not". + #[test] + fn an_intact_batch_reports_one_bucket_and_no_preemption() { + let snap = SchedBuckets { + calls: 10, + buckets: 10, + offered: 1280, + chosen: 1280, + shattered: 0, + }; + assert_eq!(snap.buckets_per_step(), Some(1.0)); + assert_eq!(snap.running_bucket_size(), Some(128.0)); + assert_eq!(snap.preempted_per_step(), Some(0.0)); + assert_eq!(snap.shattered_frac(), Some(0.0)); + } + + /// Differencing across a fence is how a cell's numbers are taken; a reset + /// in between must not read as a huge result. + #[test] + fn differencing_saturates_rather_than_wrapping() { + let early = SchedBuckets { + calls: 100, + buckets: 200, + offered: 1000, + chosen: 500, + shattered: 50, + }; + let late = SchedBuckets { + calls: 150, + buckets: 260, + offered: 1600, + chosen: 900, + shattered: 60, + }; + let d = late.since(&early); + assert_eq!(d.calls, 50); + assert_eq!(d.buckets, 60); + assert_eq!(d.running_bucket_size(), Some(8.0)); + // A counter that went backwards (process restart) reads as zero work, + // never as a wrapped `usize`. + assert_eq!(early.since(&late), SchedBuckets::default()); + } + + /// Nothing scheduled means no marker, not a marker full of zeros. + #[test] + fn no_scheduling_reports_nothing_rather_than_zero() { + assert!(SchedBuckets::default().buckets_per_step().is_none()); + assert!(SchedBuckets::default().running_bucket_size().is_none()); + } +} diff --git a/mistralrs-core/src/scheduler/default_scheduler.rs b/mistralrs-core/src/scheduler/default_scheduler.rs index 16d3466da..6769fed20 100644 --- a/mistralrs-core/src/scheduler/default_scheduler.rs +++ b/mistralrs-core/src/scheduler/default_scheduler.rs @@ -68,9 +68,22 @@ pub trait BucketingManager: Send + Sync { ) -> BucketedSeqs; } -// (cache length, (has_imgs && is_prompt), sequence offset) -// Bucket by that metric for images because if we are not a prompt, then this doesn't apply -type BucketKey = (usize, bool, usize); +// (cache length, (has_imgs && is_prompt), sequence offset, is_prompt) +// Bucket by the image flag because if we are not a prompt, then it doesn't +// apply. +// +// 🔑 `is_prompt` is a key component OF ITS OWN, not folded into the image +// flag, because the ragged-decode sentinel collides without it: under +// `ragged_decode_supported()` every decode sequence keys length 0, and a +// 1-token prompt's `cache_bucket_len()` is `len() - 1 == 0` too — so a fresh +// prompt silently merged into the pinned decode bucket, rode the +// `seq_buckets.len() <= 1` fast path into the same scheduling pass, and +// bypassed `select_running_bucket`, the floor's bookkeeping and the +// one-bucket-per-step invariant. With prompt-ness in the key a bucket is +// homogeneous by construction, which is also what lets the floor read +// prompt-ness off the key instead of off whichever sequence happens to sit +// first in the bucket. +type BucketKey = (usize, bool, usize, bool); /// How many waiting sequences may be admitted to **prefill** in one engine /// iteration. `None` (the default, and the historical behaviour) means "all of @@ -164,7 +177,7 @@ fn select_running_bucket( ) -> BucketKey { let min = *seq_buckets .keys() - .min_by_key(|(x, _, _)| *x) + .min_by_key(|(x, _, _, _)| *x) .expect("No sequence buckets."); if discrete { return min; @@ -180,11 +193,15 @@ fn select_running_bucket( } // A merge only happens when the *whole* key matches, so the coalescing - // target must agree on the image flag and token offset too. + // target must agree on the image flag, token offset and prompt-ness too — + // a decode bucket must never idle itself to "catch up" to a prompt bucket, + // which runs by a different mechanism entirely. let Some(next_len) = seq_buckets .keys() - .filter(|(len, imgs, offset)| *len > min.0 && *imgs == min.1 && *offset == min.2) - .map(|(len, _, _)| *len) + .filter(|(len, imgs, offset, prompt)| { + *len > min.0 && *imgs == min.1 && *offset == min.2 && *prompt == min.3 + }) + .map(|(len, _, _, _)| *len) .min() else { return greedy; @@ -201,9 +218,18 @@ fn select_running_bucket( } } +/// The default prompt-starvation floor, in decode steps: a prompt bucket is +/// forced to run after this many consecutive pass-overs. +/// +/// `4` was picked against this file's own fixture (see +/// `the_floor_gives_prompts_a_turn_without_giving_them_every_turn`): behind a +/// 47-sequence decode cohort a fresh prompt wins 0 of 24 steps with no floor, +/// and 1 in every 5 with floor=4 — a turn without taking every turn. +const DEFAULT_PREFILL_FLOOR_STEPS: usize = 4; + /// How many consecutive iterations prompt buckets may be passed over before one -/// is forced to run. `None` (the default) is the historical behaviour: prompts -/// win a step only by out-scoring every decode bucket on priority. +/// is forced to run. `None` means the historical behaviour: prompts win a step +/// only by out-scoring every decode bucket on priority. /// /// # Why a floor is needed as well as a cap /// @@ -223,16 +249,25 @@ fn select_running_bucket( /// waiting` and **stops**, with 48 sequences admitted and never prefilled. /// Capping admission cannot lift it; only guaranteeing prompts a turn can. /// -/// Read once from `ARC_PREFILL_FLOOR_STEPS`; unset or `0` reproduces the -/// previous selection exactly. +/// Read once from `ARC_PREFILL_FLOOR_STEPS`, **by value**: +/// * unset (or unparsable) → [`DEFAULT_PREFILL_FLOOR_STEPS`] — the floor is +/// ON by default, because the starvation it lifts is the default behaviour; +/// * `0` → `None`, the kill-switch: reproduces the pre-floor selection +/// exactly, key for key; +/// * any other `n` → `Some(n)`. fn prefill_starvation_floor() -> Option { static FLOOR: std::sync::OnceLock> = std::sync::OnceLock::new(); - *FLOOR.get_or_init(|| { - std::env::var("ARC_PREFILL_FLOOR_STEPS") - .ok() - .and_then(|v| v.parse::().ok()) - .filter(|n| *n > 0) - }) + *FLOOR.get_or_init(|| floor_from(std::env::var("ARC_PREFILL_FLOOR_STEPS").ok().as_deref())) +} + +/// The decision itself, separated from the process-global `OnceLock` and the +/// environment so both sides of the default can be pinned by a test. +fn floor_from(raw: Option<&str>) -> Option { + match raw.map(str::parse::) { + Some(Ok(0)) => None, + Some(Ok(n)) => Some(n), + Some(Err(_)) | None => Some(DEFAULT_PREFILL_FLOOR_STEPS), + } } #[derive(Default)] @@ -268,6 +303,13 @@ impl BucketingManager for FixedBucketingManager { discrete: bool, ) -> BucketedSeqs { // Now, get the sequences with the smallest sequence lengths, and allow them to catch up. + // + // How wide the step the engine is about to run actually is. A batch + // that shatters into N buckets runs one of them and preempts the rest, + // so "B=128" in a harness can be 3-wide in the engine — and a + // throughput number is unreadable without knowing which. See + // [`super::bucket_telemetry`]. + let offered = running.len(); let mut seq_buckets: HashMap> = HashMap::new(); let mut seq_priorities: HashMap = HashMap::new(); let ragged_decode = crate::kv_cache::ragged_decode_supported(); @@ -289,42 +331,26 @@ impl BucketingManager for FixedBucketingManager { } else { seq.cache_bucket_len() }; - match seq_buckets.get_mut(&( + // `is_prompt` closes the sentinel collision documented on + // `BucketKey`: a 1-token prompt also computes `len == 0` here. + let key: BucketKey = ( len, seq.images().is_some() && seq.is_prompt(), seq.token_offset(), - )) { + seq.is_prompt(), + ); + match seq_buckets.get_mut(&key) { Some(bucket) => { if !discrete { - *seq_priorities - .get_mut(&( - len, - seq.images().is_some() && seq.is_prompt(), - seq.token_offset(), - )) - .unwrap() += seq.compute_priority(); + *seq_priorities.get_mut(&key).unwrap() += seq.compute_priority(); } bucket.push(seq); } None => { if !discrete { - seq_priorities.insert( - ( - len, - seq.images().is_some() && seq.is_prompt(), - seq.token_offset(), - ), - seq.compute_priority(), - ); + seq_priorities.insert(key, seq.compute_priority()); } - seq_buckets.insert( - ( - len, - seq.images().is_some() && seq.is_prompt(), - seq.token_offset(), - ), - vec![seq], - ); + seq_buckets.insert(key, vec![seq]); } } } @@ -355,6 +381,7 @@ impl BucketingManager for FixedBucketingManager { ); } } + let n_buckets = seq_buckets.len(); let running = if seq_buckets.len() <= 1 { // Full steam ahead or have everything seq_buckets @@ -372,12 +399,10 @@ impl BucketingManager for FixedBucketingManager { // floor unset — or before it expires — the selection is the // pre-change one, key for key. if let Some(floor) = self.floor.filter(|_| !discrete) { - let is_prompt_bucket = |k: &BucketKey| { - seq_buckets - .get(k) - .and_then(|v| v.first()) - .is_some_and(Sequence::is_prompt) - }; + // Prompt-ness is a key component, so a bucket is homogeneous + // by construction and the key answers directly — no reading it + // off whichever sequence happens to sit first in the bucket. + let is_prompt_bucket = |k: &BucketKey| k.3; if is_prompt_bucket(&len) { self.steps_since_prompt = 0; } else if self.steps_since_prompt >= floor { @@ -387,7 +412,7 @@ impl BucketingManager for FixedBucketingManager { if let Some(k) = seq_buckets .keys() .filter(|k| is_prompt_bucket(k)) - .min_by_key(|(l, _, _)| *l) + .min_by_key(|(l, _, _, _)| *l) .copied() { len = k; @@ -430,6 +455,7 @@ impl BucketingManager for FixedBucketingManager { // Know min_seqs.len < running.len() <= max highest_priority_seqs }; + super::bucket_telemetry::record_bucketing(n_buckets, offered, running.len()); BucketedSeqs { running, waiting } } } @@ -769,7 +795,6 @@ mod tests { seq } - /// A sequence the scheduler has admitted but not yet prefilled. fn prompt_of_len(id: usize, n_toks: usize) -> Sequence { let seq = seq_of_len(id, n_toks); @@ -835,13 +860,37 @@ mod tests { ); } - /// The knob off is the previous selection, key for key — asserted against - /// the same fixture rather than argued. + /// The kill-switch (`ARC_PREFILL_FLOOR_STEPS=0` → `floor: None`) is the + /// previous selection, key for key — asserted against the same fixture + /// rather than argued. Note "unset" no longer means this: unset is the + /// default floor of 4 (see `floor_from`); `0` is how today's behaviour is + /// restored. #[test] - fn the_floor_unset_reproduces_the_previous_selection() { + fn the_floor_disabled_reproduces_the_previous_selection() { assert_eq!(prompt_steps_won(None, 12), 0); } + /// The env mapping behind the default, both sides pinned: + /// unset → the default floor (ON), `0` → the kill-switch (OFF, historical + /// selection), a number → itself, garbage → the default rather than a + /// silent disable. + /// + /// Mutation check (run 2026-08-21): revert `floor_from`'s `None` arm to the + /// pre-change `.filter(|n| *n > 0)` shape (unset → `None`) and this fails + /// on the first assertion with `left: None, right: Some(4)`. + #[test] + fn the_floor_defaults_on_and_zero_is_the_kill_switch() { + assert_eq!(floor_from(None), Some(DEFAULT_PREFILL_FLOOR_STEPS)); + assert_eq!(floor_from(Some("0")), None, "0 must cleanly disable"); + assert_eq!(floor_from(Some("7")), Some(7)); + assert_eq!( + floor_from(Some("not-a-number")), + Some(DEFAULT_PREFILL_FLOOR_STEPS), + "garbage must fall back to the default, not silently disable the floor" + ); + assert_eq!(DEFAULT_PREFILL_FLOOR_STEPS, 4); + } + /// Advance a scheduled sequence by one decoded token, as the engine would. fn decode_one(seq: &mut Sequence) { seq.add_token( @@ -858,7 +907,7 @@ mod tests { fn bucket_of(len: usize, n: usize, id_base: usize) -> (BucketKey, Vec) { let seqs = (0..n).map(|i| seq_of_len(id_base + i, len)).collect(); - ((len, false, 0), seqs) + ((len, false, 0, false), seqs) } /// 🔑 The dense equivalent of the paged regression: eight sequences at @@ -913,6 +962,136 @@ mod tests { ); } + /// 🔑 The cohort-can-only-shrink defect, and its fix, in one fixture: a + /// formed decode cohort at length 500 plus ONE newcomer at length 100. + /// + /// With exact-length bucketing the newcomer can never join: coalescing is + /// refused (`(9-1)*400 = 3200 > 1*256`), greedy priority keeps the cohort + /// (`8*log2(500) >> 1*log2(100)`), and running the cohort *grows* the gap + /// by one every step — so the newcomer idles forever and cohorts only + /// shrink. Every published B=256 number under this regime came from a + /// synthetic uniform burst, because a real arrival process can never refill + /// a cohort. + /// + /// With ragged decode on, every decode sequence keys to one bucket and the + /// newcomer is admitted into the running set the same step. + /// + /// Mutation check (run 2026-08-21): pin `len` to `seq.cache_bucket_len()` + /// (reverting the `ragged_decode` arm at the top of + /// `bucket_and_waitlist_seqs_waiting`) and the ON arm fails with + /// `on: 8/9` — the newcomer is waitlisted again. The OFF arm is asserted + /// in the same run, so a fixture that stopped being ragged also fails. + #[test] + fn a_formed_cohort_admits_a_newcomer_when_ragged_decode_is_on() { + use crate::kv_cache::ragged_decode_test_override as flag; + let logger = IntervalLogger::new(Duration::from_secs(3600), None); + + let running_ids_with = |on: bool| -> Vec { + flag::with(on, || { + let mut sched: DefaultScheduler> = DefaultScheduler::new( + DefaultSchedulerMethod::Fixed(NonZeroUsize::new(128).unwrap()), + ); + // The formed cohort: eight decoding sequences at length 500. + for i in 0..8 { + sched.add_seq(seq_of_len(i, 500)); + } + // The newcomer: decoding at length 100 (freshly prefilled and + // far behind the cohort — the shape coalescing refuses). + sched.add_seq(seq_of_len(8, 100)); + // A waiting request, so selection runs the non-discrete + // arrival path (the discrete no-waiting path is min-length + // catch-up, a different regime with its own test). + let waiter = seq_of_len(9, 300); + waiter.set_state(SequenceState::Waiting); + sched.add_seq(waiter); + let out = DefaultScheduler::schedule(&mut sched, &logger); + out.completion.into_vec().iter().map(|s| *s.id()).collect() + }) + }; + + let off = running_ids_with(false); + assert!( + !off.contains(&8) && off.len() == 8, + "fixture guard: with exact-length bucketing the newcomer must be \ + waitlisted behind the cohort, got {off:?} — otherwise the ON arm \ + proves nothing" + ); + + let on = running_ids_with(true); + assert_eq!( + on.len(), + 9, + "ragged decode must run the cohort AND the newcomer together, got {on:?}" + ); + assert!( + on.contains(&8), + "the newcomer must be in the running set, got {on:?}" + ); + } + + /// 🔑 The pinned-bucket collision probe: with ragged decode published, a + /// decode sequence keys its bucket at the sentinel length 0 — and a + /// 1-token PROMPT's `cache_bucket_len()` is `len() - 1 == 0`, the same + /// number. Without prompt-ness in the bucket key the two share + /// `(0, false, 0)` and the prompt silently rides into the decode cohort: + /// scheduled in the same pass through the merged single bucket's + /// `seq_buckets.len() <= 1` fast path, which bypasses `select_running_bucket`, + /// the starvation floor's bookkeeping, and the one-bucket-per-step + /// invariant the `ArcSched: decode bucketing` log line reports on. + /// + /// Asserts the prompt is NOT folded in: the decode cohort runs whole and + /// alone, and the prompt waits its turn through ordinary selection (or the + /// floor). + /// + /// Mutation check (run 2026-08-21): drop `seq.is_prompt()` from the bucket + /// key (pin the fourth component to `false`) and this fails with + /// `prompt = [4]` scheduled alongside `completion = [0, 1, 2, 3]`. + #[test] + fn a_one_token_prompt_is_not_folded_into_the_pinned_decode_bucket() { + use crate::kv_cache::ragged_decode_test_override as flag; + let logger = IntervalLogger::new(Duration::from_secs(3600), None); + + flag::with(true, || { + let mut sched: DefaultScheduler> = DefaultScheduler::new( + DefaultSchedulerMethod::Fixed(NonZeroUsize::new(128).unwrap()), + ); + // A ragged decode cohort — under the flag they all key length 0. + for (i, l) in [60usize, 71, 82, 93].iter().enumerate() { + sched.add_seq(seq_of_len(i, *l)); + } + // A 1-token prompt: `cache_bucket_len()` is `1 - 1 == 0`, the + // exact value of the ragged sentinel. + let p = seq_of_len(4, 1); + p.set_state(SequenceState::Waiting); + sched.add_seq(p); + + let (completion, prompt) = { + let out = DefaultScheduler::schedule(&mut sched, &logger); + let mut completion: Vec = + out.completion.iter().map(|s| *s.id()).collect(); + completion.sort_unstable(); + let prompt: Vec = out.prompt.iter().map(|s| *s.id()).collect(); + (completion, prompt) + }; + + assert_eq!( + completion, + vec![0, 1, 2, 3], + "the decode cohort must run whole and alone" + ); + assert!( + prompt.is_empty(), + "a 1-token prompt must not share the pinned decode bucket; it \ + must wait for selection (or the floor), got prompt = {prompt:?}" + ); + assert_eq!( + Scheduler::waiting_len(&sched), + 1, + "the prompt must be waitlisted, not lost" + ); + }); + } + /// A PROMPT batch keeps the exact-length key even with the flag on. Prefill /// is right-padded and its `cu_seqlens` are built from those padded widths /// — a different mechanism from the decode one, and not covered by the diff --git a/mistralrs-core/src/scheduler/mod.rs b/mistralrs-core/src/scheduler/mod.rs index 34c8d56fb..f3606ac3b 100644 --- a/mistralrs-core/src/scheduler/mod.rs +++ b/mistralrs-core/src/scheduler/mod.rs @@ -1,7 +1,9 @@ +pub mod bucket_telemetry; mod default_scheduler; use std::sync::Arc; +pub use bucket_telemetry::{sched_bucket_marker, sched_buckets, SchedBuckets}; pub use default_scheduler::{DefaultScheduler, DefaultSchedulerMethod, DefaultSchedulerOutput}; use tokio::sync::Mutex; 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; } diff --git a/mistralrs-core/tests/capability_reachability.rs b/mistralrs-core/tests/capability_reachability.rs index bffe418f7..2ba4de18f 100644 --- a/mistralrs-core/tests/capability_reachability.rs +++ b/mistralrs-core/tests/capability_reachability.rs @@ -119,6 +119,39 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + // ── SHIPPED OFF, WITH A NAMED EXPERIMENT ──────────────────────────────── + // + // The entry below is `Live` because the GATE is reached on every step — + // that is what this file checks, and it must not go dark. What is switched + // off is the gate's DEFAULT ANSWER, which no reachability check can see. + // + // Recording it here anyway, because "built, correct, tested, and switched + // off" is this repo's most expensive failure mode and this file is where + // someone goes looking for it. 26 pieces of finished work are parked that + // way right now, including a complete GPU sampler with zero callers; we + // shipped 34 tok/s of landed work while master served 17, for exactly this + // reason. A flag that ships off with no named experiment behind it is how + // that happens, so the experiment is named: + // + // flag ARC_V4_XS_PIN_WINDOW=1 + // gate kv_cache/xs_rolling.rs :: xs_pin_window_enabled_from + // test arc-tools/arcspec_perseq_ladder.sh — ON vs ON_PINNED, uniform + // B=32, one binary, same prompt and seed + // pass ON_PINNED faster on aggregate tok/s, generated tokens identical + // was +23.2% at uniform B=32 on #121's own pre-rebase branch — NOT a + // current fact, do not quote it as one + // then the default flips to ON in the same change that records the + // number. Leaving it off after a passing measurement is a failure + // state, not the finish line. + Capability { + name: "xs window pin: built and tested, default OFF pending one measurement", + parent: "ArcInfer / ArcKV", + check: Check::Symbol { + symbol: "xs_pin_window_enabled", + defined_in: "mistralrs-core/src/kv_cache/xs_rolling.rs", + }, + status: Status::Live, + }, Capability { name: "ragged decode: per-row lengths on the xs cache", parent: "ArcInfer / ArcKV", @@ -137,6 +170,116 @@ static REGISTRY: &[Capability] = &[ }, status: Status::Live, }, + // ── SHIPPED OFF, WITH A NAMED EXPERIMENT ──────────────────────────────── + // + // This entry is `Live` because the GATE is reached on every layer of every + // step — that is what this file checks, and it must not go dark. What is + // switched off is the gate's DEFAULT ANSWER at `batch > 1`, which no + // reachability check can see. + // + // Recorded here anyway, because "built, correct, tested, and switched off" + // is this repo's most expensive failure mode and this file is where someone + // goes looking for it. 26 pieces of finished work are parked that way right + // now, including a complete GPU sampler with zero callers; we shipped + // 34 tok/s of landed work while master served 17, for exactly this reason. + // + // The kernel is batch-correct by construction and the dispatch bug that + // kept it dead above batch 1 (`seqlen_offsets.len() != 1` — the LENGTH of + // the vector where the property needed is the DISTINCTNESS of its values) + // is fixed. It is off only because fixing that gate makes a CUDA kernel + // live on a path it has never run on, and this repo has paid for that twice + // in one week: TCFRAG (default-ON, "UNVERIFIED ON HARDWARE — NEVER RUN" in + // its own header, 63 GB held, a layer permanently broken, retired by #209) + // and the fused-512 attention path (four days silently dropping the + // attention mask at 12% agreement). + // + // flag ARC_QK_FUSED_COHORT (default ON since s10; "0" is the kill switch) + // gate cuda/qk_norm_rope.rs :: fused_cohort_enabled_from + // MEASURED s10 leg 4 (box arc-s10-ledger H200, candidate 4d03b9e25): + // ENGAGED at B=256, seeded canaries byte-identical to the eager + // arm, 599.35 vs 575.12 tok/s aggregate (the +4.2% sits inside the + // ~±5% inter-leg variance measured the same session — the default + // stands on engaged + bit-identical + not-slower). + // then DONE: the default flipped to ON in the same session, as this + // entry always demanded. Leaving it off after a passing + // measurement is a failure state, not the finish line. + Capability { + name: "fused qk_norm_rope at B>1: default ON, measured s10", + parent: "ArcInfer / ArcAttention", + check: Check::Symbol { + symbol: "fused_cohort_enabled", + defined_in: "mistralrs-core/src/cuda/qk_norm_rope.rs", + }, + status: Status::Live, + }, + // Same shape as the entry above, same doctrine, one rung lower in the + // stack: the V4 rotary's RAGGED arm. Uniform cohorts already take one + // batched `rope_i` unconditionally (that dispatch is exact and shipped + // ON); this flag governs rows at genuinely different positions, which the + // per-sequence loop served at ~9 ops × 43 layers × B per step. The + // gathered rank-3 path is bit-identical to the loop + // (`deepseek_ragged_cohort_tests`) and default OFF until measured. + // + // flag ARC_ROPE_COHORT=1 (read by value via env_flag_is_set) + // gate layers.rs :: rope_cohort_enabled + // test one binary, ragged B>1 decode (ARC_MTP_PER_SEQ_KV on, or any + // producer of unequal offsets), flag set vs unset + // pass outputs bit-identical AND the gathered arm faster; + // rope_cohort_stats::counts() shows cohort > 0 on the flagged arm + // then the default flips in the same change that records the number + Capability { + name: "ragged-cohort RoPE: one gathered rope_i for unequal offsets, default OFF", + parent: "ArcInfer / ArcAttention", + check: Check::Symbol { + symbol: "rope_cohort_enabled", + defined_in: "mistralrs-core/src/layers.rs", + }, + status: Status::Live, + }, + Capability { + // The mask the entry above forces into existence, built on the DEVICE: + // the host triple-loop build (kept as + // `make_left_padded_causal_mask_host`, the test oracle) was 4.2 MB of + // host work + H2D per token at B=256 ctx-4096, i.e. the thing that + // would make ragged decode read as a regression. If production stops + // reaching this builder, ragged batches are serving unmasked dead + // prefixes — silently wrong, not slow. + name: "ragged decode: the left-padded mask is built device-side", + parent: "ArcInfer / ArcAttention", + check: Check::Symbol { + symbol: "make_left_padded_causal_mask", + defined_in: "mistralrs-core/src/layers_masker.rs", + }, + status: Status::Live, + }, + Capability { + // The channel the SCHEDULER reads to stop exact-length bucketing. + // `batch_can_be_ragged` (above) is the producer; this is the consumer + // side, and either going dark re-imposes the one-bucket-per-step + // ceiling with no error anywhere. + name: "ragged decode: the scheduler reads the published capability", + parent: "ArcInfer / ArcSched", + check: Check::Symbol { + symbol: "ragged_decode_supported", + defined_in: "mistralrs-core/src/kv_cache/mod.rs", + }, + status: Status::Live, + }, + Capability { + // The supported ON-switch for the V4 ragged pair + // (`--v4-ragged-decode` / `v4_ragged_decode` config key, with + // `ARC_V4_XS_PER_SEQ` as the env fallback). A capability reachable + // only through an undocumented env var is indistinguishable from one + // that does not exist — which is how the ragged-batching result stayed + // off the shipped configuration once already. + name: "ragged decode: the CLI/config surface latches the V4 xs flag", + parent: "ArcInfer / ArcKV", + check: Check::Symbol { + symbol: "request_xs_per_sequence", + defined_in: "mistralrs-core/src/kv_cache/xs_rolling.rs", + }, + status: Status::Live, + }, Capability { name: "prefill admission cap (--prefill-max-seqs)", parent: "ArcInfer / ArcSched", @@ -147,7 +290,11 @@ static REGISTRY: &[Capability] = &[ status: Status::Live, }, Capability { - name: "prompt-starvation floor (--prefill-floor-steps)", + // Default ON since arcsched/ragged-and-floor: unset resolves to a + // floor of 4 (`DEFAULT_PREFILL_FLOOR_STEPS`); + // `ARC_PREFILL_FLOOR_STEPS=0` is the kill-switch that restores the + // pre-floor selection key for key. + name: "prompt-starvation floor (ARC_PREFILL_FLOOR_STEPS, default 4)", parent: "ArcInfer / ArcSched", check: Check::Symbol { symbol: "prefill_starvation_floor", @@ -164,6 +311,50 @@ 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, + }, + Capability { + name: "FP8 dispatch engagement log (ARC_LOG_FP8_DISPATCH)", + parent: "ArcLab", + check: Check::Symbol { + symbol: "log_fp8_dispatch", + defined_in: "mistralrs-quant/src/blockwise_fp8/ops.rs", + }, + // The engagement line the FP8 threshold sweep + // (`arc-tools/fp8_threshold_sweep.sh`) greps per leg to prove WHICH + // kernel actually served it. Value-read (`env_flag_is_set`), latched + // once per process, one line per path per process. Referenced from + // both native dispatch arms in `ops.rs` and from the + // dequantize+cuBLASLt divert in `blockwise_fp8/mod.rs`. + status: Status::Live, + }, + // Flipped default-ON in session 10 after its hardware A/B (box + // `arc-s10-ledger`, candidate `4d03b9e25`, leg 3): b=1 decode 43.80 vs + // 40.76 tok/s (+7.5%, 3x512 produced tokens, seeded), engagement proven by + // `[arc-fp8-dispatch] path=gemv_wide`, coherence canary correct. The + // ledger's ~1.6x prediction did NOT materialise end-to-end; +7.5% is the + // measured number and the one this default stands on. `ARC_FP8_GEMV_WIDE=0` + // (exact string) is the kill switch; value-read (#212 doctrine). + Capability { + name: "FP8 wide b=1 GEMV (ARC_FP8_GEMV_WIDE): default ON, measured s10", + parent: "ArcKernels", + check: Check::Symbol { + symbol: "wide_enabled", + defined_in: "mistralrs-quant/src/blockwise_fp8/gemv_wide.rs", + }, + status: Status::Live, + }, // -- the ApiPromise class ------------------------------------------------ Capability { name: "sampling: logit_bias reaches the sampler", @@ -196,17 +387,20 @@ static REGISTRY: &[Capability] = &[ accepted_in: &["mistralrs-core/src/sampler.rs"], honoured_in: &["arc-cuda-graph/src/autonomous.rs"], }, - // `AutonomousDecodeConfig` (autonomous.rs:70-82) carries temperature, - // top_p, both penalties and `greedy` — and neither top_k nor min_p. A - // request setting them would have them silently dropped the moment this - // path is reached. It is not reached today (`mark_unreachable( - // "cuda_graph.autonomous_decode", ...)`), which is the only reason this - // is not already a user-visible defect. - status: Status::Tracked { - reason: "AutonomousDecodeConfig has no top_k field at all; the path is itself \ - unreachable today, so this is latent rather than live. Fix before \ - GPU-autonomous decode is switched on.", - }, + // Was Tracked with the reason "AutonomousDecodeConfig has no top_k + // field at all". #130 gave it one: `AutonomousDecodeConfig::top_k` + // (autonomous.rs:82), plumbed into `sampling_cpu::SamplingConfig` at + // autonomous.rs:643 and consumed by `self.sampler.sample(..)`. So the + // promise is now honoured where it is made, and the status follows the + // code. + // + // The gate flagged this itself, on the rebase, with "PROMOTE ME" — + // the same way the greedy/logit-bias entry above was promoted. That is + // the mechanism working, not a test to be silenced. + // + // `min_p` is a separate entry below and stays Tracked: it still has no + // field on `AutonomousDecodeConfig`. + status: Status::Live, }, Capability { name: "sampling: min_p survives onto the GPU-autonomous decode path", @@ -221,6 +415,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 diff --git a/mistralrs-core/tests/doc_citations.rs b/mistralrs-core/tests/doc_citations.rs index e4204e8cc..6d219d675 100644 --- a/mistralrs-core/tests/doc_citations.rs +++ b/mistralrs-core/tests/doc_citations.rs @@ -351,6 +351,30 @@ static BASELINE: &[Waiver] = &[ kind: Kind::Unresolved, why: "candle is an out-of-tree path dependency", }, + // The CUDA driver header. #130's FFI declarations cite it by line for the + // struct layouts and flag values they mirror (`CUDA_CONDITIONAL_NODE_PARAMS`, + // `CU_GRAPH_COND_ASSIGN_DEFAULT`, the conditional-node entry points). It + // ships with the CUDA Toolkit and is not vendored here, so no citation into + // it can ever resolve — and its line numbers are toolkit-version-specific + // besides (these are CUDA 13.1), which is why each citation names the + // version in its surrounding comment. + // + // Waived by bare filename rather than one row per line number: the rows + // would otherwise have to be rewritten on every toolkit bump, which is + // churn the gate cannot check either way. The layouts themselves ARE + // checked, against `cudarc`'s bindgen output — see the ABI note on #130. + Waiver { + doc: "arc-cuda-graph/src/ffi.rs", + cite: "cuda.h", + kind: Kind::Unresolved, + why: "CUDA Toolkit header, not vendored here", + }, + Waiver { + doc: "arc-cuda-graph/src/autonomous.rs", + cite: "cuda.h", + kind: Kind::Unresolved, + why: "CUDA Toolkit header, not vendored here", + }, Waiver { doc: "memory/mission/wave29-BD-rung-decision.md", cite: "STATUS.md", diff --git a/mistralrs-paged-attn/build.rs b/mistralrs-paged-attn/build.rs index 764785806..b495dc003 100644 --- a/mistralrs-paged-attn/build.rs +++ b/mistralrs-paged-attn/build.rs @@ -13,6 +13,11 @@ fn main() -> Result<()> { println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed=src/cuda/pagedattention.cuh"); println!("cargo:rerun-if-changed=src/cuda/turbo_paged_attention.cuh"); + // Every other .cu here is listed; this one was not, and because emitting + // *any* rerun-if-changed disables cargo's watch-the-whole-package default, + // edits to it were relying entirely on cudaforge's object cache to trigger + // a rebuild. + println!("cargo:rerun-if-changed=src/cuda/turbo_paged_attention.cu"); println!("cargo:rerun-if-changed=src/cuda/copy_blocks_kernel.cu"); println!("cargo:rerun-if-changed=src/cuda/reshape_and_cache_kernel.cu"); println!("cargo:rerun-if-changed=src/cuda/concat_and_cache_mla_kernel.cu"); @@ -52,7 +57,49 @@ fn main() -> Result<()> { .arg("--compiler-options") .arg("-fPIC"); + // D16: ship cubins for every arch we target, not just the build box's GPU. + // + // cudaforge derives a single `-gencode` from `CUDA_COMPUTE_CAP` or the local + // `nvidia-smi`, so a binary built on an H200 carries SM90 only and PTX-JITs + // (or fails) on Blackwell. `ARC_CUDA_ARCHS` appends further targets to the + // fat binary — e.g. `ARC_CUDA_ARCHS=90,100,103` for Hopper + Blackwell. + // + // Left unset the build is unchanged, because each extra arch recompiles + // every .cu in this crate and that cost is not worth paying on a rented box + // that only ever runs one arch. Release and publish builds set it. + println!("cargo:rerun-if-env-changed=ARC_CUDA_ARCHS"); + + // Hoisted above the `ARC_CUDA_ARCHS` loop because it decides which of the + // requested arches are actually *extra*. let compute_cap = builder.get_compute_cap().unwrap_or(80); + + if let Ok(archs) = std::env::var("ARC_CUDA_ARCHS") { + for arch in archs.split(',').map(str::trim).filter(|a| !a.is_empty()) { + // cudaforge already emits one `-gencode` of its own, derived from + // `CUDA_COMPUTE_CAP` or the local `nvidia-smi`, and suffixes it `a` + // for cap >= 90 exactly as we do below (`GpuArch::auto_suffix` -> + // `to_gencode_arg`). It then appends `extra_args` verbatim, so + // re-requesting that arch here hands nvcc the byte-identical flag + // twice: on an H200 building `ARC_CUDA_ARCHS=90,100,103` the + // duplicate is `-gencode=arch=compute_90a,code=sm_90a`. + // + // Emit only what cudaforge does not already cover. + if arch.parse::().is_ok_and(|c| c == compute_cap) { + continue; + } + // Arch-specific ("a") cubins: sm_90a/sm_100a expose the Hopper and + // Blackwell-only instructions the arch-specialised kernel paths are + // written against. Matches cudaforge's own suffixing for cap >= 90. + let suffix = if arch.parse::().is_ok_and(|c| c >= 90) { + "a" + } else { + "" + }; + builder = builder.arg(&format!( + "-gencode=arch=compute_{arch}{suffix},code=sm_{arch}{suffix}" + )); + } + } // Enable FP8 if compute capability >= 8.0 (Ampere and newer) let using_fp8 = if compute_cap >= 80 { builder = builder.arg("-DENABLE_FP8"); diff --git a/mistralrs-paged-attn/src/cuda/backend/paged_attention.rs b/mistralrs-paged-attn/src/cuda/backend/paged_attention.rs index 14865d2f8..f147d2be3 100644 --- a/mistralrs-paged-attn/src/cuda/backend/paged_attention.rs +++ b/mistralrs-paged-attn/src/cuda/backend/paged_attention.rs @@ -712,10 +712,11 @@ pub fn turbo_reshape_and_cache( let dev = k.device(); let (num_tokens, num_heads, head_size) = k_l.shape().dims3()?; - if head_size != 128 { + if !turbo_supports_head_dim(head_size) { candle::bail!( - "turbo_reshape_and_cache only supports head_size=128 (got {head_size}); \ - the TurboQuant kernel would exit without writing anything" + "turbo_reshape_and_cache has no kernel instantiation for head_size={head_size} \ + (have {TURBO_HEAD_DIMS:?}); the launch would fall through the dispatch switch \ + and write nothing" ); } @@ -770,6 +771,21 @@ pub fn turbo_reshape_and_cache( Ok(()) } +/// Head dimensions `turbo_paged_attention.cu` instantiates its kernels for. +/// +/// Must stay equal to `TURBOQUANT_CUDA_HEAD_DIMS` in +/// `mistralrs-quant::turboquant::cuda_tables`, which is where the test lives +/// that parses this declaration and the kernel's dispatch switch and fails if +/// the three disagree. A width missing here is merely refused; a width present +/// here but missing from the switch would launch nothing and leave the caller +/// reading uninitialized memory. +pub const TURBO_HEAD_DIMS: [usize; 4] = [64, 128, 256, 512]; + +/// Whether the TurboQuant kernels have an instantiation for `head_dim`. +pub fn turbo_supports_head_dim(head_dim: usize) -> bool { + TURBO_HEAD_DIMS.contains(&head_dim) +} + /// TurboQuant: paged attention over compressed U8 KV cache. #[allow(clippy::too_many_arguments)] #[allow(unused_variables)] @@ -827,10 +843,11 @@ pub fn turbo_paged_attention( let dev = q_s.device(); let (num_seqs, num_heads, head_size) = q_l.shape().dims3()?; - if head_size != 128 { + if !turbo_supports_head_dim(head_size) { candle::bail!( - "turbo_paged_attention only supports head_size=128 (got {head_size}); \ - the kernel would exit early and return an uninitialized output buffer" + "turbo_paged_attention has no kernel instantiation for head_size={head_size} \ + (have {TURBO_HEAD_DIMS:?}); the launch would fall through the dispatch switch \ + and return an uninitialized output buffer" ); } let max_num_blocks_per_seq = bt_l.shape().dims2()?.1; diff --git a/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu b/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu index e4a24348f..e6702b033 100644 --- a/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu +++ b/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu @@ -21,126 +21,92 @@ __device__ __forceinline__ float tq_to_f(__nv_bfloat16 x) { return __bfloat162fl #define TQ_WARP 32 #define TQ_DIVUP(a,b) (((a)+(b)-1)/(b)) -static __constant__ float CB4[16] = { - -0.237664013127f,-0.180836062501f,-0.141805261760f,-0.110288414632f, - -0.082828489390f,-0.057772320256f,-0.034151583096f,-0.011302500645f, - 0.011302500645f, 0.034151583096f, 0.057772320256f, 0.082828489390f, - 0.110288414632f, 0.141805261760f, 0.180836062501f, 0.237664013127f, -}; -static __constant__ float BD4[17] = { - -1.0f,-0.209250037814f,-0.161320662130f,-0.126046838196f, - -0.096558452011f,-0.070300404823f,-0.045961951676f,-0.022727041871f, - 0.0f, 0.022727041871f, 0.045961951676f, 0.070300404823f, - 0.096558452011f, 0.126046838196f, 0.161320662130f, 0.209250037814f, 1.0f, -}; -static __constant__ float CB3[8] = { - -0.188397319183f,-0.118139828402f,-0.066585638471f,-0.021604320011f, - 0.021604320011f, 0.066585638471f, 0.118139828402f, 0.188397319183f, -}; -static __constant__ float BD3[9] = { - -1.0f,-0.153268573792f,-0.092362733436f,-0.044094979241f, - 0.0f, 0.044094979241f, 0.092362733436f, 0.153268573792f, 1.0f, -}; -static __constant__ float SGN[128] = { - -1,-1, 1,-1, 1,-1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, - -1, 1,-1,-1, 1, 1, 1, 1, 1, 1,-1,-1,-1, 1, 1,-1, - -1,-1, 1,-1, 1,-1,-1,-1, 1,-1,-1,-1, 1,-1,-1,-1, - -1,-1, 1,-1,-1,-1, 1,-1,-1,-1,-1,-1, 1,-1, 1, 1, - -1,-1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, 1,-1, 1,-1, - 1, 1, 1,-1, 1, 1, 1, 1, 1, 1, 1,-1,-1,-1,-1,-1, - -1,-1,-1,-1, 1, 1, 1,-1,-1,-1, 1, 1, 1,-1,-1,-1, - -1,-1,-1,-1, 1, 1,-1, 1, 1,-1,-1, 1, 1,-1, 1,-1, -}; - -// WHT that works with any number of active threads (uses only __syncthreads) -__device__ void wht128(float* d, int tid, int nthreads) { - for (int h = 1; h < 128; h *= 2) { - __syncthreads(); - for (int i = tid; i < 64; i += nthreads) { - int bs = (i/h)*(h*2), off = i%h; - float a = d[bs+off], b = d[bs+off+h]; - d[bs+off] = a+b; d[bs+off+h] = a-b; - } - } - __syncthreads(); -} - -__device__ void rotate128(float* d, int tid, int nthreads) { - if (tid < 128) d[tid] *= SGN[tid]; - __syncthreads(); - wht128(d, tid, nthreads); - if (tid < 128) d[tid] *= 0.08838834764831845f; - __syncthreads(); - if (tid < 128) d[tid] *= SGN[tid]; - __syncthreads(); -} - -__device__ uint8_t q4(float x) { - if (x<=BD4[1]) return 0; if (x>=BD4[16]) return 15; - int lo=1,hi=16; while(lo>1;if(x=BD3[8]) return 7; - int lo=1,hi=8; while(lo>1;if(x +// Per-head-vector compression into the paged cache. +// +// One block per (token, head); NT threads cooperate on one HS-wide vector. +// HS is a template parameter so the shared buffers, the rotation order and the +// codebook are all fixed at compile time. At HS == NT == 128 the emitted work +// is the same as the original hand-written 128-only kernel. +template __global__ void tq_cache_k( const T* in, uint8_t* cache, __half* norms, const int64_t* slots, - int nh, int hs, int bs, int in_stride, int cbs, int chs, int nbs, int nhs + int nh, int bs, int in_stride, int cbs, int chs, int nbs, int nhs ) { - int vid = blockIdx.x, tid = threadIdx.x; - int tok = vid/nh, head = vid%nh; - int64_t slot = slots[tok]; if (slot<0) return; - int bi = slot/bs, bo = slot%bs; - __shared__ float s[128]; __shared__ uint8_t ix[128]; - if (tid0;r>>=1){if(tid0;r>>=1){ if(tid1e-10f)s[tid]/=nm; __syncthreads(); - rotate128(s,tid,128); - if(tid1e-10f){ for(int i=tid;i(s,tid); + for(int i=tid;i(s[i]); + __syncthreads(); + + // 4-bit nibble pack, x=16 swizzled so a decode lane's run stays contiguous. + for(int i=tid;i +template __global__ void tq_cache_v( const T* in, uint8_t* cache, __half* norms, const int64_t* slots, - int nh, int hs, int bs, int in_stride, int vbs, int vhs, int nbs, int nhs + int nh, int bs, int in_stride, int vbs, int vhs, int nbs, int nhs ) { - int vid = blockIdx.x, tid = threadIdx.x; - int tok = vid/nh, head = vid%nh; - int64_t slot = slots[tok]; if (slot<0) return; - int bi = slot/bs, bo = slot%bs; - __shared__ float s[128]; __shared__ uint8_t ix[128]; - if (tid0;r>>=1){if(tid0;r>>=1){ if(tid1e-10f)s[tid]/=nm; __syncthreads(); - rotate128(s,tid,128); - if(tid1e-10f){ for(int i=tid;i(s,tid); + for(int i=tid;i(s[i]); + __syncthreads(); + + // 3-bit pack, 10 indices per 32-bit word, bytes strided by block size. + constexpr int NG = (HS+9)/10; + for(int i=tid;i>8); cache[bb+2*bs]=(uint8_t)(w>>16); cache[bb+3*bs]=(uint8_t)(w>>24); } @@ -152,76 +118,141 @@ __global__ void tq_cache_v( // phase 2 = end of norm computation // phase 3 = end of rotate+quantize // phase 4 = end of pack+write -template +template __global__ void tq_cache_k_clocked( const T* in, uint8_t* cache, __half* norms, const int64_t* slots, unsigned long long* clocks, - int nh, int hs, int bs, int in_stride, int cbs, int chs, int nbs, int nhs + int nh, int bs, int in_stride, int cbs, int chs, int nbs, int nhs ) { - int vid = blockIdx.x, tid = threadIdx.x; + constexpr int NT = 128; + const int vid = blockIdx.x, tid = threadIdx.x; if (tid == 0) clocks[vid * 5 + 0] = clock64(); - int tok = vid/nh, head = vid%nh; - int64_t slot = slots[tok]; if (slot<0) { if (tid==0) clocks[vid*5+4]=clock64(); return; } - int bi = slot/bs, bo = slot%bs; - __shared__ float s[128]; __shared__ uint8_t ix[128]; - if (tid0;r>>=1){if(tid0;r>>=1){ if(tid1e-10f)s[tid]/=nm; __syncthreads(); + if(nm>1e-10f){ for(int i=tid;i(s,tid); + for(int i=tid;i(s[i]); + __syncthreads(); if (tid == 0) clocks[vid * 5 + 3] = clock64(); - if(tid +template __global__ void tq_cache_v_clocked( const T* in, uint8_t* cache, __half* norms, const int64_t* slots, unsigned long long* clocks, - int nh, int hs, int bs, int in_stride, int vbs, int vhs, int nbs, int nhs + int nh, int bs, int in_stride, int vbs, int vhs, int nbs, int nhs ) { - int vid = blockIdx.x, tid = threadIdx.x; + constexpr int NT = 128; + const int vid = blockIdx.x, tid = threadIdx.x; if (tid == 0) clocks[vid * 5 + 0] = clock64(); - int tok = vid/nh, head = vid%nh; - int64_t slot = slots[tok]; if (slot<0) { if (tid==0) clocks[vid*5+4]=clock64(); return; } - int bi = slot/bs, bo = slot%bs; - __shared__ float s[128]; __shared__ uint8_t ix[128]; - if (tid0;r>>=1){if(tid0;r>>=1){ if(tid1e-10f)s[tid]/=nm; __syncthreads(); + if(nm>1e-10f){ for(int i=tid;i(s,tid); + for(int i=tid;i(s[i]); + __syncthreads(); if (tid == 0) clocks[vid * 5 + 3] = clock64(); - int ng=(hs+9)/10; - if(tid>8); cache[bb+2*bs]=(uint8_t)(w>>16); cache[bb+3*bs]=(uint8_t)(w>>24); } if (tid == 0) clocks[vid * 5 + 4] = clock64(); } +// Head dimensions the kernels are instantiated for. `TURBOQUANT_CUDA_HEAD_DIMS` +// in mistralrs-quant mirrors this list and a unit test pins the two together, +// so a head dim can never reach a launcher that has no instantiation for it. +template +static void tq_launch_cache( + const void* key, const void* value, void* kc, void* vc, void* kn, void* vn, + const int64_t* slots, int nt, int nh, int bs, + int ks, int vs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream +) { + dim3 g(nt*nh), b(128); + constexpr int VPD = ((HS+9)/10)*4; + const int vvbs = nh*VPD*bs, vvhs = VPD*bs; + tq_cache_k<<>>((const T*)key,(uint8_t*)kc,(__half*)kn,slots,nh,bs,ks,kbs,khs,nbs,nhs); + tq_cache_v<<>>((const T*)value,(uint8_t*)vc,(__half*)vn,slots,nh,bs,vs,vvbs,vvhs,nbs,nhs); +} + +template +static void tq_launch_cache_clocked( + const void* key, const void* value, void* kc, void* vc, void* kn, void* vn, + const int64_t* slots, void* clocks_k, void* clocks_v, int nt, int nh, int bs, + int ks, int vs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream +) { + dim3 g(nt*nh), b(128); + constexpr int VPD = ((HS+9)/10)*4; + const int vvbs = nh*VPD*bs, vvhs = VPD*bs; + tq_cache_k_clocked<<>>((const T*)key,(uint8_t*)kc,(__half*)kn,slots,(unsigned long long*)clocks_k,nh,bs,ks,kbs,khs,nbs,nhs); + tq_cache_v_clocked<<>>((const T*)value,(uint8_t*)vc,(__half*)vn,slots,(unsigned long long*)clocks_v,nh,bs,vs,vvbs,vvhs,nbs,nhs); +} + +// dtype: 1 = BF16 input, anything else = F16. +#define TQ_CACHE_BY_DTYPE(FN, HS_, ...) \ + do { \ + if (dtype == 1) FN(__VA_ARGS__); \ + else FN(__VA_ARGS__); \ + } while (0) + +#define TQ_CACHE_BY_HEAD_DIM(FN, ...) \ + switch (hs) { \ + case 64: TQ_CACHE_BY_DTYPE(FN, 64, __VA_ARGS__); break; \ + case 128: TQ_CACHE_BY_DTYPE(FN, 128, __VA_ARGS__); break; \ + case 256: TQ_CACHE_BY_DTYPE(FN, 256, __VA_ARGS__); break; \ + case 512: TQ_CACHE_BY_DTYPE(FN, 512, __VA_ARGS__); break; \ + default: break; \ + } + extern "C" void turbo_reshape_and_cache_clocked( const void* key, const void* value, void* kc, void* vc, void* kn, void* vn, @@ -231,16 +262,9 @@ extern "C" void turbo_reshape_and_cache_clocked( int ks, int vs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t dtype ) { - if(hs!=128)return; - dim3 g(nt*nh),b(128); - int vpd=(hs+9)/10*4; int vvbs=nh*vpd*bs; int vvhs=vpd*bs; - if (dtype == 1) { - tq_cache_k_clocked<<>>((const __nv_bfloat16*)key,(uint8_t*)kc,(__half*)kn,slots,(unsigned long long*)clocks_k,nh,hs,bs,ks,kbs,khs,nbs,nhs); - tq_cache_v_clocked<<>>((const __nv_bfloat16*)value,(uint8_t*)vc,(__half*)vn,slots,(unsigned long long*)clocks_v,nh,hs,bs,vs,vvbs,vvhs,nbs,nhs); - } else { - tq_cache_k_clocked<<>>((const __half*)key,(uint8_t*)kc,(__half*)kn,slots,(unsigned long long*)clocks_k,nh,hs,bs,ks,kbs,khs,nbs,nhs); - tq_cache_v_clocked<<>>((const __half*)value,(uint8_t*)vc,(__half*)vn,slots,(unsigned long long*)clocks_v,nh,hs,bs,vs,vvbs,vvhs,nbs,nhs); - } + TQ_CACHE_BY_HEAD_DIM(tq_launch_cache_clocked, + key, value, kc, vc, kn, vn, slots, clocks_k, clocks_v, + nt, nh, bs, ks, vs, kbs, khs, nbs, nhs, stream); } extern "C" void turbo_reshape_and_cache( @@ -251,16 +275,9 @@ extern "C" void turbo_reshape_and_cache( int ks, int vs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t dtype ) { - if(hs!=128)return; - dim3 g(nt*nh),b(128); - int vpd=(hs+9)/10*4; int vvbs=nh*vpd*bs; int vvhs=vpd*bs; - if (dtype == 1) { // BF16 - tq_cache_k<<>>((const __nv_bfloat16*)key,(uint8_t*)kc,(__half*)kn,slots,nh,hs,bs,ks,kbs,khs,nbs,nhs); - tq_cache_v<<>>((const __nv_bfloat16*)value,(uint8_t*)vc,(__half*)vn,slots,nh,hs,bs,vs,vvbs,vvhs,nbs,nhs); - } else { // F16 - tq_cache_k<<>>((const __half*)key,(uint8_t*)kc,(__half*)kn,slots,nh,hs,bs,ks,kbs,khs,nbs,nhs); - tq_cache_v<<>>((const __half*)value,(uint8_t*)vc,(__half*)vn,slots,nh,hs,bs,vs,vvbs,vvhs,nbs,nhs); - } + TQ_CACHE_BY_HEAD_DIM(tq_launch_cache, + key, value, kc, vc, kn, vn, slots, + nt, nh, bs, ks, vs, kbs, khs, nbs, nhs, stream); } // ============================================================================ @@ -277,7 +294,13 @@ extern "C" void turbo_reshape_and_cache( // OutT = float (default) or __nv_bfloat16 (saves a separate F32→BF16 cast kernel // in the dedicated decode path). -template +// +// HS is the head dimension. Everything that depends on it — the rotation order, +// the sign diagonal, the codebooks, the packed stride, how many output dims each +// thread owns — is resolved at compile time. NT stays 128 for every HS so the +// warp-reduction topology (4 warps, `red[8]`) is unchanged; wider heads give +// each thread more dims rather than adding threads. +template __global__ void tq_attn( OutT* __restrict__ out, const QT* __restrict__ query, @@ -290,10 +313,10 @@ __global__ void tq_attn( int nkvh, int mbps, int nh, float scale, int kbs, int khs, int vbs, int vhs, int nbs, int nhs ) { - constexpr int HS = 128; constexpr int NT = 128; constexpr int NW = NT / TQ_WARP; - constexpr int EPT = HS / TQ_WARP; // 4 elements per thread for dot product + constexpr int BPL = TQ_BYTES_PER_LANE(HS); // packed K bytes per lane + constexpr int DPT = (HS + NT - 1) / NT; // output dims per thread const int hidx = blockIdx.x, sidx = blockIdx.y; const int tid = threadIdx.x; @@ -302,17 +325,19 @@ __global__ void tq_attn( if (clen == 0) return; const int kvh = hidx / (nh / nkvh); - // 1. Load + rotate Q + const float* cb4 = tq_cb4(); + const float* cb3 = tq_cb3(); + + // 1. Load + rotate Q. Attention runs in the rotated domain, so this is the + // only rotation the Q·K half needs. __shared__ float qr[HS]; - if (tid < HS) qr[tid] = tq_to_f(query[sidx*nh*HS + hidx*HS + tid]); + for (int i = tid; i < HS; i += NT) qr[i] = tq_to_f(query[sidx*nh*HS + hidx*HS + i]); __syncthreads(); - rotate128(qr, tid, NT); - // After this: qr is rotated, all threads synced + tq_rotate(qr, tid); // 2. Q·K — compute logits extern __shared__ char shmem[]; float* logits = (float*)shmem; - // Initialize logits to 0 for (int i = tid; i < (int)clen; i += NT) logits[i] = 0.f; __syncthreads(); @@ -320,40 +345,52 @@ __global__ void tq_attn( const uint32_t* sbt = bt + sidx * mbps; const int nblocks = TQ_DIVUP(clen, BLOCK_SIZE); - // Each warp processes blocks in parallel. - // All 32 lanes collaborate on ONE token at a time. - // Each lane reads 2 packed bytes = 4 elements of the 128-dim K vector. - // Warp shuffle reduces the 32 partial sums into one dot product. + // Each warp takes whole blocks; all 32 lanes collaborate on one token at a + // time, each owning BPL packed bytes (2*BPL coordinates). A warp shuffle + // folds the 32 partial dot products into the logit. for (int bi = warp; bi < nblocks; bi += NW) { - int pb = sbt[bi]; - int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + const int pb = sbt[bi]; + const int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + if (tib <= 0) continue; - for (int t = 0; t < tib; t++) { - int tpos = bi * BLOCK_SIZE + t; + // This lane's packed run: contiguous, and inside one x=16 group. + const int byidx0 = lane * BPL; + const uint8_t* kbase = kc + (long long)pb*kbs + (long long)kvh*khs + + (byidx0 / 16) * BLOCK_SIZE * 16 + (byidx0 % 16); + + uint64_t cur = tq_load_klane(kbase); - // Each lane reads its 2 packed bytes (4 elements) + for (int t = 0; t < tib; t++) { +#if TQ_PREFETCH > 1 + // Hopper/Blackwell: stage the next token's bytes while reducing this + // one. Same arithmetic, one fewer stall. + const uint64_t nxt = (t + 1 < tib) + ? tq_load_klane(kbase + (t + 1) * 16) : 0ull; +#endif float qk = 0.f; #pragma unroll - for (int e = 0; e < EPT/2; e++) { - int byidx = lane * (EPT/2) + e; // 0..63 - int x = 16; - int koff = pb*kbs + kvh*khs + (byidx/x)*BLOCK_SIZE*x + t*x + (byidx%x); - uint8_t pk = kc[koff]; - int d0 = byidx*2, d1 = byidx*2+1; - qk += qr[d0] * CB4[pk & 0xF]; - qk += qr[d1] * CB4[(pk>>4) & 0xF]; + for (int j = 0; j < BPL; j++) { + const uint8_t pk = (uint8_t)((cur >> (8 * j)) & 0xFFull); + const int d0 = (byidx0 + j) * 2; + qk += qr[d0] * cb4[pk & 0xF]; + qk += qr[d0 + 1] * cb4[(pk >> 4) & 0xF]; } - // Warp reduction: sum 32 partial dot products #pragma unroll for (int mask = TQ_WARP/2; mask > 0; mask >>= 1) qk += __shfl_xor_sync(0xffffffff, qk, mask); - float knorm = __half2float(kn[pb*nbs + kvh*nhs + t]); + const float knorm = __half2float(kn[pb*nbs + kvh*nhs + t]); qk *= knorm * scale; - if (lane == 0) logits[tpos] = qk; + if (lane == 0) logits[bi*BLOCK_SIZE + t] = qk; qk_max = fmaxf(qk_max, qk); + +#if TQ_PREFETCH > 1 + cur = nxt; +#else + if (t + 1 < tib) cur = tq_load_klane(kbase + (t + 1) * 16); +#endif } } @@ -361,7 +398,7 @@ __global__ void tq_attn( #pragma unroll for (int mask = TQ_WARP/2; mask >= 1; mask /= 2) qk_max = fmaxf(qk_max, __shfl_xor_sync(0xffffffff, qk_max, mask)); - __shared__ float red[8]; + __shared__ float red[2 * NW]; if (lane == 0) red[warp] = qk_max; __syncthreads(); qk_max = (lane < NW) ? red[lane] : -FLT_MAX; @@ -373,7 +410,7 @@ __global__ void tq_attn( // 3. Softmax float esum = 0.f; for (int i = tid; i < (int)clen; i += NT) { - float v = __expf(logits[i] - qk_max); + const float v = __expf(logits[i] - qk_max); logits[i] = v; esum += v; } @@ -386,53 +423,54 @@ __global__ void tq_attn( #pragma unroll for (int mask = NW/2; mask >= 1; mask /= 2) esum += __shfl_xor_sync(0xffffffff, esum, mask); - float inv = __fdividef(1.f, __shfl_sync(0xffffffff, esum, 0) + 1e-6f); + const float inv = __fdividef(1.f, __shfl_sync(0xffffffff, esum, 0) + 1e-6f); for (int i = tid; i < (int)clen; i += NT) logits[i] *= inv; __syncthreads(); - // 4. V accumulation - // Each thread accumulates EPT=4 output dimensions - // thread tid handles dims: for Q·K we used lane-based, for V we use tid-based - // tid handles dims [tid*EPT .. tid*EPT+3] (but tid goes 0..127, EPT=4 would be 512 dims) - // Actually: 128 threads, 128 dims, each thread handles 1 dim - // Simpler: each thread handles dim = tid (one dim per thread) - - float acc = 0.f; - if (tid < HS) { - int dim = tid; - int group = dim / 10; - int pos = dim % 10; - - for (int bi = 0; bi < nblocks; bi++) { - int pb = sbt[bi]; - int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); - - for (int t = 0; t < tib; t++) { - float w = logits[bi*BLOCK_SIZE + t]; - if (w < 1e-8f) continue; - - int bb = pb*vbs + kvh*vhs + group*4*BLOCK_SIZE + t; - uint32_t word = (uint32_t)vc[bb] | - ((uint32_t)vc[bb+BLOCK_SIZE]<<8) | - ((uint32_t)vc[bb+2*BLOCK_SIZE]<<16) | - ((uint32_t)vc[bb+3*BLOCK_SIZE]<<24); - uint8_t vi = (word>>(pos*3))&7; - float vnorm = __half2float(vn[pb*nbs+kvh*nhs+t]); - acc += w * CB3[vi] * vnorm; + // 4. V accumulation — each thread owns DPT output dims, strided by NT. + float acc[DPT]; + #pragma unroll + for (int r = 0; r < DPT; r++) acc[r] = 0.f; + + for (int bi = 0; bi < nblocks; bi++) { + const int pb = sbt[bi]; + const int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + + for (int t = 0; t < tib; t++) { + const float w = logits[bi*BLOCK_SIZE + t]; + if (w < 1e-8f) continue; + const float vnorm = __half2float(vn[pb*nbs + kvh*nhs + t]); + + #pragma unroll + for (int r = 0; r < DPT; r++) { + const int dim = tid + r * NT; + if (dim >= HS) continue; + const int group = dim / 10, pos = dim % 10; + const long long bb = (long long)pb*vbs + (long long)kvh*vhs + group*4*BLOCK_SIZE + t; + const uint32_t word = (uint32_t)vc[bb] + | ((uint32_t)vc[bb + BLOCK_SIZE] << 8) + | ((uint32_t)vc[bb + 2*BLOCK_SIZE] << 16) + | ((uint32_t)vc[bb + 3*BLOCK_SIZE] << 24); + const uint8_t vi = (word >> (pos * 3)) & 7; + acc[r] += w * cb3[vi] * vnorm; } } } - // 5. Inverse rotation + output + // 5. Inverse rotation + output. V was quantized in the rotated basis, so the + // accumulator lands there too; D·H·D is its own inverse. __shared__ float obuf[HS]; - if (tid < HS) obuf[tid] = acc; + #pragma unroll + for (int r = 0; r < DPT; r++) { + const int dim = tid + r * NT; + if (dim < HS) obuf[dim] = acc[r]; + } __syncthreads(); - rotate128(obuf, tid, NT); + tq_rotate(obuf, tid); - if (tid < HS) { - // Cast on store: float (no-op), __half, or __nv_bfloat16. - const int oi = sidx*nh*HS + hidx*HS + tid; - const float v = obuf[tid]; + for (int i = tid; i < HS; i += NT) { + const int oi = sidx*nh*HS + hidx*HS + i; + const float v = obuf[i]; if constexpr (std::is_same::value) { out[oi] = __float2bfloat16(v); } else if constexpr (std::is_same::value) { @@ -450,7 +488,7 @@ __global__ void tq_attn( // phase 3 = after softmax // phase 4 = after V accumulate // phase 5 = after inverse rotate + write -template +template __global__ void tq_attn_clocked( __nv_bfloat16* __restrict__ out, const QT* __restrict__ query, @@ -464,10 +502,12 @@ __global__ void tq_attn_clocked( int nkvh, int mbps, int nh, float scale, int kbs, int khs, int vbs, int vhs, int nbs, int nhs ) { - constexpr int HS = 128; constexpr int NT = 128; constexpr int NW = NT / TQ_WARP; - constexpr int EPT = HS / TQ_WARP; + constexpr int BPL = TQ_BYTES_PER_LANE(HS); + constexpr int DPT = (HS + NT - 1) / NT; + const float* cb4 = tq_cb4(); + const float* cb3 = tq_cb3(); const int hidx = blockIdx.x, sidx = blockIdx.y; const int tid = threadIdx.x; @@ -483,9 +523,9 @@ __global__ void tq_attn_clocked( const int kvh = hidx / (nh / nkvh); __shared__ float qr[HS]; - if (tid < HS) qr[tid] = tq_to_f(query[sidx*nh*HS + hidx*HS + tid]); + for (int i = tid; i < HS; i += NT) qr[i] = tq_to_f(query[sidx*nh*HS + hidx*HS + i]); __syncthreads(); - rotate128(qr, tid, NT); + tq_rotate(qr, tid); if (tid == 0) clocks[block_id * 6 + 1] = clock64(); extern __shared__ char shmem[]; @@ -496,29 +536,45 @@ __global__ void tq_attn_clocked( float qk_max = -FLT_MAX; const uint32_t* sbt = bt + sidx * mbps; const int nblocks = TQ_DIVUP(clen, BLOCK_SIZE); + // This loop must stay a mirror of tq_attn's: the whole point of the clocked + // variant is that its phase-2 stamps describe the kernel that actually + // runs. A scalar gather here against a vectorized one there would have made + // the timings describe a kernel nobody launches. for (int bi = warp; bi < nblocks; bi += NW) { int pb = sbt[bi]; int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + if (tib <= 0) continue; + + const int byidx0 = lane * BPL; + const uint8_t* kbase = kc + (long long)pb*kbs + (long long)kvh*khs + + (byidx0 / 16) * BLOCK_SIZE * 16 + (byidx0 % 16); + uint64_t cur = tq_load_klane(kbase); + for (int t = 0; t < tib; t++) { - int tpos = bi * BLOCK_SIZE + t; +#if TQ_PREFETCH > 1 + const uint64_t nxt = (t + 1 < tib) + ? tq_load_klane(kbase + (t + 1) * 16) : 0ull; +#endif float qk = 0.f; #pragma unroll - for (int e = 0; e < EPT/2; e++) { - int byidx = lane * (EPT/2) + e; - int x = 16; - int koff = pb*kbs + kvh*khs + (byidx/x)*BLOCK_SIZE*x + t*x + (byidx%x); - uint8_t pk = kc[koff]; - int d0 = byidx*2, d1 = byidx*2+1; - qk += qr[d0] * CB4[pk & 0xF]; - qk += qr[d1] * CB4[(pk>>4) & 0xF]; + for (int j = 0; j < BPL; j++) { + const uint8_t pk = (uint8_t)((cur >> (8 * j)) & 0xFFull); + const int d0 = (byidx0 + j) * 2; + qk += qr[d0] * cb4[pk & 0xF]; + qk += qr[d0 + 1] * cb4[(pk >> 4) & 0xF]; } #pragma unroll for (int mask = TQ_WARP/2; mask > 0; mask >>= 1) qk += __shfl_xor_sync(0xffffffff, qk, mask); float knorm = __half2float(kn[pb*nbs + kvh*nhs + t]); qk *= knorm * scale; - if (lane == 0) logits[tpos] = qk; + if (lane == 0) logits[bi * BLOCK_SIZE + t] = qk; qk_max = fmaxf(qk_max, qk); +#if TQ_PREFETCH > 1 + cur = nxt; +#else + if (t + 1 < tib) cur = tq_load_klane(kbase + (t + 1) * 16); +#endif } } if (tid == 0) clocks[block_id * 6 + 2] = clock64(); @@ -526,7 +582,7 @@ __global__ void tq_attn_clocked( #pragma unroll for (int mask = TQ_WARP/2; mask >= 1; mask /= 2) qk_max = fmaxf(qk_max, __shfl_xor_sync(0xffffffff, qk_max, mask)); - __shared__ float red[8]; + __shared__ float red[2 * NW]; if (lane == 0) red[warp] = qk_max; __syncthreads(); qk_max = (lane < NW) ? red[lane] : -FLT_MAX; @@ -555,40 +611,157 @@ __global__ void tq_attn_clocked( __syncthreads(); if (tid == 0) clocks[block_id * 6 + 3] = clock64(); - float acc = 0.f; - if (tid < HS) { - int dim = tid; - int group = dim / 10; - int pos = dim % 10; - for (int bi = 0; bi < nblocks; bi++) { - int pb = sbt[bi]; - int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); - for (int t = 0; t < tib; t++) { - float w = logits[bi*BLOCK_SIZE + t]; - if (w < 1e-8f) continue; - int bb = pb*vbs + kvh*vhs + group*4*BLOCK_SIZE + t; - uint32_t word = (uint32_t)vc[bb] | + float acc[DPT]; + #pragma unroll + for (int r = 0; r < DPT; r++) acc[r] = 0.f; + for (int bi = 0; bi < nblocks; bi++) { + int pb = sbt[bi]; + int tib = min(BLOCK_SIZE, (int)clen - bi*BLOCK_SIZE); + for (int t = 0; t < tib; t++) { + float w = logits[bi*BLOCK_SIZE + t]; + if (w < 1e-8f) continue; + float vnorm = __half2float(vn[pb*nbs+kvh*nhs+t]); + #pragma unroll + for (int r = 0; r < DPT; r++) { + const int dim = tid + r * NT; + if (dim >= HS) continue; + const int group = dim / 10, pos = dim % 10; + const long long bb = (long long)pb*vbs + (long long)kvh*vhs + group*4*BLOCK_SIZE + t; + const uint32_t word = (uint32_t)vc[bb] | ((uint32_t)vc[bb+BLOCK_SIZE]<<8) | ((uint32_t)vc[bb+2*BLOCK_SIZE]<<16) | ((uint32_t)vc[bb+3*BLOCK_SIZE]<<24); - uint8_t vi = (word>>(pos*3))&7; - float vnorm = __half2float(vn[pb*nbs+kvh*nhs+t]); - acc += w * CB3[vi] * vnorm; + const uint8_t vi = (word>>(pos*3))&7; + acc[r] += w * cb3[vi] * vnorm; } } } if (tid == 0) clocks[block_id * 6 + 4] = clock64(); __shared__ float obuf[HS]; - if (tid < HS) obuf[tid] = acc; + #pragma unroll + for (int r = 0; r < DPT; r++) { + const int dim = tid + r * NT; + if (dim < HS) obuf[dim] = acc[r]; + } __syncthreads(); - rotate128(obuf, tid, NT); - if (tid < HS) { - out[sidx*nh*HS + hidx*HS + tid] = __float2bfloat16(obuf[tid]); + tq_rotate(obuf, tid); + for (int i = tid; i < HS; i += NT) { + out[sidx*nh*HS + hidx*HS + i] = __float2bfloat16(obuf[i]); } if (tid == 0) clocks[block_id * 6 + 5] = clock64(); } +// ============================================================================ +// Launch wrappers +// ============================================================================ + +/// Launch one instantiation, opting in to the arch's real shared-memory budget +/// first. +/// +/// The logits array is dynamic shared memory sized by the padded context +/// length, so it grows without bound as context grows. CUDA caps dynamic shared +/// memory at 48 KB per block unless a kernel explicitly opts in — 164 KB on +/// Ampere, 227 KB on Hopper (SM90) and Blackwell (SM100/SM103). Without this +/// call, any context past roughly 12k tokens fails to launch. `granted` is a +/// per-instantiation static, so the driver call happens once per kernel rather +/// than once per decode step. +template +static void tq_launch_attn( + void* out, const void* query, const void* kc, const void* vc, + const void* kn, const void* vn, const uint32_t* bt, const uint32_t* cl, + int nkvh, int mbps, int nh, float scale, + int kbs, int khs, int vbs, int vhs, int nbs, int nhs, + dim3 grid, dim3 block, int smem, cudaStream_t stream +) { + // Not cached: the requestable maximum differs by architecture, so a single + // process-wide flag would be wrong the moment two devices differ. A failure + // here means the request exceeds the device's opt-in cap; the launch below + // then fails with cudaErrorInvalidValue, which surfaces on the next sync. + if (smem > 48 * 1024) { + (void)cudaFuncSetAttribute( + (const void*)(tq_attn), + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + tq_attn<<>>( + (OutT*)out, (const QT*)query, (const uint8_t*)kc, (const uint8_t*)vc, + (const __half*)kn, (const __half*)vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vbs, vhs, nbs, nhs); +} + +template +static void tq_launch_attn_clocked( + void* out_bf16, const void* query, const void* kc, const void* vc, + const void* kn, const void* vn, const uint32_t* bt, const uint32_t* cl, + void* clocks, int nkvh, int mbps, int nh, float scale, + int kbs, int khs, int vbs, int vhs, int nbs, int nhs, + dim3 grid, dim3 block, int smem, cudaStream_t stream +) { + // Not cached: the requestable maximum differs by architecture, so a single + // process-wide flag would be wrong the moment two devices differ. A failure + // here means the request exceeds the device's opt-in cap; the launch below + // then fails with cudaErrorInvalidValue, which surfaces on the next sync. + if (smem > 48 * 1024) { + (void)cudaFuncSetAttribute( + (const void*)(tq_attn_clocked), + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + tq_attn_clocked<<>>( + (__nv_bfloat16*)out_bf16, (const QT*)query, + (const uint8_t*)kc, (const uint8_t*)vc, + (const __half*)kn, (const __half*)vn, bt, cl, + (unsigned long long*)clocks, + nkvh, mbps, nh, scale, kbs, khs, vbs, vhs, nbs, nhs); +} + +// Paged block sizes the cache engine emits. +#define TQ_ATTN_BY_BS(HS_, QT, OUTT, ...) \ + switch (bs) { \ + case 8: tq_launch_attn(__VA_ARGS__); break; \ + case 16: tq_launch_attn(__VA_ARGS__); break; \ + case 32: tq_launch_attn(__VA_ARGS__); break; \ + default: break; \ + } + +// Head dimensions with codebooks and sign diagonals in turbo_paged_attention.cuh. +// `TURBOQUANT_CUDA_HEAD_DIMS` in mistralrs-quant mirrors this set and a unit +// test pins the two together, so the Rust side refuses a head dim before it can +// reach a launcher with no instantiation for it. +#define TQ_ATTN_BY_HS(QT, OUTT, ...) \ + switch (hs) { \ + case 64: TQ_ATTN_BY_BS(64, QT, OUTT, __VA_ARGS__) break; \ + case 128: TQ_ATTN_BY_BS(128, QT, OUTT, __VA_ARGS__) break; \ + case 256: TQ_ATTN_BY_BS(256, QT, OUTT, __VA_ARGS__) break; \ + case 512: TQ_ATTN_BY_BS(512, QT, OUTT, __VA_ARGS__) break; \ + default: break; \ + } + +#define TQ_ATTN_CLOCKED_BY_BS(HS_, QT, ...) \ + switch (bs) { \ + case 8: tq_launch_attn_clocked(__VA_ARGS__); break; \ + case 16: tq_launch_attn_clocked(__VA_ARGS__); break; \ + case 32: tq_launch_attn_clocked(__VA_ARGS__); break; \ + default: break; \ + } + +#define TQ_ATTN_CLOCKED_BY_HS(QT, ...) \ + switch (hs) { \ + case 64: TQ_ATTN_CLOCKED_BY_BS(64, QT, __VA_ARGS__) break; \ + case 128: TQ_ATTN_CLOCKED_BY_BS(128, QT, __VA_ARGS__) break; \ + case 256: TQ_ATTN_CLOCKED_BY_BS(256, QT, __VA_ARGS__) break; \ + case 512: TQ_ATTN_CLOCKED_BY_BS(512, QT, __VA_ARGS__) break; \ + default: break; \ + } + +// Geometry shared by every entry point below. +#define TQ_ATTN_PROLOGUE() \ + const int vpd = (hs + 9) / 10 * 4; \ + const int vvbs = nkvh * vpd * bs, vvhs = vpd * bs; \ + const int padded = TQ_DIVUP(mcl, bs) * bs; \ + const int smem = padded * (int)sizeof(float); \ + dim3 grid(nh, ns, 1); \ + dim3 block(128); + extern "C" void turbo_paged_attention_v1_bf16out_clocked( void* out_bf16, const void* query, const void* kc, const void* vc, @@ -600,34 +773,17 @@ extern "C" void turbo_paged_attention_v1_bf16out_clocked( int mbps, int qs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t qdtype ) { - if (hs != 128) return; - int vpd = (hs + 9) / 10 * 4; - int vvbs = nkvh * vpd * bs, vvhs = vpd * bs; - int padded = TQ_DIVUP(mcl, bs) * bs; - int smem = padded * sizeof(float); - dim3 grid(nh, ns, 1); - dim3 block(128); + TQ_ATTN_PROLOGUE() if (qdtype == 1) { - switch (bs) { - case 8: tq_attn_clocked<8, __nv_bfloat16><<>>( - (__nv_bfloat16*)out_bf16, (const __nv_bfloat16*)query, - (const uint8_t*)kc, (const uint8_t*)vc, - (const __half*)kn, (const __half*)vn, - bt, cl, (unsigned long long*)clocks, - nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs); break; - case 16: tq_attn_clocked<16, __nv_bfloat16><<>>( - (__nv_bfloat16*)out_bf16, (const __nv_bfloat16*)query, - (const uint8_t*)kc, (const uint8_t*)vc, - (const __half*)kn, (const __half*)vn, - bt, cl, (unsigned long long*)clocks, - nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs); break; - case 32: tq_attn_clocked<32, __nv_bfloat16><<>>( - (__nv_bfloat16*)out_bf16, (const __nv_bfloat16*)query, - (const uint8_t*)kc, (const uint8_t*)vc, - (const __half*)kn, (const __half*)vn, - bt, cl, (unsigned long long*)clocks, - nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs); break; - } + TQ_ATTN_CLOCKED_BY_HS(__nv_bfloat16, + out_bf16, query, kc, vc, kn, vn, bt, cl, clocks, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } else { + TQ_ATTN_CLOCKED_BY_HS(__half, + out_bf16, query, kc, vc, kn, vn, bt, cl, clocks, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) } } @@ -641,20 +797,11 @@ extern "C" void turbo_paged_attention_v1_f16( int mbps, int qs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream ) { - if (hs != 128) return; - int vpd=(hs+9)/10*4; - int vvbs=nkvh*vpd*bs, vvhs=vpd*bs; - int padded = TQ_DIVUP(mcl,bs)*bs; - int smem = padded * sizeof(float); - // Ensure enough for output buffer too (128 floats = 512 bytes, always < logits) - dim3 grid(nh, ns, 1); - dim3 block(128); - - switch(bs){ - case 8: tq_attn<8><<>>((float*)out,(const __half*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; - case 16: tq_attn<16><<>>((float*)out,(const __half*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; - case 32: tq_attn<32><<>>((float*)out,(const __half*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; - } + TQ_ATTN_PROLOGUE() + TQ_ATTN_BY_HS(__half, float, + out, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) } // Dtype-generic entry point for the dedicated decode path @@ -668,22 +815,18 @@ extern "C" void turbo_paged_attention_v1( int mbps, int qs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t dtype ) { - if (hs != 128) return; - int vpd=(hs+9)/10*4; - int vvbs=nkvh*vpd*bs, vvhs=vpd*bs; - int padded = TQ_DIVUP(mcl,bs)*bs; - int smem = padded * sizeof(float); - dim3 grid(nh, ns, 1); - dim3 block(128); - #define TQ_LAUNCH(QT) \ - switch(bs){ \ - case 8: tq_attn<8, QT><<>>((float*)out,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - case 16: tq_attn<16, QT><<>>((float*)out,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - case 32: tq_attn<32, QT><<>>((float*)out,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - } - if (dtype == 1) { TQ_LAUNCH(__nv_bfloat16); } - else { TQ_LAUNCH(__half); } - #undef TQ_LAUNCH + TQ_ATTN_PROLOGUE() + if (dtype == 1) { + TQ_ATTN_BY_HS(__nv_bfloat16, float, + out, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } else { + TQ_ATTN_BY_HS(__half, float, + out, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } } // BF16-output entry point: writes BF16 directly, eliminating the F32→BF16 cast @@ -698,20 +841,16 @@ extern "C" void turbo_paged_attention_v1_bf16out( int mbps, int qs, int kbs, int khs, int nbs, int nhs, cudaStream_t stream, uint32_t qdtype ) { - if (hs != 128) return; - int vpd=(hs+9)/10*4; - int vvbs=nkvh*vpd*bs, vvhs=vpd*bs; - int padded = TQ_DIVUP(mcl,bs)*bs; - int smem = padded * sizeof(float); - dim3 grid(nh, ns, 1); - dim3 block(128); - #define TQ_LAUNCH_BF16OUT(QT) \ - switch(bs){ \ - case 8: tq_attn<8, QT, __nv_bfloat16><<>>((__nv_bfloat16*)out_bf16,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - case 16: tq_attn<16, QT, __nv_bfloat16><<>>((__nv_bfloat16*)out_bf16,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - case 32: tq_attn<32, QT, __nv_bfloat16><<>>((__nv_bfloat16*)out_bf16,(const QT*)query,(const uint8_t*)kc,(const uint8_t*)vc,(const __half*)kn,(const __half*)vn,bt,cl,nkvh,mbps,nh,scale,kbs,khs,vvbs,vvhs,nbs,nhs); break; \ - } - if (qdtype == 1) { TQ_LAUNCH_BF16OUT(__nv_bfloat16); } - else { TQ_LAUNCH_BF16OUT(__half); } - #undef TQ_LAUNCH_BF16OUT + TQ_ATTN_PROLOGUE() + if (qdtype == 1) { + TQ_ATTN_BY_HS(__nv_bfloat16, __nv_bfloat16, + out_bf16, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } else { + TQ_ATTN_BY_HS(__half, __nv_bfloat16, + out_bf16, query, kc, vc, kn, vn, bt, cl, + nkvh, mbps, nh, scale, kbs, khs, vvbs, vvhs, nbs, nhs, + grid, block, smem, stream) + } } diff --git a/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh b/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh index 7e037b615..1e8debb31 100644 --- a/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh +++ b/mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh @@ -1,391 +1,379 @@ /** - * TurboQuant Paged Attention Kernel + * TurboQuant device tables and generic rotation — shared by every TurboQuant + * CUDA kernel in this crate. * - * A variant of the standard paged attention kernel that reads compressed KV cache - * using TurboQuant codebook lookups instead of direct memory loads or FP8 dequantization. + * WHAT THIS FILE IS + * ----------------- + * TurboQuant compresses a KV head vector by (1) recording its L2 norm, (2) + * applying the involutory rotation D·H·D to the unit vector, and (3) scalar + * quantizing each rotated coordinate against a Lloyd-Max codebook. Attention + * then runs *in the rotated domain*: because D·H·D is orthogonal, q·k is + * preserved under it, so the K dot product needs no inverse rotation at all. + * Only the V accumulation, which lands in the rotated basis, is rotated back + * once at the end. * - * Key differences from standard paged attention: - * 1. Q is pre-rotated via WHT before the Q·K loop (once per query) - * 2. K is stored as packed indices; lookup centroids from codebook in shared memory - * 3. K dot product operates in the rotated domain (no inverse rotation needed) - * 4. V is stored as packed indices; dequantized inline via codebook lookup - * 5. Norms are stored per-token and applied after the dot product + * DIMENSION SUPPORT + * ----------------- + * The tables below cover head_dim 64, 128, 256 and 512. Both `D` (the sign + * diagonal) and the codebook are functions of the dimension, so a kernel + * cannot simply reuse the 128 tables at another width — the codebook centroids + * shrink as sqrt(1/d) because the rotated coordinates concentrate. * - * This kernel handles the 4-bit key / 3-bit value (default) configuration. - * The cache layout per block: - * k_cache: [num_blocks, num_kv_heads, packed_k_bytes_per_head, block_size] - * v_cache: [num_blocks, num_kv_heads, packed_v_bytes_per_head, block_size] - * k_norms: [num_blocks, num_kv_heads, block_size] (half) - * v_norms: [num_blocks, num_kv_heads, block_size] (half) + * PROVENANCE — DO NOT HAND-EDIT THE TABLES + * ---------------------------------------- + * Every `__constant__` array below was emitted by + * cargo run -p mistralrs-quant --example emit_turboquant_cuda_tables -- 64 128 256 512 + * which calls the same `turboquant::wht::generate_signs` and + * `turboquant::codebook::get_codebook` that the Rust compressor uses. The + * values are therefore correct by construction rather than by transcription. * - * Reference: TurboQuant ICLR 2026 (arXiv:2504.19874), 0xSero/turboquant Triton kernels. + * `turboquant::cuda_tables` re-derives them at test time and diffs against this + * file, so a hand-edit here fails `cargo test -p mistralrs-quant`. That test is + * the only thing standing between a silent sign-table drift and a cache that + * decodes to noise, which is why it parses this header directly. + * + * ARCH SPECIALISATION (D16) + * ------------------------- + * `TQ_PREFETCH` selects the software-pipeline depth of the packed-K gather. + * Hopper and Blackwell have deeper memory latency to hide and enough registers + * to hide it, so they run a 2-deep pipeline; pre-Hopper parts run the flat + * loop. Both produce identical arithmetic — this is a scheduling choice, not a + * numerical one. */ +#pragma once + #include #include #include -#ifndef WARP_SIZE -#define WARP_SIZE 32 +// --------------------------------------------------------------------------- +// Arch specialisation +// --------------------------------------------------------------------------- +// Hopper (SM90) and Blackwell (SM100/SM103) hide global latency better with a +// software-pipelined gather; older parts are latency-bound on the extra +// registers instead. Identical results either way. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +#define TQ_PREFETCH 2 +#else +#define TQ_PREFETCH 1 #endif -#define MAX_TQ(a, b) ((a) > (b) ? (a) : (b)) -#define MIN_TQ(a, b) ((a) < (b) ? (a) : (b)) -#define DIVIDE_ROUND_UP_TQ(a, b) (((a) + (b) - 1) / (b)) -// ============================================================================ -// Codebooks (same values as turbo_wht.cu) -// ============================================================================ +// --------------------------------------------------------------------------- +// Tables — generated, see PROVENANCE above. +// --------------------------------------------------------------------------- -__constant__ float TQ_CODEBOOK_4BIT[16] = { - -0.237664013127f, -0.180836062501f, -0.141805261760f, -0.110288414632f, - -0.082828489390f, -0.057772320256f, -0.034151583096f, -0.011302500645f, - 0.011302500645f, 0.034151583096f, 0.057772320256f, 0.082828489390f, - 0.110288414632f, 0.141805261760f, 0.180836062501f, 0.237664013127f, +// ==== head_dim = 64 (seed 42) ==== +static __constant__ float TQ_SGN_64[64] = { + 1,-1, 1,-1, 1, 1,-1,-1,-1,-1, 1, 1, 1,-1, 1, 1, + 1,-1, 1,-1,-1,-1, 1, 1, 1,-1, 1,-1, 1, 1, 1, 1, + 1,-1,-1, 1,-1,-1, 1, 1,-1, 1, 1,-1,-1,-1, 1,-1, + -1,-1,-1, 1, 1, 1, 1,-1, 1,-1, 1,-1, 1, 1, 1, 1, +}; +static __constant__ float TQ_CB3_64[8] = { + -0.263914018869f,-0.166167899966f,-0.093832284212f,-0.030469184741f, + 0.030469184741f,0.093832284212f,0.166167899966f,0.263914018869f, +}; +static __constant__ float TQ_BD3_64[9] = { + -1.000000000000f,-0.215040951967f,-0.130000084639f,-0.062150731683f, + 0.000000000000f,0.062150731683f,0.130000084639f,0.215040951967f, + 1.000000000000f, +}; +static __constant__ float TQ_CB4_64[16] = { + -0.330796420574f,-0.252913773060f,-0.198856174946f,-0.154925554991f, + -0.116486772895f,-0.081311784685f,-0.048089791089f,-0.015919025987f, + 0.015919025987f,0.048089791089f,0.081311784685f,0.116486772895f, + 0.154925554991f,0.198856174946f,0.252913773060f,0.330796420574f, +}; +static __constant__ float TQ_BD4_64[17] = { + -1.000000000000f,-0.291855096817f,-0.225884988904f,-0.176890864968f, + -0.135706171393f,-0.098899275064f,-0.064700789750f,-0.032004408538f, + 0.000000000000f,0.032004408538f,0.064700789750f,0.098899275064f, + 0.135706171393f,0.176890864968f,0.225884988904f,0.291855096817f, + 1.000000000000f, }; -__constant__ float TQ_CODEBOOK_3BIT[8] = { - -0.188397319183f, -0.118139828402f, -0.066585638471f, -0.021604320011f, - 0.021604320011f, 0.066585638471f, 0.118139828402f, 0.188397319183f, +// ==== head_dim = 128 (seed 42) ==== +static __constant__ float TQ_SGN_128[128] = { + -1,-1, 1,-1, 1,-1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, + -1, 1,-1,-1, 1, 1, 1, 1, 1, 1,-1,-1,-1, 1, 1,-1, + -1,-1, 1,-1, 1,-1,-1,-1, 1,-1,-1,-1, 1,-1,-1,-1, + -1,-1, 1,-1,-1,-1, 1,-1,-1,-1,-1,-1, 1,-1, 1, 1, + -1,-1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, 1,-1, 1,-1, + 1, 1, 1,-1, 1, 1, 1, 1, 1, 1, 1,-1,-1,-1,-1,-1, + -1,-1,-1,-1, 1, 1, 1,-1,-1,-1, 1, 1, 1,-1,-1,-1, + -1,-1,-1,-1, 1, 1,-1, 1, 1,-1,-1, 1, 1,-1, 1,-1, +}; +static __constant__ float TQ_CB3_128[8] = { + -0.188397318125f,-0.118139825761f,-0.066585637629f,-0.021604320034f, + 0.021604320034f,0.066585637629f,0.118139825761f,0.188397318125f, +}; +static __constant__ float TQ_BD3_128[9] = { + -1.000000000000f,-0.153268575668f,-0.092362731695f,-0.044094979763f, + 0.000000000000f,0.044094979763f,0.092362731695f,0.153268575668f, + 1.000000000000f, +}; +static __constant__ float TQ_CB4_128[16] = { + -0.237664014101f,-0.180836066604f,-0.141805261374f,-0.110288411379f, + -0.082828491926f,-0.057772319764f,-0.034151583910f,-0.011302500963f, + 0.011302500963f,0.034151583910f,0.057772319764f,0.082828491926f, + 0.110288411379f,0.141805261374f,0.180836066604f,0.237664014101f, +}; +static __constant__ float TQ_BD4_128[17] = { + -1.000000000000f,-0.209250032902f,-0.161320656538f,-0.126046836376f, + -0.096558451653f,-0.070300407708f,-0.045961949974f,-0.022727042437f, + 0.000000000000f,0.022727042437f,0.045961949974f,0.070300407708f, + 0.096558451653f,0.126046836376f,0.161320656538f,0.209250032902f, + 1.000000000000f, }; -__constant__ float TQ_CODEBOOK_2BIT[4] = { - -0.133041590561f, -0.039991612341f, 0.039991612341f, 0.133041590561f, +// ==== head_dim = 256 (seed 42) ==== +static __constant__ float TQ_SGN_256[256] = { + 1,-1,-1, 1,-1, 1,-1, 1,-1, 1,-1,-1,-1, 1, 1, 1, + 1,-1,-1,-1, 1,-1,-1,-1,-1, 1,-1, 1,-1, 1, 1, 1, + 1, 1, 1, 1,-1, 1, 1,-1,-1,-1, 1, 1, 1,-1,-1, 1, + -1,-1, 1,-1, 1,-1,-1,-1,-1,-1,-1,-1,-1,-1,-1, 1, + -1,-1, 1,-1, 1, 1, 1, 1,-1, 1, 1,-1,-1, 1,-1, 1, + 1,-1,-1, 1,-1, 1,-1,-1, 1,-1,-1,-1,-1,-1,-1,-1, + 1, 1, 1, 1, 1, 1,-1,-1, 1,-1,-1,-1,-1,-1, 1, 1, + 1,-1,-1,-1, 1, 1,-1, 1, 1, 1, 1,-1, 1,-1, 1,-1, + -1, 1,-1, 1,-1, 1, 1,-1,-1, 1,-1, 1,-1, 1,-1, 1, + -1,-1,-1, 1, 1,-1,-1, 1,-1,-1,-1,-1, 1,-1,-1, 1, + -1, 1, 1, 1, 1,-1, 1,-1,-1, 1, 1, 1, 1, 1, 1,-1, + -1, 1, 1,-1, 1,-1,-1,-1,-1, 1,-1, 1,-1, 1, 1,-1, + 1, 1,-1, 1, 1,-1, 1, 1, 1,-1, 1,-1, 1, 1,-1, 1, + -1,-1, 1, 1,-1,-1, 1, 1, 1, 1, 1, 1,-1,-1,-1, 1, + -1,-1,-1,-1,-1,-1, 1, 1, 1,-1, 1,-1,-1, 1,-1, 1, + 1, 1, 1, 1,-1, 1,-1, 1,-1,-1, 1, 1,-1,-1,-1,-1, +}; +static __constant__ float TQ_CB3_256[8] = { + -0.133854493499f,-0.083765551448f,-0.047166757286f,-0.015297502279f, + 0.015297502279f,0.047166757286f,0.083765551448f,0.133854493499f, +}; +static __constant__ float TQ_BD3_256[9] = { + -1.000000000000f,-0.108810022473f,-0.065466158092f,-0.031232129782f, + 0.000000000000f,0.031232129782f,0.065466158092f,0.108810022473f, + 1.000000000000f, +}; +static __constant__ float TQ_CB4_256[16] = { + -0.169410735369f,-0.128588363528f,-0.100698113441f,-0.078249387443f, + -0.058732159436f,-0.040949229151f,-0.024200897664f,-0.008008380421f, + 0.008008380421f,0.024200897664f,0.040949229151f,0.058732159436f, + 0.078249387443f,0.100698113441f,0.128588363528f,0.169410735369f, +}; +static __constant__ float TQ_BD4_256[17] = { + -1.000000000000f,-0.148999556899f,-0.114643231034f,-0.089473746717f, + -0.068490773439f,-0.049840692431f,-0.032575063407f,-0.016104638577f, + 0.000000000000f,0.016104638577f,0.032575063407f,0.049840692431f, + 0.068490773439f,0.089473746717f,0.114643231034f,0.148999556899f, + 1.000000000000f, }; -// Sign array for WHT rotation (seed=42, dim=128) -__constant__ float TQ_SIGNS_128[128] = { - -1, -1, 1, -1, 1, -1, -1, 1, -1, -1, 1, 1, -1, -1, -1, -1, - -1, 1, -1, -1, 1, 1, 1, 1, 1, 1, -1, -1, -1, 1, 1, -1, - -1, -1, 1, -1, 1, -1, -1, -1, 1, -1, -1, -1, 1, -1, -1, -1, - -1, -1, 1, -1, -1, -1, 1, -1, -1, -1, -1, -1, 1, -1, 1, 1, - -1, -1, -1, 1, -1, -1, 1, 1, -1, -1, -1, -1, 1, -1, 1, -1, - 1, 1, 1, -1, 1, 1, 1, 1, 1, 1, 1, -1, -1, -1, -1, -1, - -1, -1, -1, -1, 1, 1, 1, -1, -1, -1, 1, 1, 1, -1, -1, -1, - -1, -1, -1, -1, 1, 1, -1, 1, 1, -1, -1, 1, 1, -1, 1, -1, +// ==== head_dim = 512 (seed 42) ==== +static __constant__ float TQ_SGN_512[512] = { + 1,-1,-1, 1, 1,-1, 1, 1, 1,-1,-1, 1, 1, 1,-1, 1, + -1, 1, 1,-1,-1,-1,-1, 1, 1,-1,-1,-1,-1,-1,-1, 1, + -1,-1,-1, 1,-1, 1,-1, 1,-1,-1,-1, 1,-1,-1, 1,-1, + 1,-1, 1, 1,-1, 1,-1,-1,-1, 1,-1,-1, 1,-1, 1,-1, + 1, 1,-1, 1, 1, 1, 1, 1,-1, 1, 1, 1,-1, 1, 1,-1, + -1,-1,-1,-1, 1,-1,-1,-1,-1,-1,-1, 1, 1,-1, 1,-1, + 1, 1,-1,-1, 1,-1,-1, 1,-1, 1,-1, 1,-1,-1,-1,-1, + 1, 1, 1, 1, 1, 1,-1,-1,-1,-1,-1, 1,-1,-1,-1,-1, + -1, 1,-1,-1, 1, 1, 1,-1,-1, 1,-1,-1, 1, 1, 1, 1, + -1,-1, 1, 1, 1, 1, 1,-1,-1, 1,-1,-1, 1,-1, 1,-1, + 1,-1,-1, 1,-1, 1,-1,-1, 1,-1, 1,-1,-1, 1,-1, 1, + 1,-1,-1,-1,-1,-1,-1, 1, 1, 1,-1,-1,-1, 1,-1, 1, + -1, 1, 1,-1,-1,-1, 1,-1, 1,-1, 1, 1,-1, 1,-1,-1, + -1, 1,-1, 1,-1, 1, 1,-1, 1,-1,-1, 1, 1,-1, 1,-1, + 1, 1, 1, 1,-1,-1,-1,-1, 1, 1, 1,-1,-1,-1,-1, 1, + 1, 1,-1, 1, 1, 1,-1,-1, 1,-1,-1, 1, 1, 1, 1, 1, + -1,-1,-1, 1, 1,-1,-1,-1,-1,-1, 1, 1,-1, 1,-1,-1, + 1, 1,-1, 1, 1,-1,-1, 1, 1,-1, 1,-1,-1, 1, 1, 1, + 1,-1,-1, 1, 1, 1, 1,-1,-1, 1,-1,-1,-1, 1,-1, 1, + 1, 1,-1,-1, 1, 1, 1, 1, 1,-1,-1, 1, 1, 1,-1,-1, + 1,-1, 1, 1,-1, 1, 1,-1,-1, 1, 1, 1, 1,-1, 1, 1, + 1, 1,-1,-1, 1, 1,-1, 1,-1, 1, 1,-1, 1, 1, 1, 1, + -1, 1,-1,-1, 1,-1, 1,-1,-1,-1, 1, 1, 1,-1, 1, 1, + -1, 1, 1,-1, 1,-1,-1,-1,-1, 1, 1, 1,-1, 1, 1,-1, + 1, 1, 1, 1,-1, 1, 1, 1,-1, 1, 1, 1,-1, 1,-1,-1, + -1,-1, 1,-1, 1,-1,-1,-1, 1,-1,-1,-1,-1, 1,-1, 1, + 1, 1,-1,-1,-1,-1, 1,-1, 1,-1, 1, 1, 1, 1, 1,-1, + 1,-1, 1, 1,-1,-1,-1, 1, 1, 1,-1,-1, 1,-1,-1, 1, + 1,-1, 1,-1, 1,-1, 1, 1,-1,-1,-1, 1,-1, 1, 1, 1, + -1, 1,-1, 1,-1, 1, 1,-1, 1,-1,-1,-1,-1, 1, 1, 1, + 1,-1, 1, 1,-1, 1, 1, 1,-1,-1,-1, 1,-1,-1, 1, 1, + -1, 1,-1,-1, 1,-1,-1,-1, 1,-1, 1,-1,-1, 1, 1,-1, +}; +static __constant__ float TQ_CB3_512[8] = { + -0.094875931740f,-0.059311967343f,-0.033381462097f,-0.010824348778f, + 0.010824348778f,0.033381462097f,0.059311967343f,0.094875931740f, +}; +static __constant__ float TQ_BD3_512[9] = { + -1.000000000000f,-0.077093951404f,-0.046346712857f,-0.022102905437f, + -0.000000000000f,0.022102905437f,0.046346712857f,0.077093951404f, + 1.000000000000f, +}; +static __constant__ float TQ_CB4_512[16] = { + -0.120276063681f,-0.091181337833f,-0.071355909109f,-0.055424209684f, + -0.041587848216f,-0.028990162537f,-0.017131028697f,-0.005668541417f, + 0.005668541417f,0.017131028697f,0.028990162537f,0.041587848216f, + 0.055424209684f,0.071355909109f,0.091181337833f,0.120276063681f, +}; +static __constant__ float TQ_BD4_512[17] = { + -1.000000000000f,-0.105728700757f,-0.081268623471f,-0.063390061259f, + -0.048506028950f,-0.035289004445f,-0.023060595617f,-0.011399785057f, + -0.000000000000f,0.011399785057f,0.023060595617f,0.035289004445f, + 0.048506028950f,0.063390061259f,0.081268623471f,0.105728700757f, + 1.000000000000f, }; -// ============================================================================ -// WHT in shared memory (for query rotation) -// ============================================================================ -__device__ void tq_wht_butterfly_128(float* data, int tid) { - for (int h = 1; h < 128; h *= 2) { - __syncthreads(); - for (int idx = tid; idx < 64; idx += 128) { - int block_start = (idx / h) * (h * 2); - int offset = idx % h; - int i = block_start + offset; - int j = i + h; - float a = data[i]; - float b = data[j]; - data[i] = a + b; - data[j] = a - b; - } - } - __syncthreads(); -} +// --------------------------------------------------------------------------- +// Table selection by head dimension +// --------------------------------------------------------------------------- +// `if constexpr` chains rather than a class template: __constant__ arrays decay +// to __constant__-space pointers, so this compiles to a compile-time address +// with no indirection at run time. -// Apply D·H·D rotation to a 128-dim vector in shared memory -__device__ void tq_rotate_128(float* data, int tid) { - // D · x - if (tid < 128) data[tid] *= TQ_SIGNS_128[tid]; - __syncthreads(); - // H · D · x - tq_wht_butterfly_128(data, tid); - // Scale by 1/sqrt(128) - if (tid < 128) data[tid] *= 0.08838834764831845f; - __syncthreads(); - // D · H · D · x - if (tid < 128) data[tid] *= TQ_SIGNS_128[tid]; - __syncthreads(); +template +__device__ __forceinline__ const float* tq_signs() { + if constexpr (HS == 64) return TQ_SGN_64; + else if constexpr (HS == 128) return TQ_SGN_128; + else if constexpr (HS == 256) return TQ_SGN_256; + else return TQ_SGN_512; } -// ============================================================================ -// TurboQuant Paged Attention Kernel (simplified, HEAD_SIZE=128) -// ============================================================================ - -/** - * Simplified TurboQuant attention for decode (single query token). - * - * This kernel processes one sequence per block, one head per block. - * It computes attention over compressed KV cache using codebook lookups. - * - * Parameters: - * out: [num_seqs, num_heads, HEAD_SIZE] output - * q: [num_seqs, num_heads, HEAD_SIZE] query (not rotated) - * k_packed: [total_tokens, num_kv_heads, k_packed_bytes] packed 4-bit K indices - * v_packed: [total_tokens, num_kv_heads, v_packed_bytes] packed 3-bit V indices - * k_norms: [total_tokens, num_kv_heads] half norms for K - * v_norms: [total_tokens, num_kv_heads] half norms for V - * token_table: [num_seqs, max_context_len] maps (seq, pos) → token index - * context_lens: [num_seqs] actual context length per sequence - * - * Grid: (num_heads, num_seqs, 1) - * Block: (128, 1, 1) - */ -template -__global__ void turbo_paged_attention_kernel( - float* __restrict__ out, // [num_seqs, num_heads, HEAD_SIZE] - const float* __restrict__ q, // [num_seqs, num_heads, HEAD_SIZE] - const uint8_t* __restrict__ k_packed, // [total_tokens, num_kv_heads, k_bytes_per_head] - const uint8_t* __restrict__ v_packed, // [total_tokens, num_kv_heads, v_bytes_per_head] - const __half* __restrict__ k_norms, // [total_tokens, num_kv_heads] - const __half* __restrict__ v_norms, // [total_tokens, num_kv_heads] - const int32_t* __restrict__ token_table, // [num_seqs, max_tokens_per_seq] - const int32_t* __restrict__ context_lens,// [num_seqs] - const int num_kv_heads, - const int max_tokens_per_seq, - const int num_heads, - const float scale -) { - static_assert(HEAD_SIZE == 128, "TurboQuant currently supports HEAD_SIZE=128 only"); - - const int head_idx = blockIdx.x; - const int seq_idx = blockIdx.y; - const int tid = threadIdx.x; - const int context_len = context_lens[seq_idx]; - - if (context_len == 0) return; +template +__device__ __forceinline__ const float* tq_cb4() { + if constexpr (HS == 64) return TQ_CB4_64; + else if constexpr (HS == 128) return TQ_CB4_128; + else if constexpr (HS == 256) return TQ_CB4_256; + else return TQ_CB4_512; +} - const int kv_head_idx = head_idx / (num_heads / num_kv_heads); +template +__device__ __forceinline__ const float* tq_cb3() { + if constexpr (HS == 64) return TQ_CB3_64; + else if constexpr (HS == 128) return TQ_CB3_128; + else if constexpr (HS == 256) return TQ_CB3_256; + else return TQ_CB3_512; +} - // Packed sizes per head - constexpr int K_BYTES = (K_BITS == 4) ? (HEAD_SIZE / 2) : - (K_BITS == 3) ? (((HEAD_SIZE + 9) / 10) * 4) : - (HEAD_SIZE / 4); - constexpr int V_BYTES = (V_BITS == 4) ? (HEAD_SIZE / 2) : - (V_BITS == 3) ? (((HEAD_SIZE + 9) / 10) * 4) : - (HEAD_SIZE / 4); +template +__device__ __forceinline__ const float* tq_bd4() { + if constexpr (HS == 64) return TQ_BD4_64; + else if constexpr (HS == 128) return TQ_BD4_128; + else if constexpr (HS == 256) return TQ_BD4_256; + else return TQ_BD4_512; +} - // Select codebooks - const float* k_codebook = (K_BITS == 4) ? TQ_CODEBOOK_4BIT : - (K_BITS == 3) ? TQ_CODEBOOK_3BIT : - TQ_CODEBOOK_2BIT; - const float* v_codebook = (V_BITS == 4) ? TQ_CODEBOOK_4BIT : - (V_BITS == 3) ? TQ_CODEBOOK_3BIT : - TQ_CODEBOOK_2BIT; +template +__device__ __forceinline__ const float* tq_bd3() { + if constexpr (HS == 64) return TQ_BD3_64; + else if constexpr (HS == 128) return TQ_BD3_128; + else if constexpr (HS == 256) return TQ_BD3_256; + else return TQ_BD3_512; +} - // ========================================================= - // Step 1: Load and rotate query - // ========================================================= - __shared__ float q_rotated[128]; - const float* q_ptr = q + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE; - if (tid < HEAD_SIZE) { - q_rotated[tid] = q_ptr[tid]; - } +/// 1/sqrt(HS) — the normalisation that makes the Hadamard transform, and +/// therefore all of D·H·D, its own inverse. A plain constant expression rather +/// than a constexpr function so it needs no relaxed-constexpr support. +#define TQ_WHT_SCALE(HS) ((HS) == 64 ? 0.125f \ + : (HS) == 128 ? 0.08838834764831845f \ + : (HS) == 256 ? 0.0625f \ + : 0.04419417382415922f) + +// --------------------------------------------------------------------------- +// Rotation +// --------------------------------------------------------------------------- + +/// In-place D·H·D over `HS` contiguous floats in shared memory, cooperatively +/// across `NT` threads. +/// +/// Self-inverse, so the same call rotates a query forward and rotates the V +/// accumulator back. The butterfly and the 1/sqrt(HS) scaling are ordered to +/// match `turboquant::wht::fwht_inplace` exactly (all butterflies, then one +/// scale) — reordering them would still be orthogonal but would no longer +/// agree with the cache written by the Rust compressor. +/// +/// At HS == NT == 128 this emits the same sequence of operations the original +/// hand-written `rotate128` did. +template +__device__ __forceinline__ void tq_rotate(float* d, int tid) { + const float* sgn = tq_signs(); + + for (int i = tid; i < HS; i += NT) d[i] *= sgn[i]; __syncthreads(); - // Apply D·H·D rotation to the query - tq_rotate_128(q_rotated, tid); - - // ========================================================= - // Step 2: Compute Q·K attention scores - // ========================================================= - // We process one token at a time. Each thread handles one dimension. - // This is simpler than the vectorized paged attention but correct. - - extern __shared__ char shared_mem[]; - float* logits = reinterpret_cast(shared_mem); - float* reduce_buf = reinterpret_cast(shared_mem + context_len * sizeof(float)); - - float qk_max = -FLT_MAX; - - for (int t = 0; t < context_len; t++) { - int token_idx = token_table[seq_idx * max_tokens_per_seq + t]; - const uint8_t* k_ptr = k_packed + (token_idx * num_kv_heads + kv_head_idx) * K_BYTES; - float k_norm = __half2float(k_norms[token_idx * num_kv_heads + kv_head_idx]); - - // Each thread loads and dequantizes one element of the key - float k_val = 0.0f; - if (tid < HEAD_SIZE) { - uint8_t idx; - if constexpr (K_BITS == 4) { - // Nibble packing: 2 per byte - uint8_t byte = k_ptr[tid / 2]; - idx = (tid & 1) ? ((byte >> 4) & 0xF) : (byte & 0xF); - } else if constexpr (K_BITS == 3) { - // 10-in-32 packing - int group = tid / 10; - int pos_in_group = tid % 10; - uint32_t word = *reinterpret_cast(k_ptr + group * 4); - idx = (word >> (pos_in_group * 3)) & 0x7; - } else { - // 2-bit: 4 per byte - uint8_t byte = k_ptr[tid / 4]; - int shift = 6 - (tid % 4) * 2; - idx = (byte >> shift) & 0x3; - } - k_val = k_codebook[idx]; - } - - // Dot product: q_rotated · k_centroids (all in rotated domain) - // Parallel reduction across 128 threads - float partial = (tid < HEAD_SIZE) ? q_rotated[tid] * k_val : 0.0f; - - // Warp-level reduction - for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - partial += __shfl_xor_sync(0xffffffff, partial, mask); - } - - // Cross-warp reduction via shared memory - __shared__ float warp_sums[4]; // 128 threads / 32 = 4 warps - if (tid % WARP_SIZE == 0) { - warp_sums[tid / WARP_SIZE] = partial; - } + for (int h = 1; h < HS; h *= 2) { __syncthreads(); - - float qk = 0.0f; - if (tid == 0) { - for (int w = 0; w < 4; w++) qk += warp_sums[w]; - qk = qk * k_norm * scale; - logits[t] = qk; - qk_max = fmaxf(qk_max, qk); + for (int i = tid; i < HS / 2; i += NT) { + const int bs = (i / h) * (h * 2), off = i % h; + const float a = d[bs + off], b = d[bs + off + h]; + d[bs + off] = a + b; + d[bs + off + h] = a - b; } - __syncthreads(); - } - - // Broadcast qk_max to all threads - if (tid == 0) { - reduce_buf[0] = qk_max; - } - __syncthreads(); - qk_max = reduce_buf[0]; - - // ========================================================= - // Step 3: Softmax - // ========================================================= - float exp_sum = 0.0f; - for (int t = tid; t < context_len; t += blockDim.x) { - float val = __expf(logits[t] - qk_max); - logits[t] = val; - exp_sum += val; - } - // Reduce exp_sum across threads - for (int mask = WARP_SIZE / 2; mask > 0; mask >>= 1) { - exp_sum += __shfl_xor_sync(0xffffffff, exp_sum, mask); - } - __shared__ float warp_exp_sums[4]; - if (tid % WARP_SIZE == 0) { - warp_exp_sums[tid / WARP_SIZE] = exp_sum; - } - __syncthreads(); - if (tid == 0) { - float total = 0.0f; - for (int w = 0; w < 4; w++) total += warp_exp_sums[w]; - reduce_buf[0] = __fdividef(1.0f, total + 1e-6f); } __syncthreads(); - float inv_sum = reduce_buf[0]; - for (int t = tid; t < context_len; t += blockDim.x) { - logits[t] *= inv_sum; - } + constexpr float s = TQ_WHT_SCALE(HS); + for (int i = tid; i < HS; i += NT) d[i] *= s; __syncthreads(); - // ========================================================= - // Step 4: Weighted sum of values (attention · V) - // ========================================================= - // Each thread accumulates its dimension across all tokens - float acc = 0.0f; - - if (tid < HEAD_SIZE) { - for (int t = 0; t < context_len; t++) { - float weight = logits[t]; - if (weight < 1e-8f) continue; // Skip near-zero weights - - int token_idx = token_table[seq_idx * max_tokens_per_seq + t]; - const uint8_t* v_ptr = v_packed + (token_idx * num_kv_heads + kv_head_idx) * V_BYTES; - float v_norm = __half2float(v_norms[token_idx * num_kv_heads + kv_head_idx]); - - // Dequantize value element - uint8_t idx; - if constexpr (V_BITS == 4) { - uint8_t byte = v_ptr[tid / 2]; - idx = (tid & 1) ? ((byte >> 4) & 0xF) : (byte & 0xF); - } else if constexpr (V_BITS == 3) { - int group = tid / 10; - int pos_in_group = tid % 10; - uint32_t word = *reinterpret_cast(v_ptr + group * 4); - idx = (word >> (pos_in_group * 3)) & 0x7; - } else { - uint8_t byte = v_ptr[tid / 4]; - int shift = 6 - (tid % 4) * 2; - idx = (byte >> shift) & 0x3; - } - - float v_val = v_codebook[idx] * v_norm; - acc += weight * v_val; - } - } - - // ========================================================= - // Step 5: Inverse rotation and write output - // ========================================================= - // V was compressed in the rotated domain, so the accumulated output - // is also in the rotated domain. Apply inverse rotation. - __shared__ float out_rotated[128]; - if (tid < HEAD_SIZE) { - out_rotated[tid] = acc; - } + for (int i = tid; i < HS; i += NT) d[i] *= sgn[i]; __syncthreads(); - - // D·H·D is self-inverse - tq_rotate_128(out_rotated, tid); - - // Write output - float* out_ptr = out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE; - if (tid < HEAD_SIZE) { - out_ptr[tid] = out_rotated[tid]; - } } -// ============================================================================ -// Launch wrapper -// ============================================================================ - -inline void launch_turbo_paged_attention( - float* out, - const float* q, - const uint8_t* k_packed, - const uint8_t* v_packed, - const __half* k_norms, - const __half* v_norms, - const int32_t* token_table, - const int32_t* context_lens, - int num_kv_heads, - int max_tokens_per_seq, - int num_seqs, - int num_heads, - int head_size, - float scale, - int k_bits, - int v_bits, - cudaStream_t stream -) { - // Shared memory: context_len * sizeof(float) for logits + scratch - // Use max_tokens_per_seq as upper bound - int shared_mem_size = (max_tokens_per_seq + 1) * sizeof(float); +// --------------------------------------------------------------------------- +// Scalar quantizers +// --------------------------------------------------------------------------- +// Binary search over the codebook's decision boundaries. `boundaries[0]` and +// `boundaries[n]` are the sentinels -1 and +1, so a unit-sphere coordinate can +// never fall outside. + +template +__device__ __forceinline__ uint8_t tq_q4(float x) { + const float* bd = tq_bd4(); + if (x <= bd[1]) return 0; + if (x >= bd[16]) return 15; + int lo = 1, hi = 16; + while (lo < hi) { const int m = (lo + hi) >> 1; if (x < bd[m]) hi = m; else lo = m + 1; } + return (uint8_t)(lo - 1); +} - dim3 grid(num_heads, num_seqs, 1); - dim3 block(128); // One thread per head dimension +template +__device__ __forceinline__ uint8_t tq_q3(float x) { + const float* bd = tq_bd3(); + if (x <= bd[1]) return 0; + if (x >= bd[8]) return 7; + int lo = 1, hi = 8; + while (lo < hi) { const int m = (lo + hi) >> 1; if (x < bd[m]) hi = m; else lo = m + 1; } + return (uint8_t)(lo - 1); +} - // Dispatch based on bit-width combination - if (head_size == 128) { - if (k_bits == 4 && v_bits == 3) { - turbo_paged_attention_kernel<128, 4, 3><<>>( - out, q, k_packed, v_packed, k_norms, v_norms, - token_table, context_lens, num_kv_heads, max_tokens_per_seq, - num_heads, scale); - } else if (k_bits == 3 && v_bits == 3) { - turbo_paged_attention_kernel<128, 3, 3><<>>( - out, q, k_packed, v_packed, k_norms, v_norms, - token_table, context_lens, num_kv_heads, max_tokens_per_seq, - num_heads, scale); - } else if (k_bits == 3 && v_bits == 2) { - turbo_paged_attention_kernel<128, 3, 2><<>>( - out, q, k_packed, v_packed, k_norms, v_norms, - token_table, context_lens, num_kv_heads, max_tokens_per_seq, - num_heads, scale); - } +// --------------------------------------------------------------------------- +// Packed-K gather +// --------------------------------------------------------------------------- + +/// Bytes of packed 4-bit K each lane owns: (HS/2 bytes per head) / 32 lanes. +/// 1 at HS=64, 2 at 128, 4 at 256, 8 at 512. +#define TQ_BYTES_PER_LANE(HS) ((HS) / 64) + +/// Load a lane's whole packed-K run as one vector access. +/// +/// The run is contiguous and naturally aligned: the cache's x=16 swizzle keeps +/// each group of 16 bytes contiguous, and a lane's run starts at `lane*BPL` +/// with BPL a power of two dividing 16, so it never straddles a group. +template +__device__ __forceinline__ uint64_t tq_load_klane(const uint8_t* p) { + if constexpr (BPL == 1) { + return (uint64_t)(*p); + } else if constexpr (BPL == 2) { + return (uint64_t)(*reinterpret_cast(p)); + } else if constexpr (BPL == 4) { + return (uint64_t)(*reinterpret_cast(p)); + } else { + const uint2 v = *reinterpret_cast(p); + return ((uint64_t)v.y << 32) | (uint64_t)v.x; } } diff --git a/mistralrs-quant/examples/emit_turboquant_cuda_tables.rs b/mistralrs-quant/examples/emit_turboquant_cuda_tables.rs new file mode 100644 index 000000000..6f3e53332 --- /dev/null +++ b/mistralrs-quant/examples/emit_turboquant_cuda_tables.rs @@ -0,0 +1,70 @@ +//! Emit the CUDA `__constant__` tables for TurboQuant at a given head dimension. +//! +//! The CUDA kernels cannot call `generate_signs`/`cached_generated_codebook` at +//! runtime, so the tables are checked in as `__constant__` arrays. This example +//! is the *only* sanctioned way to produce them: it uses the same crate code the +//! Rust compressor uses, so the emitted tables cannot drift from the format by +//! construction. +//! +//! `turboquant_cuda_tables_match_crate` in `mistralrs-quant/src/turboquant/cuda_tables.rs` +//! re-derives these values at test time and fails if the checked-in CUDA source +//! disagrees, so a hand-edit of the `.cu` is caught by CI. +//! +//! Usage: `cargo run -p mistralrs-quant --example emit_turboquant_cuda_tables -- 512` + +// `get_codebook` is the same entry point `TurboQuantLayout::new` uses, so the +// emitted table is the one the Rust compressor actually quantizes against — +// static for 64/128/256, numerically generated elsewhere. +use mistralrs_quant::turboquant::codebook::get_codebook; +use mistralrs_quant::turboquant::wht::generate_signs; + +const SEED: u64 = 42; + +fn emit_signs(dim: usize) { + let signs = generate_signs(SEED, dim); + println!("__constant__ float TQ_SGN_{dim}[{dim}] = {{"); + for row in signs.chunks(16) { + let cells: Vec = row + .iter() + .map(|s| format!("{:>2}", if *s > 0.0 { "1" } else { "-1" })) + .collect(); + println!(" {},", cells.join(",")); + } + println!("}};"); +} + +fn emit_codebook(dim: usize, bits: u32) { + let cb = get_codebook(dim, bits); + let n = cb.centroids.len(); + println!("__constant__ float TQ_CB{bits}_{dim}[{n}] = {{"); + for row in cb.centroids.chunks(4) { + let cells: Vec = row.iter().map(|c| format!("{c:.12}f")).collect(); + println!(" {},", cells.join(",")); + } + println!("}};"); + + let nb = cb.boundaries.len(); + println!("__constant__ float TQ_BD{bits}_{dim}[{nb}] = {{"); + for row in cb.boundaries.chunks(4) { + let cells: Vec = row.iter().map(|c| format!("{c:.12}f")).collect(); + println!(" {},", cells.join(",")); + } + println!("}};"); +} + +fn main() { + let dims: Vec = std::env::args() + .skip(1) + .map(|a| a.parse().expect("head dim must be an integer")) + .collect(); + let dims = if dims.is_empty() { vec![512] } else { dims }; + + for dim in dims { + println!("// ==== head_dim = {dim} (seed {SEED}) ===="); + emit_signs(dim); + for bits in [3u32, 4u32] { + emit_codebook(dim, bits); + } + println!(); + } +} diff --git a/mistralrs-quant/examples/qtip_lut_grouped_bench.rs b/mistralrs-quant/examples/qtip_lut_grouped_bench.rs index 98f710cb7..59f430db9 100644 --- a/mistralrs-quant/examples/qtip_lut_grouped_bench.rs +++ b/mistralrs-quant/examples/qtip_lut_grouped_bench.rs @@ -32,7 +32,8 @@ use candle_core::{DType, Device, Result, Tensor}; use mistralrs_quant::{ - QtipCodebook, QtipLayer, QtipSearchDetail, QtipSearchStamp, QuantMethod, QTIP_GROUPED_TILE_K, + QtipCodebook, QtipGeometry, QtipLayer, QtipSearchDetail, QtipSearchStamp, QuantMethod, + QTIP_GROUPED_TILE_K, }; use std::time::Instant; @@ -138,6 +139,9 @@ fn build_layer( QtipSearchStamp::Unstamped, QtipSearchDetail::Unknown, QtipCodebook::Gaussian, + // This bench synthesises a `[65536, 2]` F32 table above; that is the + // K=4/V=2/L=16 geometry by construction. + QtipGeometry::K4V2L16, ) } diff --git a/mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu b/mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu index dad9e8330..077b7c473 100644 --- a/mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu +++ b/mistralrs-quant/kernels/blockwise_fp8/blockwise_fp8_gemm.cu @@ -40,17 +40,54 @@ __device__ __forceinline__ float get_scale(const float *__restrict__ scale, return __ldg(&scale[sr * scale_stride + sc]); } +// Shift variant of `get_scale`, for power-of-two scale-block sizes. +// +// `block_size_y`/`block_size_x` arrive as kernel ARGUMENTS, so `n / +// block_size_y` above is `int / int` by a non-constant divisor. NVIDIA GPUs +// have no integer-divide instruction; nvcc expands each division inline to a +// reciprocal-multiply-and-fixup sequence of ~15-20 instructions, and +// `fp8_matmul_tiled` pays TWO of them per weight element it stages into +// shared memory. The shipped `weight_block_size` is [128, 128] — both powers +// of two — so the host launcher computes `shift = log2(block_size)` once per +// launch and the kernel indexes with two SHFs instead. +// +// Bit-identical: for non-negative `n` and a power-of-two divisor, +// `n >> log2(d)` IS `n / d` — the same quotient for every input, no rounding +// involved anywhere. The two paths read the same scale word; only the +// instruction count changes. Non-power-of-two geometries keep the division +// path above. +__device__ __forceinline__ float get_scale_pow2(const float *__restrict__ scale, + int n, int k, int scale_stride, + int scale_shift_y, + int scale_shift_x) { + int sr = n >> scale_shift_y; + int sc = k >> scale_shift_x; + return __ldg(&scale[sr * scale_stride + sc]); +} + // ============================================================================ // FP8 Matmul Kernel // ============================================================================ -template +// `POW2_SCALE` selects, at compile time, how the per-element weight scale is +// indexed while staging the weight tile: +// * true -> `get_scale_pow2`: two SHFs, using `scale_shift_y`/`scale_shift_x` +// (host-computed `log2(block_size_y)`/`log2(block_size_x)`). +// * false -> `get_scale`: two runtime integer divisions (the general path, +// kept for non-power-of-two scale geometries). +// The two instantiations produce BIT-IDENTICAL output (`n >> log2(d) == n / d` +// for the non-negative indices used here), so dispatching on it can never +// change what a leg of an A/B measures — only how fast the scalar arm runs. +// Both `block_size_*` and `scale_shift_*` are always passed; each +// instantiation reads only its own pair. +template __global__ void fp8_matmul_tiled(const T *__restrict__ input, const __nv_fp8_e4m3 *__restrict__ weight, const float *__restrict__ weight_scale, T *__restrict__ output, int M, int N, int K, int scale_row_stride, int block_size_y, - int block_size_x) { + int block_size_x, int scale_shift_y, + int scale_shift_x) { // Padding is +1, not +4. The inner product reads `s_weight[tx][k]` with `tx` // varying fastest within a warp, so consecutive lanes are `BLOCK_K + pad` // floats apart. At BLOCK_K=32 a +4 pad gives stride 36; 36 mod 32 = 4, so @@ -106,8 +143,14 @@ __global__ void fp8_matmul_tiled(const T *__restrict__ input, if (gn < N && gk < K) { __nv_fp8_e4m3 w; w.__x = __ldg(reinterpret_cast(&weight[gn * K + gk])); - float s = get_scale(weight_scale, gn, gk, scale_row_stride, - block_size_y, block_size_x); + float s; + if constexpr (POW2_SCALE) { + s = get_scale_pow2(weight_scale, gn, gk, scale_row_stride, + scale_shift_y, scale_shift_x); + } else { + s = get_scale(weight_scale, gn, gk, scale_row_stride, block_size_y, + block_size_x); + } val = fp8_to_float(w) * s; } s_weight[ln][lk] = val; @@ -625,41 +668,66 @@ __global__ void fp8_moe_gemm(const T *__restrict__ input, // C API // ============================================================================ -extern "C" void launch_fp8_matmul_f16(const __half *input, - const __nv_fp8_e4m3 *weight, - const float *weight_scale, __half *output, - int M, int N, int K, int scale_row_stride, - int block_size_y, int block_size_x, - cudaStream_t stream) { +// Host-side pow2 test + exact log2 for the tiled kernel's scale fast path. +// Computed ONCE per launch, on the host, so the device never divides. +static inline bool fp8_is_pow2(int v) { return v > 0 && (v & (v - 1)) == 0; } +static inline int fp8_ilog2_exact(int v) { + int s = 0; + while ((1 << s) < v) + ++s; + return s; // caller guarantees fp8_is_pow2(v), so (1 << s) == v +} + +// Shared launcher: picks the POW2_SCALE instantiation when BOTH scale-block +// sizes are powers of two (the shipped geometry is [128, 128], so serving +// always takes it), otherwise the general division path. The two +// instantiations are bit-identical in output — see the kernel comment. +template +static void launch_fp8_matmul_impl(const T *input, const __nv_fp8_e4m3 *weight, + const float *weight_scale, T *output, int M, + int N, int K, int scale_row_stride, + int block_size_y, int block_size_x, + cudaStream_t stream) { constexpr int TILE = 32; constexpr int TILE_K = 32; dim3 block(TILE, TILE); dim3 grid(CEILDIV(N, TILE), CEILDIV(M, TILE)); - fp8_gemm::fp8_matmul_tiled - <<>>(input, weight, weight_scale, output, M, N, K, - scale_row_stride, block_size_y, - block_size_x); + if (fp8_is_pow2(block_size_y) && fp8_is_pow2(block_size_x)) { + fp8_gemm::fp8_matmul_tiled + <<>>(input, weight, weight_scale, output, M, N, + K, scale_row_stride, block_size_y, + block_size_x, fp8_ilog2_exact(block_size_y), + fp8_ilog2_exact(block_size_x)); + } else { + fp8_gemm::fp8_matmul_tiled + <<>>(input, weight, weight_scale, output, M, N, + K, scale_row_stride, block_size_y, + block_size_x, 0, 0); + } CUDA_CHECK(cudaGetLastError()); } +extern "C" void launch_fp8_matmul_f16(const __half *input, + const __nv_fp8_e4m3 *weight, + const float *weight_scale, __half *output, + int M, int N, int K, int scale_row_stride, + int block_size_y, int block_size_x, + cudaStream_t stream) { + launch_fp8_matmul_impl(input, weight, weight_scale, output, M, N, K, + scale_row_stride, block_size_y, block_size_x, + stream); +} + extern "C" void launch_fp8_matmul_bf16(const __nv_bfloat16 *input, const __nv_fp8_e4m3 *weight, const float *weight_scale, __nv_bfloat16 *output, int M, int N, int K, int scale_row_stride, int block_size_y, int block_size_x, cudaStream_t stream) { - constexpr int TILE = 32; - constexpr int TILE_K = 32; - - dim3 block(TILE, TILE); - dim3 grid(CEILDIV(N, TILE), CEILDIV(M, TILE)); - - fp8_gemm::fp8_matmul_tiled<__nv_bfloat16, TILE, TILE, TILE_K> - <<>>(input, weight, weight_scale, output, M, N, K, - scale_row_stride, block_size_y, - block_size_x); - CUDA_CHECK(cudaGetLastError()); + launch_fp8_matmul_impl<__nv_bfloat16>(input, weight, weight_scale, output, M, + N, K, scale_row_stride, block_size_y, + block_size_x, stream); } // Warp-per-row GEMV launcher for the decode regime (M <= 4). Row-grouping diff --git a/mistralrs-quant/src/blockwise_fp8/gemv_wide.rs b/mistralrs-quant/src/blockwise_fp8/gemv_wide.rs index e173f4f01..0f38c182f 100644 --- a/mistralrs-quant/src/blockwise_fp8/gemv_wide.rs +++ b/mistralrs-quant/src/blockwise_fp8/gemv_wide.rs @@ -29,14 +29,17 @@ //! contraction — contraction is identical in both kernels and so cancels out //! of the comparison. //! -//! # Why the gate is opt-in +//! # Why the gate is default-ON //! -//! 🔴 The kernel is UNVERIFIED ON HARDWARE. [`wide_enabled`] reads -//! `ARC_FP8_GEMV_WIDE` and is true **only** for the exact value `"1"`. -//! `ARC_FP8_GEMV_WIDE=0`, `=off`, `=true` and unset all leave it OFF. That -//! strictness is deliberate: this crate already ships `ARC_NO_FP8_GEMV`, read -//! with `std::env::var(..).is_ok()` (`blockwise_fp8/ops.rs`), where `=0` -//! *enables* the disable — the polarity bug this form cannot have. +//! ✅ MEASURED ON HARDWARE 2026-08-22, box `arc-s10-ledger` (H200), candidate +//! `4d03b9e25`, session-10 leg 3: b=1 decode 43.80 tok/s vs 40.76 baseline +//! (+7.5%, 3×512 produced tokens, seeded t=0.7 p=0.95), engagement proven by +//! `[arc-fp8-dispatch] path=gemv_wide`, coherence canary correct (numerics +//! shift the sampled continuation within the stated f32 bound; canary1 was +//! byte-identical, canary2 coherent and arithmetically right). Per the flip +//! doctrine the default turned ON in the same session as the measurement. +//! `ARC_FP8_GEMV_WIDE=0` — and only the exact string `"0"` — restores the +//! warp GEMV from the same binary. Value-read, never presence-read (#212). // This module is exactly two things: a gate that only a CUDA build calls, and // a CPU model of the kernel that only `mod tests` calls. On a non-CUDA build @@ -56,16 +59,19 @@ pub(crate) const WARP_K_PER_ITER: usize = 128; /// K-elements one lane of `fp8_gemv_wide` owns per iteration. pub(crate) const WIDE_K_PER_LANE: usize = 16; -/// `ARC_FP8_GEMV_WIDE` — opt-in switch for the wide b=1 GEMV. +/// `ARC_FP8_GEMV_WIDE` — kill switch for the wide b=1 GEMV (default ON). /// /// Split from the env read so the polarity is unit-testable. Only the exact -/// string `"1"` enables. +/// string `"0"` disables; everything else (unset included) leaves the +/// measured-faster wide kernel on. Measured: s10 leg 3, `arc-s10-ledger`, +/// `4d03b9e25`, 43.80 vs 40.76 tok/s b=1. pub(crate) fn wide_enabled_from(value: Option<&str>) -> bool { - value == Some("1") + value != Some("0") } -/// `ARC_FP8_GEMV_WIDE=1` enables `fp8_gemv_wide` for eligible b=1 shapes. -/// Default OFF. Read once; the decode hot loop never touches the env. +/// `ARC_FP8_GEMV_WIDE=0` disables `fp8_gemv_wide` for eligible b=1 shapes. +/// Default ON (measured +7.5% b=1). Read once; the decode hot loop never +/// touches the env. #[cfg(feature = "cuda")] pub(crate) fn wide_enabled() -> bool { use std::sync::OnceLock; @@ -636,21 +642,20 @@ mod tests { ); } - /// Env polarity: only the exact string "1" turns the kernel on. Guards + /// Env polarity after the s10 default flip: only the exact string "0" + /// turns the kernel OFF. Still value-read, never presence-read — guards /// against the `ARC_NO_DEDICATED_DECODE` failure, where `var_os(..) /// .is_some()` makes `=0` *enable* the flag. #[test] - fn only_literal_one_enables() { - assert!(wide_enabled_from(Some("1"))); - for v in [ - "0", "", "on", "off", "true", "false", "yes", "no", "2", " 1", "1 ", - ] { + fn only_literal_zero_disables() { + assert!(!wide_enabled_from(Some("0")), "`ARC_FP8_GEMV_WIDE=0` must disable"); + for v in ["1", "", "on", "off", "true", "false", "yes", "no", "2", " 0", "0 "] { assert!( - !wide_enabled_from(Some(v)), - "`ARC_FP8_GEMV_WIDE={v}` must not enable" + wide_enabled_from(Some(v)), + "`ARC_FP8_GEMV_WIDE={v}` must leave the measured default ON" ); } - assert!(!wide_enabled_from(None), "unset must not enable"); + assert!(wide_enabled_from(None), "unset must stay ON (the measured default)"); } /// The gate must reject every shape whose preconditions the kernel relies diff --git a/mistralrs-quant/src/blockwise_fp8/mod.rs b/mistralrs-quant/src/blockwise_fp8/mod.rs index 2cf3b9fa5..897f1e068 100644 --- a/mistralrs-quant/src/blockwise_fp8/mod.rs +++ b/mistralrs-quant/src/blockwise_fp8/mod.rs @@ -95,6 +95,20 @@ impl QuantMethod for BlockwiseFP8Linear { // A/B runs from a single binary. let m_rows: usize = x.dims().iter().rev().skip(1).product(); let prefer_cublas = m_rows >= arc_fp8_cublas_min_m(); + if prefer_cublas + && matches!(x.device(), candle_core::Device::Cuda(_)) + && ffi::HAVE_BLOCKWISE_GEMM_KERNELS + { + // Engagement line for the threshold sweep: this forward was + // DIVERTED from the native FP8 kernels to dequantize+cuBLASLt + // by `ARC_FP8_CUBLAS_MIN_M`. `ARC_LOG_FP8_DISPATCH=1` only. + ops::log_fp8_dispatch( + ops::Fp8DispatchPath::DequantCublaslt, + m_rows, + self.weight.dims().first().copied().unwrap_or(0), + x.dims().last().copied().unwrap_or(0), + ); + } if !prefer_cublas && matches!(x.device(), candle_core::Device::Cuda(_)) && ffi::HAVE_BLOCKWISE_GEMM_KERNELS @@ -143,6 +157,34 @@ impl QuantMethod for BlockwiseFP8Linear { weight, self.bias.clone(), )))?; + + // 🔴 FLATTEN TO 2-D FIRST — the native branch above already does this at + // the top of `forward`, and this branch did not. + // + // V4's activation is rank-3 `[B, T, hidden]`. Passed through untouched, + // `UnquantLinear::forward` takes `w.broadcast_left(B)` and, because bias + // is `None` on every V4 linear, dispatches `cublaslt.batch_matmul` with + // `stride_b = 0`. At decode T=1, so that is a **B-way batched GEMM with + // m=1** — a batched GEMV wearing a GEMM's name, which is the one shape + // cuBLASLt has no advantage in. The documented 27x was measured at + // prefill, where `[1, 512, hidden]` collapses to batch=1 and the call is + // a genuine `[512, hidden] x [hidden, N]` GEMM. + // + // Without this, lowering `ARC_FP8_CUBLAS_MIN_M` measures the batched-GEMV + // path and concludes "cuBLASLt loses at decode" — a wrong conclusion + // drawn from a shape the threshold change was never meant to select. + // Flattening makes the two branches comparable, which is the whole + // premise of the A/B. + let orig_dims = x.dims().to_vec(); + if orig_dims.len() > 2 { + let features = orig_dims[orig_dims.len() - 1]; + let rows: usize = orig_dims[..orig_dims.len() - 1].iter().product(); + let out = unquant.forward(&x.reshape((rows, features))?)?; + let out_features = out.dim(out.rank() - 1)?; + let mut new_dims = orig_dims[..orig_dims.len() - 1].to_vec(); + new_dims.push(out_features); + return out.reshape(new_dims); + } unquant.forward(x) } @@ -644,10 +686,30 @@ pub fn blockwise_fp8_linear_b( /// **The measurement that settles this constant:** re-run the M sweep with the /// flatten fix present, three arms, on one binary — /// `ARC_FP8_CUBLAS_MIN_M`/`ARC_NO_FP8_WMMA` toggle all three with no rebuild — -/// and set the default to the first M at which (a) beats `min((b), (c))`. Until -/// then the default stays where it is, because moving it would be a *second* -/// constant with no measurement behind it, which is the exact defect this doc -/// exists to record. +/// and set the default to the first M at which (a) beats `min((b), (c))`. +/// +/// # ✅ THAT SWEEP HAS NOW RUN — s10, and 64 is the measured crossover +/// +/// 2026-08-22, box `arc-s10-ledger` (H200), candidate `4d03b9e25`, V4-Flash, +/// `arc-tools/fp8_threshold_sweep.sh`, LADDER {8, 64, 256}, three arms, one +/// binary, flatten fix present, >=1000-token floor per row (M=8 rows VOID on +/// the floor and make no claim): +/// +/// ```text +/// M arm tokens wall_s tok/s status +/// 64 tiled 2811 12.6 223.1 CLEAN +/// 64 wmma 2811 8.6 328.2 CLEAN +/// 64 cublaslt 3072 5.6 552.4 CLEAN <- (a) beats min((b),(c)) +/// 256 tiled 12288 43.0 285.8 CLEAN +/// 256 wmma 11244 18.1 620.4 CLEAN +/// 256 cublaslt 12288 17.2 715.9 CLEAN <- and again +/// ``` +/// +/// First clean M where (a) beats min((b), (c)): **64**. The default moves to +/// that measured point. M=8 stays with the native arms (its rows were void, so +/// no claim is made about 5..63 — the conservative side of a void row is to +/// leave it with the incumbent). `fp8_gemv_warp` still owns M <= 4; the M=1 +/// floor row from the earlier sweep (cuBLASLt 0.72x at M=1) still stands. /// /// `fp8_gemv_warp` still owns `M <= 4` and is not affected by this constant. /// @@ -656,11 +718,11 @@ pub fn blockwise_fp8_linear_b( #[cfg(feature = "cuda")] fn arc_fp8_cublas_min_m() -> usize { use std::sync::OnceLock; - // Still 512. NOT because 512 is right -- the doc above shows it is a single - // prefill point -- but because the only sweep in hand measured the scalar - // kernel that #200 displaced, and lowering this to 5 on that evidence would - // make the tensor-core GEMM dead code at every M >= 5. Re-sweep three arms. - const DEFAULT_MIN_M: usize = 512; + // 64 = the first CLEAN M at which dequantize+cuBLASLt beats min(WMMA, + // tiled) in the s10 three-arm sweep (arc-s10-ledger, 4d03b9e25; rows in + // the doc above). M=8 rows were VOID on the token floor, so 5..63 stays + // with the native arms rather than being claimed by interpolation. + const DEFAULT_MIN_M: usize = 64; static CACHE: OnceLock = OnceLock::new(); *CACHE.get_or_init(|| { std::env::var("ARC_FP8_CUBLAS_MIN_M") diff --git a/mistralrs-quant/src/blockwise_fp8/ops.rs b/mistralrs-quant/src/blockwise_fp8/ops.rs index 75a1d1f42..347eb8288 100644 --- a/mistralrs-quant/src/blockwise_fp8/ops.rs +++ b/mistralrs-quant/src/blockwise_fp8/ops.rs @@ -1008,8 +1008,7 @@ fn fp8_wmma_enabled() -> bool { // Under the presence test this replaced, it silently selected the control // arm — which is the arm this kernel's own docs tell you to run first, so // an operator following them would have recorded the control twice. - static ENABLED: LazyLock = - LazyLock::new(|| !crate::env_flag_is_set("ARC_NO_FP8_WMMA")); + static ENABLED: LazyLock = LazyLock::new(|| !crate::env_flag_is_set("ARC_NO_FP8_WMMA")); *ENABLED } @@ -1070,11 +1069,12 @@ fn fp8_wmma_eligible(block_size_y: i32, block_size_x: i32) -> bool { /// (see [`gemv_wide::wide_eligible`]); any shape that fails one keeps /// `fp8_gemv_warp`, unchanged, with no error. /// -/// 🔴 Default OFF. `ARC_FP8_GEMV_WIDE=1` — and *only* the literal `"1"` — -/// turns it on. The kernel has never run on a GPU and is not bit-identical to -/// `fp8_gemv_warp` (it re-associates the f32 accumulation; the exact claim and -/// its CPU proof are in `blockwise_fp8/gemv_wide.rs`). It must be A/B'd on -/// hardware against the default before anyone considers flipping it. +/// ✅ Default ON since the s10 hardware A/B (box `arc-s10-ledger`, candidate +/// `4d03b9e25`, leg 3: 43.80 vs 40.76 tok/s b=1, engagement line +/// `path=gemv_wide`, coherent canary). `ARC_FP8_GEMV_WIDE=0` — and *only* the +/// literal `"0"` — restores `fp8_gemv_warp` from the same binary. It is not +/// bit-identical to `fp8_gemv_warp` (it re-associates the f32 accumulation; +/// the exact claim and its CPU proof are in `blockwise_fp8/gemv_wide.rs`). /// /// The pointers passed here are the ones the kernel will receive, so the /// 16-byte alignment a `uint4` load requires is checked on the real address @@ -1099,6 +1099,81 @@ fn wide_gemv_shift( ) } +/// Parent system: ArcLab. +/// +/// Which blockwise-FP8 kernel a forward actually selected. Used only by +/// [`log_fp8_dispatch`]; the discriminants index its once-per-process latch. +#[cfg(feature = "cuda")] +#[derive(Clone, Copy)] +pub(crate) enum Fp8DispatchPath { + GemvWide = 0, + GemvWarp = 1, + Wmma = 2, + Tiled = 3, + DequantCublaslt = 4, +} + +#[cfg(feature = "cuda")] +impl Fp8DispatchPath { + fn name(self) -> &'static str { + match self { + Fp8DispatchPath::GemvWide => "gemv_wide", + Fp8DispatchPath::GemvWarp => "gemv_warp", + Fp8DispatchPath::Wmma => "wmma", + Fp8DispatchPath::Tiled => "tiled", + Fp8DispatchPath::DequantCublaslt => "dequant_cublaslt", + } + } +} + +/// `ARC_LOG_FP8_DISPATCH=1` — engagement logging for the blockwise-FP8 +/// dispatch. Read by VALUE via `env_flag_is_set` (`=0` keeps it off), latched +/// once per process so the decode hot loop never touches the env. +/// +/// Registered in `mistralrs-core/tests/capability_reachability.rs`. +#[cfg(feature = "cuda")] +fn fp8_dispatch_log_enabled() -> bool { + use std::sync::LazyLock; + static ENABLED: LazyLock = + LazyLock::new(|| crate::env_flag_is_set("ARC_LOG_FP8_DISPATCH")); + *ENABLED +} + +/// Parent system: ArcLab. +/// +/// Print ONE line per dispatch path per process naming the kernel that +/// actually ran — the engagement assertion every threshold sweep must grep +/// for. A sweep leg that cannot prove WHICH kernel served it is measuring an +/// assumption (`ARC_NO_DEDICATED_DECODE` produced exactly that null result: +/// both arms of two A/Bs ran the same code and nobody could tell). +/// +/// Cost when `ARC_LOG_FP8_DISPATCH` is unset (the default): one latched bool +/// read. When set: one relaxed atomic swap per call — no locks, no +/// allocation — so enabling it does not perturb the rates the sweep records. +/// The shape printed is the FIRST shape that reached the path, which is +/// enough to identify the projection; per-call logging would be ~430 lines +/// per V4 step at B=256 and is deliberately not offered. +#[cfg(feature = "cuda")] +pub(crate) fn log_fp8_dispatch(path: Fp8DispatchPath, m: usize, n: usize, k: usize) { + use std::sync::atomic::{AtomicBool, Ordering}; + static LOGGED: [AtomicBool; 5] = [ + AtomicBool::new(false), + AtomicBool::new(false), + AtomicBool::new(false), + AtomicBool::new(false), + AtomicBool::new(false), + ]; + if !fp8_dispatch_log_enabled() { + return; + } + if !LOGGED[path as usize].swap(true, Ordering::Relaxed) { + eprintln!( + "[arc-fp8-dispatch] path={} first_shape=m{m}_n{n}_k{k}", + path.name() + ); + } +} + /// FP8 blockwise matmul. /// Computes output = input @ weight.T where weight is FP8 blockwise quantized. /// - input: [M, K] in fp16/bf16 @@ -1221,8 +1296,7 @@ fn fp8_blockwise_matmul_impl( match input.dtype() { DType::F16 => { - let output = - crate::arc_outbuf::alloc_out_fully_written::(&dev, (m * n) as usize)?; + let output = crate::arc_outbuf::alloc_out_fully_written::(&dev, (m * n) as usize)?; let input_s = match &*input_storage { Storage::Cuda(cuda_storage) => cuda_storage.as_cuda_slice::()?, @@ -1234,8 +1308,23 @@ fn fp8_blockwise_matmul_impl( let (input_ptr, _input_guard) = slice_ptr(input_s, input_l.start_offset()); let wide_shift = wide_gemv_shift(use_gemv, k, block_size_x, weight_ptr, input_ptr); + // Mirrors the if/else chain below exactly; keep the two in step. + log_fp8_dispatch( + if wide_shift.is_some() { + Fp8DispatchPath::GemvWide + } else if use_gemv { + Fp8DispatchPath::GemvWarp + } else if use_wmma { + Fp8DispatchPath::Wmma + } else { + Fp8DispatchPath::Tiled + }, + m as usize, + n as usize, + k as usize, + ); if let Some(scale_shift) = wide_shift { - // 🔴 UNVERIFIED ON HARDWARE. `ARC_FP8_GEMV_WIDE=1` only. + // ✅ Measured s10 leg 3 (arc-s10-ledger, 4d03b9e25): +7.5% b=1. Default ON; `ARC_FP8_GEMV_WIDE=0` restores warp. unsafe { ffi::launch_fp8_gemv_wide_f16( input_ptr as *const _, @@ -1323,8 +1412,23 @@ fn fp8_blockwise_matmul_impl( let (input_ptr, _input_guard) = slice_ptr(input_s, input_l.start_offset()); let wide_shift = wide_gemv_shift(use_gemv, k, block_size_x, weight_ptr, input_ptr); + // Mirrors the if/else chain below exactly; keep the two in step. + log_fp8_dispatch( + if wide_shift.is_some() { + Fp8DispatchPath::GemvWide + } else if use_gemv { + Fp8DispatchPath::GemvWarp + } else if use_wmma { + Fp8DispatchPath::Wmma + } else { + Fp8DispatchPath::Tiled + }, + m as usize, + n as usize, + k as usize, + ); if let Some(scale_shift) = wide_shift { - // 🔴 UNVERIFIED ON HARDWARE. `ARC_FP8_GEMV_WIDE=1` only. + // ✅ Measured s10 leg 3 (arc-s10-ledger, 4d03b9e25): +7.5% b=1. Default ON; `ARC_FP8_GEMV_WIDE=0` restores warp. unsafe { ffi::launch_fp8_gemv_wide_bf16( input_ptr as *const _, diff --git a/mistralrs-quant/src/lib.rs b/mistralrs-quant/src/lib.rs index 2ec9ae67f..b489f148a 100644 --- a/mistralrs-quant/src/lib.rs +++ b/mistralrs-quant/src/lib.rs @@ -120,8 +120,8 @@ pub use qtip::tune::{ pub use qtip::{ bake_cache, gpu_quantize_cpu_fallback_count, grouped_launch_counts, grouped_variant, hessian_row_weights, qtip_expected_distinct_experts, qtip_expected_pairs_per_distinct_expert, - qtip_grouped_gemm_tile_fill, set_grouped_variant, viterbi_quantize_row, - BakeCacheError, BakeKey, ExpertBpwTable, Qtip2bLayer, QtipBakeConfig, QtipCodebook, QtipLayer, + qtip_grouped_gemm_tile_fill, set_grouped_variant, viterbi_quantize_row, BakeCacheError, + BakeKey, ExpertBpwTable, Qtip2bLayer, QtipBakeConfig, QtipCodebook, QtipGeometry, QtipLayer, QtipMode, QtipPackedView, QtipRotation, QtipSearchDetail, QtipSearchStamp, TrellisBpw, TrellisSearch, QTIP2B_MCG_MULT, QTIP_GATHER_GEMV_MAX_PAIRS, QTIP_GROUPED_TILE_K, QTIP_GROUPED_TILE_M, QTIP_GROUPED_TILE_N, QTIP_GROUPED_VARIANT_BASELINE, diff --git a/mistralrs-quant/src/qtip/bitshift.rs b/mistralrs-quant/src/qtip/bitshift.rs index b914c5e6e..308ee1428 100644 --- a/mistralrs-quant/src/qtip/bitshift.rs +++ b/mistralrs-quant/src/qtip/bitshift.rs @@ -1902,37 +1902,53 @@ impl QuantMethod for Qtip2bLayer { // On-device only, propagate errors (the CPU fallback would // D2H-sync under graph capture and read garbage indices). // - // Above this boundary the rung goes to the trellis grouped GEMM - // below — tokens sorted by expert on-device, each expert's bytes - // read once per m-tile — which is the amortizing kernel the - // fleet math wants. But it only amortizes when its m-tiles are - // FULL: it stages a woken expert's bytes once per - // `GROUPED_TILE_M` pairs, so at V4's top-6-of-256 routing a - // 9-token step wakes ~54 experts for 54 pairs and every 16-row - // tile carries one useful row. The gate is therefore the tile - // fill, not a token count — same predicate, same reasons, as - // the LUT rung; see `super::gather_policy` §4 for the - // derivation and for the 5.7x end-to-end regression the raw - // token gate caused there. + // 🔴 TILE FILL WAS THE WRONG QUANTITY — RETRACTED ON MEASUREMENT. + // + // This gate used to also require `grouped_gemm_tiles_amortize`: + // switch to the grouped GEMM only once each woken expert draws + // `GROUPED_TILE_M` (=16) pairs. At V4's top-6-of-256 that needs + // `6n/256 >= 16`, i.e. **n >= 683** — a token count decode never + // reaches. The clause therefore pinned every decode step, and + // every batch up to 682, on the gather-GEMV arm. + // + // Two independent measurements retired it: + // * decode, H200, V4-Flash qtip2b, clean rows only (>=97% + // achieved concurrency, <7% derived-vs-measured spread): + // grouped beats GEMV at B=48 (1.10x), 64 (1.14x), 80 + // (1.22x), 96 (1.29x), 128 (1.66x), 512 (1.51x). + // * prefill, same gate over prompt lengths: +16% at 16 + // tokens, +27% at 32, +50% at 64, +78% at 128, +181% at 512. + // The gate's own justifying A/B ("1.00x at N=128 and N=512") + // does not reproduce. + // + // **A 7.5%-full tile still wins by 16%.** So fill is not the + // deciding quantity. The win is that the grouped kernel stages a + // woken expert's bytes ONCE PER M-TILE instead of once per + // (token, expert) pair — the GEMV arm's reads are exactly linear + // in `n_tokens`, measured at 12.00x redundancy at n=512, while + // the grouped arm holds flat at one pass over the expert set — + // and that it runs on tensor cores rather than a scalar GEMV. + // Occupancy was never the mechanism being traded against. + // + // What survives is `DECODE_REGIME_MAX_TOKENS` as a FLOOR: at + // n<=8 there genuinely is not enough work to group (b=1 measures + // 0.47x), and it is also the RUN-161 capture-safety rule. + // 9..=15 is UNMEASURED — bounded below by the n=8 floor and + // above by the +16% at n=16, so it rides with the grouped arm. // // `DECODE_REGIME_MAX_TOKENS` survives as a FLOOR (the RUN-161 // capture-safety rule), and an explicit // `ARC_QTIP_ONDEVICE_MOE_MAX_TOKENS` still pins the GEMV arm // outright so a harness can A/B the two kernels. - let ondevice_preferred = match super::gather_policy::ondevice_max_tokens_override() - { - Some(cap) => n_tokens <= cap, - None => { - n_tokens <= super::DECODE_REGIME_MAX_TOKENS - || !super::gather_policy::grouped_gemm_tiles_amortize( - n_tokens, - n_experts_per_tok, - self.num_experts_count(), - ) - } - }; + let ondevice_preferred = super::gather_policy::prefer_gather_gemv( + n_tokens, + super::gather_policy::ondevice_max_tokens_override(), + ); // Read by VALUE: `ARC_NO_QTIP_ONDEVICE_MOE=0` leaves the - // on-device path enabled (the default). + // on-device path enabled (the default). #197 was cut before + // #212 and still had the presence-check `std::env::var(..) + // .is_ok()` here; keeping #212's parser is deliberate, so this + // retraction does not re-break the flag polarity. let ondevice_disabled = crate::env_flag_is_set("ARC_NO_QTIP_ONDEVICE_MOE"); if !ondevice_disabled && ondevice_preferred @@ -1964,6 +1980,15 @@ impl QuantMethod for Qtip2bLayer { self.num_experts_count(), ); let total_pairs = n_tokens * n_experts_per_tok; + // Same accounting on the amortising arm, so the two paths + // are comparable in the same units (ARC_MOE_BYTE_PROBE=1). + super::byte_probe::record_grouped( + total_pairs, + self.rows_per_expert()?, + self.blocks.dim(2)?, + self.num_experts_count(), + super::grouped::GROUPED_TILE_M, + ); let out_flat = self.gather_forward_batched( &a.reshape((total_pairs, cols))?, &indices.reshape((total_pairs,))?, diff --git a/mistralrs-quant/src/qtip/byte_probe.rs b/mistralrs-quant/src/qtip/byte_probe.rs new file mode 100644 index 000000000..a1eee5196 --- /dev/null +++ b/mistralrs-quant/src/qtip/byte_probe.rs @@ -0,0 +1,273 @@ +//! Parent system: ArcQuant / QTIP — expert-weight byte accounting. +//! +//! Answers one question: **how many bytes of expert weight does the MoE path +//! read per decode step, and does that number stay flat as the batch grows?** +//! +//! - Flat with batch ⇒ amortisation works; a step's expert bytes are shared. +//! - Linear in batch ⇒ every token drags its own copy of the expert weights, +//! and aggregate throughput cannot scale no matter how many users arrive. +//! +//! ## Method (state this whenever you quote the number) +//! +//! These are **kernel-issued bytes**, computed exactly from each launch's own +//! geometry — not a profiler sample, not an estimate, and not DRAM traffic. +//! `ncu` cannot attach on the rented H200 ("Failed to prepare kernel for +//! profiling / Unknown Error on device 0") and the box has no DCGM, so a +//! hardware DRAM counter is not available; `nvidia-smi dmon` reports a +//! utilisation percentage, not bytes. +//! +//! The gather GEMV (`qtip2b_gemv_tuned_kernel`, `qtip_gather_gemv.cu`) is +//! indexed `pair = blockIdx.y`, `row-block = blockIdx.x`, and each block loads +//! its own expert's packed rows. Summed over `blockIdx.x`, **one pair reads one +//! whole expert**: `n_rows * packed_per_row` bytes. There is no dedup between +//! pairs, so a launch reads `n_pairs * n_rows * packed_per_row` bytes — which is +//! what [`record_gemv`] accumulates. +//! +//! The grouped GEMM (`qtip_grouped_gemm.cu`) sorts pairs by expert and stages a +//! woken expert's bytes once per `GROUPED_TILE_M` rows, so it reads +//! `ceil(pairs_e / TILE_M)` copies of expert `e` — [`record_grouped`] uses the +//! same per-expert accounting. +//! +//! Because these are loads *issued*, they are an upper bound on DRAM traffic: +//! repeated reads of the same expert within a short window can be served by L2. +//! That distinction is the point — a card at 99% SM occupancy with a 3% memory +//! controller is exactly what re-reading a cached-but-redundant working set +//! looks like. The redundancy factor below is the quantity that has to reach +//! 1.0 for batching to amortise, whichever cache serves it. +//! +//! Enabled with `ARC_MOE_BYTE_PROBE=1`. Off by default and it costs one relaxed +//! atomic add per expert-GEMM launch when on — no device sync, no D2H read, so +//! it is safe under CUDA-graph capture and does not disturb the dispatch it is +//! measuring. + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::OnceLock; + +static INVOCATIONS: AtomicU64 = AtomicU64::new(0); +static GEMV_INVOCATIONS: AtomicU64 = AtomicU64::new(0); +static GROUPED_INVOCATIONS: AtomicU64 = AtomicU64::new(0); +static PAIRS: AtomicU64 = AtomicU64::new(0); +/// Bytes of expert weight the launches were pointed at. +static ISSUED_BYTES: AtomicU64 = AtomicU64::new(0); +/// Lower bound: each expert woken by a launch read exactly once. +static FLOOR_BYTES: AtomicU64 = AtomicU64::new(0); +/// Bytes of a single expert, for reporting the working-set size. +static BYTES_PER_EXPERT: AtomicU64 = AtomicU64::new(0); +static NUM_EXPERTS: AtomicU64 = AtomicU64::new(0); + +/// `ARC_MOE_BYTE_PROBE=1`. Read once — `std::env::var` takes a global lock and +/// this sits on the per-layer decode path. +pub fn enabled() -> bool { + static ON: OnceLock = OnceLock::new(); + *ON.get_or_init(|| matches!(std::env::var("ARC_MOE_BYTE_PROBE").as_deref(), Ok("1"))) +} + +/// How many distinct experts a launch can possibly have woken. Exact whenever +/// every expert is woken (the saturated regime that matters for the fleet +/// question) and a conservative *over*-estimate otherwise, so the redundancy +/// factor this yields is never flattering. +#[inline] +fn distinct_upper_bound(n_pairs: usize, num_experts: usize) -> usize { + n_pairs.min(num_experts) +} + +/// The non-amortising gather GEMV: one pair reads one whole expert. +pub fn record_gemv(n_pairs: usize, n_rows: usize, packed_per_row: usize, num_experts: usize) { + if !enabled() { + return; + } + let per_expert = (n_rows as u64).saturating_mul(packed_per_row as u64); + let issued = per_expert.saturating_mul(n_pairs as u64); + let floor = per_expert.saturating_mul(distinct_upper_bound(n_pairs, num_experts) as u64); + INVOCATIONS.fetch_add(1, Ordering::Relaxed); + GEMV_INVOCATIONS.fetch_add(1, Ordering::Relaxed); + PAIRS.fetch_add(n_pairs as u64, Ordering::Relaxed); + ISSUED_BYTES.fetch_add(issued, Ordering::Relaxed); + FLOOR_BYTES.fetch_add(floor, Ordering::Relaxed); + BYTES_PER_EXPERT.store(per_expert, Ordering::Relaxed); + NUM_EXPERTS.store(num_experts as u64, Ordering::Relaxed); + maybe_dump(); +} + +/// The amortising grouped GEMM: a woken expert's bytes are staged once per +/// `tile_m` pairs that landed on it. Without the on-device sort order here, use +/// the even-spread estimate `ceil(pairs/distinct / tile_m)` copies per expert — +/// exact when routing is balanced, which is the regime the gate targets. +pub fn record_grouped( + n_pairs: usize, + n_rows: usize, + packed_per_row: usize, + num_experts: usize, + tile_m: usize, +) { + if !enabled() { + return; + } + let per_expert = (n_rows as u64).saturating_mul(packed_per_row as u64); + let distinct = distinct_upper_bound(n_pairs, num_experts).max(1); + let pairs_each = n_pairs.div_ceil(distinct); + let copies = pairs_each.div_ceil(tile_m.max(1)); + let issued = per_expert + .saturating_mul(distinct as u64) + .saturating_mul(copies as u64); + let floor = per_expert.saturating_mul(distinct as u64); + INVOCATIONS.fetch_add(1, Ordering::Relaxed); + GROUPED_INVOCATIONS.fetch_add(1, Ordering::Relaxed); + PAIRS.fetch_add(n_pairs as u64, Ordering::Relaxed); + ISSUED_BYTES.fetch_add(issued, Ordering::Relaxed); + FLOOR_BYTES.fetch_add(floor, Ordering::Relaxed); + BYTES_PER_EXPERT.store(per_expert, Ordering::Relaxed); + NUM_EXPERTS.store(num_experts as u64, Ordering::Relaxed); + maybe_dump(); +} + +/// Periodic cumulative dump so a run yields a number even if it is killed. +/// `ARC_MOE_BYTE_PROBE_EVERY` invocations, default 4096. +fn maybe_dump() { + static EVERY: OnceLock = OnceLock::new(); + let every = *EVERY.get_or_init(|| { + std::env::var("ARC_MOE_BYTE_PROBE_EVERY") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|v| *v > 0) + .unwrap_or(4096) + }); + let n = INVOCATIONS.load(Ordering::Relaxed); + if n.is_multiple_of(every) { + eprintln!("{}", report()); + } +} + +/// One machine-greppable cumulative line. Divide by the tokens the run actually +/// produced to get bytes/token; multiply by the batch to get bytes/step. +pub fn report() -> String { + let inv = INVOCATIONS.load(Ordering::Relaxed); + let gemv = GEMV_INVOCATIONS.load(Ordering::Relaxed); + let grouped = GROUPED_INVOCATIONS.load(Ordering::Relaxed); + let pairs = PAIRS.load(Ordering::Relaxed); + let issued = ISSUED_BYTES.load(Ordering::Relaxed); + let floor = FLOOR_BYTES.load(Ordering::Relaxed); + let per_expert = BYTES_PER_EXPERT.load(Ordering::Relaxed); + let experts = NUM_EXPERTS.load(Ordering::Relaxed); + let redundancy = if floor > 0 { + issued as f64 / floor as f64 + } else { + 0.0 + }; + format!( + "ARC_MOE_BYTES invocations={inv} gemv={gemv} grouped={grouped} pairs={pairs} \ + issued_bytes={issued} floor_bytes={floor} redundancy={redundancy:.3} \ + bytes_per_expert={per_expert} num_experts={experts}" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The GEMV reads one whole expert per pair, so its redundancy is exactly + /// pairs/distinct — the quantity that has to fall to 1.0 for batching to + /// amortise. These are the shapes V4 actually runs (top-6 of 256). + #[test] + fn gemv_redundancy_is_pairs_per_distinct_expert() { + let per_expert = 1024usize * 96; + for (n_tokens, top_k, experts) in [(1usize, 6usize, 256usize), (512, 6, 256)] { + let pairs = n_tokens * top_k; + let distinct = distinct_upper_bound(pairs, experts); + let issued = per_expert * pairs; + let floor = per_expert * distinct; + assert_eq!(issued / floor.max(1), pairs / distinct.max(1)); + } + // B=1 wakes 6 experts with 6 pairs: nothing to share, redundancy 1. + assert_eq!(distinct_upper_bound(6, 256), 6); + // B=512 wakes all 256 with 3072 pairs: 12x redundant. + assert_eq!(3072 / distinct_upper_bound(3072, 256), 12); + } + + /// MEASURED on H200, V4-Flash qtip2b (ARC_MOE_BYTE_PROBE=1): one expert is + /// 2 MiB and a step is 129 gather invocations (43 layers x 3). Expert bytes + /// per step scale LINEARLY with the batch on the GEMV arm and are FLAT on + /// the grouped arm -- which is why the dispatch gate was changed to select + /// grouped above the capture-safety floor. This pins the arithmetic that + /// motivated the fix; `default_gate_selects_the_flat_arm_above_the_floor` + /// below pins the dispatch that resulted from it. + #[test] + fn v4_expert_bytes_per_step_scale_with_batch_on_the_gemv_arm() { + const PER_EXPERT: u64 = 2 * 1024 * 1024; // measured + const INV_PER_STEP: u64 = 129; // measured: 43 layers x 3 gathers + const TOP_K: usize = 6; + const EXPERTS: usize = 256; + let gemv_gib = |n_tokens: usize| -> f64 { + let pairs = (n_tokens * TOP_K) as u64; + (INV_PER_STEP * pairs * PER_EXPERT) as f64 / (1024.0 * 1024.0 * 1024.0) + }; + let floor_gib = |n_tokens: usize| -> f64 { + let distinct = distinct_upper_bound(n_tokens * TOP_K, EXPERTS) as u64; + (INV_PER_STEP * distinct * PER_EXPERT) as f64 / (1024.0 * 1024.0 * 1024.0) + }; + // GEMV: linear in batch, 512x the users reads 512x the expert bytes. + for (n, want) in [(1usize, 1.5f64), (8, 12.1), (64, 96.8), (512, 774.0)] { + assert!( + (gemv_gib(n) - want).abs() < 0.1, + "n_tokens={n}: {:.1} GiB/step, expected {want}", + gemv_gib(n) + ); + } + // Amortised floor saturates at the whole expert working set. + assert!((floor_gib(256) - 64.5).abs() < 0.1); + assert!((floor_gib(512) - 64.5).abs() < 0.1); + // The redundancy the fleet pays at B=512. + assert!((gemv_gib(512) / floor_gib(512) - 12.0).abs() < 0.01); + } + + /// The fix, as an assertion. The gate now selects the arm whose expert + /// bytes are FLAT in the batch, for every width above the capture-safety + /// floor. MEASURED aggregate at the shipped default: B=8 60.28, B=64 + /// 156.61, B=256 193.46 tok/s, against 57.18 / 106.26 / 114.76 under the + /// old tile-fill gate. + /// + /// This branch originally asserted against its own replacement predicate, + /// `gather_policy::grouped_gemm_preferred` — a byte-traffic model gated on + /// `GROUPED_MIN_PAIRS_PER_EXPERT = 1.09`. That predicate is not what ships: + /// #197 arrived first and retired the tile-fill gate outright rather than + /// loosening it, leaving [`prefer_gather_gemv`] — the GEMV arm below + /// [`DECODE_REGIME_MAX_TOKENS`], grouped above it — as the whole dispatch. + /// The two agree exactly at V4's top-6-of-256 (1.09 pairs per woken expert + /// flips at n = 9; the floor is n > 8), so the boundary this test was + /// written to pin is unchanged; only the predicate naming it is. Asserting + /// against the shipped one is the point of the test. + #[test] + fn default_gate_selects_the_flat_arm_above_the_floor() { + use crate::qtip::gather_policy::prefer_gather_gemv; + // At and below the floor the GEMV is mandatory (CUDA-graph capture + // safety) and measured a wash, so nothing is lost. + assert!(prefer_gather_gemv(8, None)); + // Above it, every measured width favours the flat arm. + for n in [9usize, 12, 16, 24, 32, 64, 128, 256, 512] { + assert!( + !prefer_gather_gemv(n, None), + "n={n}: expert bytes on the GEMV arm are {}x the floor here", + (n * 6) / distinct_upper_bound(n * 6, 256).max(1) + ); + } + } + /// The grouped GEMM stages an expert once per TILE_M pairs, so at the same + /// shape it must read strictly fewer bytes than the GEMV once the tiles + /// fill — that is the whole claim the dispatch gate rests on. + #[test] + fn grouped_reads_less_than_gemv_once_tiles_fill() { + let (n_rows, ppr, experts, tile) = (1024usize, 96usize, 256usize, 16usize); + let per_expert = (n_rows * ppr) as u64; + for n_tokens in [512usize, 683, 2048] { + let pairs = n_tokens * 6; + let distinct = distinct_upper_bound(pairs, experts); + let gemv = per_expert * pairs as u64; + let copies = pairs.div_ceil(distinct).div_ceil(tile); + let grouped = per_expert * distinct as u64 * copies as u64; + assert!( + grouped < gemv, + "grouped {grouped} !< gemv {gemv} at n_tokens={n_tokens}" + ); + } + } +} diff --git a/mistralrs-quant/src/qtip/cuda_ops.rs b/mistralrs-quant/src/qtip/cuda_ops.rs index 6a329d952..049b9e8eb 100644 --- a/mistralrs-quant/src/qtip/cuda_ops.rs +++ b/mistralrs-quant/src/qtip/cuda_ops.rs @@ -1973,6 +1973,12 @@ pub(crate) fn gather_gemv_2b_cuda( // One pair per `grid.y` index; past `maxGridSize[1]` the launch fails // silently and the zero-filled output buffer is returned as if valid. super::gather_policy::check_gather_gemv_pairs(n_pairs, "qtip2b gather gemv CUDA")?; + // Expert-weight byte accounting (ARC_MOE_BYTE_PROBE=1). This kernel is + // indexed `pair = blockIdx.y` and each block loads its own expert's packed + // rows, so summed over `blockIdx.x` one pair reads one whole expert and a + // launch reads `n_pairs * n_rows * packed_per_row` bytes with no sharing + // between pairs. Relaxed atomics only: no device sync, capture-safe. + super::byte_probe::record_gemv(n_pairs, n_rows, packed_per_row, num_experts); let x_2d = x_rotated.contiguous()?; let indices = indices.flatten_all()?.contiguous()?; diff --git a/mistralrs-quant/src/qtip/gather_policy.rs b/mistralrs-quant/src/qtip/gather_policy.rs index 7cf8c60a1..27ef73d1d 100644 --- a/mistralrs-quant/src/qtip/gather_policy.rs +++ b/mistralrs-quant/src/qtip/gather_policy.rs @@ -139,8 +139,44 @@ //! 1024 6144 256.0 24.00 150.0% //! ``` //! -//! and the in-tree end-to-end A/B brackets exactly that curve: **1.00× at -//! N=128 and N=512, 2.41× at N=1024**. +//! and the in-tree end-to-end A/B was said to bracket exactly that curve: +//! **1.00× at N=128 and N=512, 2.41× at N=1024**. +//! +//! # 🔴 4a. THE TILE-FILL GATE IS RETRACTED — IT MEASURED THE WRONG QUANTITY +//! +//! **That A/B does not reproduce, and the model above is not what decides the +//! winner.** Two independent measurements on landed master (H200, V4-Flash +//! qtip2b) put the grouped GEMM ahead at every size either could test: +//! +//! ```text +//! decode, clean rows only (>=97% achieved concurrency, <7% spread): +//! B=48 1.10x B=64 1.14x B=80 1.22x B=96 1.29x B=128 1.66x B=512 1.51x +//! prefill, over prompt length: +//! n=16 +16% n=32 +27% n=64 +50% n=128 +78% n=512 +181% +//! ``` +//! +//! At `n = 16` the tile is **7.5% full** and the grouped kernel still wins by +//! 16%. A predicate that demands 100% fill therefore cannot be the right gate: +//! it required `6n/256 >= 16`, i.e. **n >= 683**, which decode never reaches, so +//! it pinned every decode step and every batch up to 682 on the GEMV arm. A +//! profile of the shipped default confirmed the consequence directly — +//! `qtip_gather_gemv_warp_kernel` at **57.6% of GPU time with 25× the launches** +//! of the grouped kernel. +//! +//! **Why fill was the wrong quantity.** The ratio above counts only *weight* +//! bytes and assumes the two kernels are otherwise alike. They are not. The +//! grouped kernel stages a woken expert's bytes once per m-tile rather than once +//! per (token, expert) pair — the GEMV arm's reads are exactly linear in +//! `n_tokens`, measured at **12.00× redundancy at n=512**, while the grouped arm +//! holds flat at one pass over the expert set — and it runs the tensor-core tile +//! loop instead of a scalar per-pair GEMV. A partly-empty tile still buys both. +//! Occupancy was never the mechanism being traded against. +//! +//! What survives is [`super::DECODE_REGIME_MAX_TOKENS`] as a floor: at `n <= 8` +//! there is genuinely nothing to group (b=1 measures **0.47×**), and it is the +//! RUN-161 capture-safety rule. `9..=15` is **unmeasured**, bounded below by +//! that floor and above by the +16% at 16, and rides with the grouped arm. +//! [`prefer_gather_gemv`] is the shipped decision and is unit-tested. //! //! Shipping the gate as a raw `n_tokens > DECODE_REGIME_MAX_TOKENS` switched //! the grouped kernel on at n=9 — onto a tile that is ~9% full, roughly two @@ -305,6 +341,40 @@ pub fn grouped_gemm_tile_fill(n_tokens: usize, top_k: usize, num_experts: usize) /// `GROUPED_TILE_M` changes, and the last one that was not re-derived is the /// regression this replaces. #[cfg_attr(not(feature = "cuda"), allow(dead_code))] +/// Does the fused gather-GEMV arm win at this token count? +/// +/// The pure half of the MoE dispatch in `bitshift.rs`, so the decision that +/// actually ships can be tested without a GPU. It has to live here and be +/// exercised by a test, because the dispatch itself sits inside a +/// `#[cfg(feature = "cuda")]` block that a normal `cargo test` cannot compile — +/// a non-CUDA lane will report a full green over code it never saw. +/// +/// 🔴 This used to also consult [`grouped_gemm_tiles_amortize`]: switch to the +/// grouped GEMM only once each woken expert fills an m-tile. At V4's +/// top-6-of-256 that needs `6n/256 >= GROUPED_TILE_M`, i.e. **n >= 683**, a +/// token count decode never reaches — so the clause pinned every decode step, +/// and every batch up to 682, on the GEMV arm. +/// +/// Two independent measurements retired it: decode at B=48..512 on an H200, and +/// prefill over 16..512-token prompts. Grouped wins everywhere at and above 16, +/// including **+16% on a tile that is only 7.5% full**. +/// +/// **Tile fill was therefore the wrong quantity.** The grouped kernel wins +/// because it stages each woken expert's bytes once per m-tile instead of once +/// per (token, expert) pair — the GEMV arm's reads are exactly linear in +/// `n_tokens` — and because it runs on tensor cores rather than a scalar GEMV. +/// Occupancy was never the mechanism being traded against. +/// +/// [`super::DECODE_REGIME_MAX_TOKENS`] survives as a FLOOR: below it there is +/// genuinely nothing to group (b=1 measures 0.47x), and it is the RUN-161 +/// capture-safety rule besides. +pub(crate) fn prefer_gather_gemv(n_tokens: usize, override_cap: Option) -> bool { + match override_cap { + Some(cap) => n_tokens <= cap, + None => n_tokens <= super::DECODE_REGIME_MAX_TOKENS, + } +} + pub(crate) fn grouped_gemm_tiles_amortize( n_tokens: usize, top_k: usize, @@ -740,3 +810,55 @@ mod tests { } } } + +#[cfg(test)] +mod crossover_tests { + use super::*; + + /// The shipped dispatch: the GEMV arm only at or below the capture-safety + /// floor, and the grouped GEMM reachable at every batch size we measured. + #[test] + fn the_grouped_gemm_is_reachable_at_decode_batch_sizes() { + for n in [1, 2, 4, 8] { + assert!( + prefer_gather_gemv(n, None), + "n={n} is below the floor and must stay on the GEMV arm" + ); + } + // Exactly the points measured on hardware, all of which the grouped + // kernel won: decode B=48 1.10x, 64 1.14x, 80 1.22x, 96 1.29x, + // 128 1.66x, 512 1.51x (clean rows: >=97% achieved concurrency, <7% + // derived-vs-measured spread), and prefill +16% as low as n=16. + for n in [16, 32, 48, 64, 80, 96, 128, 512] { + assert!( + !prefer_gather_gemv(n, None), + "n={n} must reach the grouped GEMM; it was measured faster there" + ); + } + } + + /// NON-VACUITY CONTROL, and the whole reason the clause was removed: the + /// retired tile-fill predicate would have sent every one of those token + /// counts to the GEMV arm. If this stops holding, the control is broken and + /// the test above proves nothing about the change. + #[test] + fn the_retired_tile_fill_clause_would_have_excluded_all_of_them() { + for n in [16, 32, 48, 64, 80, 96, 128, 512] { + assert!( + !grouped_gemm_tiles_amortize(n, 6, 256), + "control broken: the retired clause was supposed to reject n={n}" + ); + } + // It only relented at 683 -- a token count decode never reaches, which + // is precisely why the grouped kernel was unreachable in serving. + assert_eq!(grouped_gemm_min_tokens(6, 256), Some(683)); + } + + /// The A/B override still pins the GEMV arm outright, so both kernels stay + /// comparable from one binary. + #[test] + fn the_override_still_pins_the_gemv_arm() { + assert!(prefer_gather_gemv(512, Some(1024))); + assert!(!prefer_gather_gemv(512, Some(256))); + } +} diff --git a/mistralrs-quant/src/qtip/grouped.rs b/mistralrs-quant/src/qtip/grouped.rs index 17c55f231..b4481ef99 100644 --- a/mistralrs-quant/src/qtip/grouped.rs +++ b/mistralrs-quant/src/qtip/grouped.rs @@ -109,11 +109,43 @@ pub fn grouped_variant() -> i32 { } let init = match std::env::var(QTIP_GROUPED_VARIANT_ENV) { Ok(s) => match s.trim().to_ascii_lowercase().as_str() { + "baseline" | "0" => QTIP_GROUPED_VARIANT_BASELINE, "tuned" | "1" => QTIP_GROUPED_VARIANT_TUNED, "ldst" | "2" => QTIP_GROUPED_VARIANT_LDST, - _ => QTIP_GROUPED_VARIANT_BASELINE, + // An unrecognised value used to fall through to the baseline. Now + // that the baseline is the SLOW arm rather than the default, a typo + // would silently cost ~20% of prefill, so say so and keep the + // default instead of quietly selecting the worst kernel. + other => { + tracing::warn!( + "{QTIP_GROUPED_VARIANT_ENV}=`{other}` is not a known grouped-GEMM \ + variant (baseline|0, tuned|1, ldst|2); using the default." + ); + QTIP_GROUPED_VARIANT_LDST + } }, - Err(_) => QTIP_GROUPED_VARIANT_BASELINE, + // 🔴 The DEFAULT is the tuned kernel, not the baseline. + // + // This used to be `QTIP_GROUPED_VARIANT_BASELINE`, and nothing in the + // serving path ever calls `set_grouped_variant` -- its only callers are + // in `examples/qtip_grouped_curve.rs`. So every production prefill ran + // variant 0 and the tuned kernels were reachable only by setting an env + // var nobody sets. They are bit-identical in output and were already + // measured at +39.4% (v1) and +41.6% (v2) per m-tile. + // + // ✅ MEASURED END-TO-END 2026-08-20, H200, DeepSeek-V4-Flash qtip2b, + // pp2048 b=1, --no-paged-attn, arm proved from the runtime launch + // counters ([0,0,129] on v2 vs [129,0,0] on v0, i.e. 43 layers x 3 + // matrices, and zero launches on the arms not under test): + // variant 0: 561.7 prompt tok/s, 1.780 ms/prompt-token, TTFT 3.675 s + // variant 2: 671.3 prompt tok/s, 1.490 ms/prompt-token, TTFT 3.080 s + // => +19.5% prefill throughput, -16.2% TTFT. + // An engine-independent wall clock agrees (557.2 -> 664.9 tok/s, +19.3%). + // + // The end-to-end gain is smaller than the kernel gain because the gather + // is ~56% of a 2048-token prefill step, not 100% -- backing +41.6% on + // that share out of a measured +19.5% overall is consistent. + Err(_) => QTIP_GROUPED_VARIANT_LDST, }; let _ = GROUPED_VARIANT.compare_exchange( GROUPED_VARIANT_UNINIT, diff --git a/mistralrs-quant/src/qtip/mod.rs b/mistralrs-quant/src/qtip/mod.rs index af6c8f295..9f96d72fe 100644 --- a/mistralrs-quant/src/qtip/mod.rs +++ b/mistralrs-quant/src/qtip/mod.rs @@ -61,6 +61,10 @@ mod bake_memory_tests; #[cfg(test)] mod bake_quality_tests; pub mod bitshift; +// Every call site is behind `#[cfg(feature = "cuda")]`, so a CPU-only build +// sees the whole probe as dead. Same convention as `gather_policy`. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] +pub mod byte_probe; #[cfg(test)] mod codebook_tests; #[cfg(feature = "cuda")] @@ -686,6 +690,10 @@ impl QtipCodebook { /// one. Unknown tags fail closed (DOCTRINE: refuse, never launder). fn from_wire(tag: u8, read_mult: impl FnOnce() -> Result) -> Result { match tag { + GEOMETRY_WIRE_TAG => candle_core::bail!( + "QtipLayer: internal error — the geometry section reached the codebook parser. \ + The trailing-section loop must consume tag {GEOMETRY_WIRE_TAG} itself." + ), 2 => { let mult = read_mult()?; if mult == 0 || mult % 2 == 0 { @@ -705,6 +713,334 @@ impl QtipCodebook { } } +/// Trailing-section tag for the trellis geometry. +/// +/// **It is 3, and it is written BEFORE the codebook section, and both of those +/// are load-bearing.** A build that predates this field parses the trailing +/// region by reading one byte and handing it to [`QtipCodebook::from_wire`], +/// which refuses every tag it does not know. So an old Arc handed a +/// non-default-geometry artifact sees tag 3 first and fails closed with +/// "written by a newer Arc … refusing rather than decoding its symbols against +/// the wrong codebook" — which is exactly the right outcome, because it cannot +/// decode this geometry either. +/// +/// Put the section *after* the codebook and that property is lost: an old +/// reader would consume the codebook section, stop, and decode a K=8/V=4 symbol +/// stream as K=4/V=2. It would not even fault — the packed byte count is +/// identical at 2 bpw, so every index stays in bounds and the output is +/// plausible garbage. [`tests::geometry_section_precedes_the_codebook_section`] +/// pins the order. +const GEOMETRY_WIRE_TAG: u8 = 3; + +/// Which trellis geometry a [`QtipLayer`]'s symbols were baked at. +/// +/// **This is a format discriminator, not a tuning knob.** Symbols are a +/// bit-stream interpreted by a specific `(K, L, V)` and a specific reproduction +/// table; read at the wrong geometry they decode to different weights, silently. +/// +/// `K` is carried as data, not baked into the variant, because that is what the +/// wire already says — the section is `[tag, K, L, V]`. Within the V=4/L=12 +/// family the table is K-independent, so a new K is a new value here and not a +/// new variant. See [`crate::trellis_v4l12`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum QtipGeometry { + /// K=4 / V=2 / L=16 with a `[65536, 2]` F32 table — every QTIP artifact + /// Arc has ever written. The default, and it serializes as **nothing**. + #[default] + K4V2L16, + /// The V=4 / L=12 family: a `[4096, 4]` **BF16** table (32,768 B — fits + /// static shared memory), with `k` bits per symbol. + /// + /// `k = 8` is the byte-aligned control at 2.00 bpw; `k = 9` is 2.25 bpw and + /// is the quality winner. Constructed only through + /// [`QtipGeometry::trellis_v4l12`], which refuses a `k` with no decoder. + TrellisV4L12 { + /// The decode rung, which owns the K-dependent arithmetic. Stored + /// rather than a bare `k` so the packed-size formula has exactly one + /// implementation — see [`QtipGeometry::packed_len`]. + rung: trellis_v4l12::Rung, + }, +} + +impl QtipGeometry { + /// The V=4/L=12 family at symbol width `k`, refusing a `k` this build has + /// no decoder for. + /// + /// Refuses rather than storing it: `k` arrives from an artifact, and a + /// geometry we cannot decode must not be constructible as a value that + /// later code will try to use. + pub fn trellis_v4l12(k: u32) -> Result { + let rung = trellis_v4l12::Rung::new(k).map_err(candle_core::Error::Msg)?; + Ok(QtipGeometry::TrellisV4L12 { rung }) + } + + /// Bits per symbol. + pub fn k(self) -> u32 { + match self { + QtipGeometry::K4V2L16 => K, + QtipGeometry::TrellisV4L12 { rung } => rung.k(), + } + } + + /// Trellis state width in bits. + pub fn l(self) -> u32 { + match self { + QtipGeometry::K4V2L16 => L, + QtipGeometry::TrellisV4L12 { .. } => trellis_v4l12::L, + } + } + + /// Reproduction values per symbol. + pub fn v(self) -> u32 { + match self { + QtipGeometry::K4V2L16 => V, + QtipGeometry::TrellisV4L12 { .. } => trellis_v4l12::V, + } + } + + /// Bits per weight × 100 (`100·K/V`), as an integer so it compares exactly. + /// + /// 200 for the shipped rung and for K=8/V=4; **225 for K=9/V=4**; 250 for + /// K=10/V=4. Not every geometry in this enum is the same bit rate any more, + /// which is a change from when the family was K=8-only. + pub fn bpw_x100(self) -> u32 { + self.k() * 100 / self.v() + } + + /// Total values in the reproduction table: `2^L × V`. + pub fn lut_values(self) -> usize { + (1usize << self.l()) * self.v() as usize + } + + /// Element type of the reproduction table. + /// + /// Part of the discriminator, not an implementation detail: the V=4 table + /// is BF16 precisely so it lands at 32,768 B, and an F32 table of the same + /// shape is a different artifact. + pub fn lut_dtype(self) -> DType { + match self { + QtipGeometry::K4V2L16 => DType::F32, + QtipGeometry::TrellisV4L12 { .. } => DType::BF16, + } + } + + /// Trellis symbols in a row of `in_features` weights. Depends only on V. + pub fn num_symbols(self, in_features: usize) -> usize { + in_features / self.v() as usize + } + + /// Bytes of a row that actually hold symbols, governed by the bit rate. + /// + /// **Not the row stride.** [`QtipGeometry::packed_len`] is what a tensor is + /// allocated at and what this format validates; it additionally carries + /// tail padding and a round-up to 4 for the V=4/L=12 family. Conflating + /// them is how a bits-per-weight claim ends up four bytes wrong. + pub fn data_bytes(self, in_features: usize) -> usize { + match self { + QtipGeometry::K4V2L16 => self.num_symbols(in_features) / 2, + QtipGeometry::TrellisV4L12 { rung } => rung.data_bytes(rung.num_symbols(in_features)), + } + } + + /// The allocated length of a row of `in_features` weights — the tensor's + /// last dimension, and what this format validates. + /// + /// **Delegated, not restated.** `num_symbols / (8 / K)` is wrong at K=9 — + /// there is no whole number of symbols per byte — and divides by zero, so + /// the formula has to be the general one. It is also invisible to test at + /// realistic shapes: every plausible `in_features` is a multiple of 32, so + /// `num_symbols · K` is a multiple of 8 and floor equals ceil. A duplicate + /// of this formula here would therefore be unguarded in practice (measured: + /// mutation W3 floored it and every format test stayed green), so the + /// family's own [`trellis_v4l12::Rung`] is the single implementation and it + /// is exercised at non-byte-aligned symbol counts by that module's tests. + pub fn packed_len(self, in_features: usize) -> usize { + match self { + QtipGeometry::K4V2L16 => self.num_symbols(in_features) / 2, + QtipGeometry::TrellisV4L12 { rung } => rung.packed_len(in_features), + } + } + + /// What a fresh bake picks when `ARC_QTIP_GEOMETRY` says nothing. + /// + /// **`K4V2L16`, and the flip is this one line.** Deliberately not taken + /// here: nothing about the V=4/L=12 family has touched a GPU — its kernel + /// has never been compiled at any K — and the precedent for shipping a + /// CPU-validated storage change default-on is specific and bad. Its + /// arithmetic was proven and its device-time cost was not, and it killed + /// every request on the first V4 forward that met a real device. Ship the + /// path; leave the default alone. Pinned by + /// [`tests::the_new_geometry_is_not_the_default`]. + pub const DEFAULT: Self = QtipGeometry::K4V2L16; + + /// Read the geometry choice from the environment. + /// + /// `ARC_QTIP_GEOMETRY` is `k4v2l16` (the default) or `kv4l12` for a K + /// this build implements — `k9v4l12` is the one worth asking for. Anything + /// else is refused rather than silently resolved, exactly as + /// [`QtipCodebook::from_env`] refuses: picking the wrong one changes what + /// every symbol in the artifact means. + pub fn from_env() -> Result { + match std::env::var("ARC_QTIP_GEOMETRY") { + Err(_) => Ok(Self::DEFAULT), + Ok(v) => Self::from_spec(&v), + } + } + + /// The parse rule behind [`QtipGeometry::from_env`], as a pure function. + /// + /// Split out so it can be tested directly. Testing a *replica* of this rule + /// instead — which is what the first version of the test did — asserts only + /// that the replica behaves, and passes happily while the real reader + /// silently accepts an unsupported K (measured: mutation S7). Mutating the + /// process environment in a unit test races the other tests in the binary, + /// so the function is the seam, not the variable. + pub fn from_spec(spec: &str) -> Result { + match spec { + "" => Ok(Self::DEFAULT), + "k4v2l16" | "k4" | "shipped" => Ok(QtipGeometry::K4V2L16), + other => { + // `kv4l12`, so a new K needs no new arm here — the family + // decides whether it has a decoder. + let parsed = other + .strip_prefix('k') + .and_then(|r| r.strip_suffix("v4l12")) + .and_then(|k| k.parse::().ok()); + match parsed { + Some(k) => Self::trellis_v4l12(k), + None => candle_core::bail!( + "ARC_QTIP_GEOMETRY={other:?} is not a geometry. Use `k4v2l16` (the \ + default) or `kv4l12` with K in {:?}. Refusing rather than guessing \ + — the choice changes what every symbol in the artifact means.", + trellis_v4l12::K_SUPPORTED + ), + } + } + } + } + + /// Short label for bake headers and error messages. + pub fn tag(self) -> String { + match self { + QtipGeometry::K4V2L16 => "k4v2l16".to_string(), + QtipGeometry::TrellisV4L12 { rung } => format!("k{}v4l12", rung.k()), + } + } + + /// Wire encoding. **The default geometry writes nothing**, so every + /// artifact Arc has already produced stays byte-identical and no existing + /// checksum moves. Only a geometry that cannot be inferred from absence + /// announces itself. + fn to_wire(self) -> Option<(u8, [u8; 3])> { + match self { + QtipGeometry::K4V2L16 => None, + QtipGeometry::TrellisV4L12 { rung } => Some(( + GEOMETRY_WIRE_TAG, + [ + rung.k() as u8, + trellis_v4l12::L as u8, + trellis_v4l12::V as u8, + ], + )), + } + } + + /// Parse a `(K, L, V)` triple into a known geometry. + /// + /// The triple is stored rather than an opaque ordinal so a payload is + /// self-describing to a human with a hex editor, and so an unsupported but + /// well-formed geometry produces a diagnosis instead of "unknown tag". + fn from_wire_body(k: u8, l: u8, v: u8) -> Result { + // Written as comparisons, not as match patterns: a bare `K` in pattern + // position is a const pattern only for as long as the const exists, and + // silently becomes a catch-all binding if it is ever renamed away. + let got = (k as u32, l as u32, v as u32); + if got == (K, L, V) { + return Ok(QtipGeometry::K4V2L16); + } + if got.1 == trellis_v4l12::L && got.2 == trellis_v4l12::V { + // In-family: K is the free parameter, so let the family say + // whether it has a decoder for this one. + return QtipGeometry::trellis_v4l12(got.0); + } + let (ok, ol, ov) = got; + candle_core::bail!( + "QtipLayer: UQFF payload declares trellis geometry K={ok}/L={ol}/V={ov}, which this \ + build has no decoder for. Refusing rather than decoding its symbols at the wrong \ + geometry — the packed byte count alone does not identify a geometry, so a wrong \ + guess need not even fault.", + ) + } + + /// Check that a layer's tensors are shaped the way this geometry requires. + /// + /// The tag alone is a claim; this is what makes it true. The table's dtype + /// and length are the load-bearing half: K=4/V=2 and K=8/V=4 are both + /// 2 bits per weight, so `blocks` has the identical size either way and + /// cannot discriminate them. + fn validate_shapes(self, blocks: &Tensor, lut: &Tensor, in_features: usize) -> Result<()> { + if !in_features.is_multiple_of(self.v() as usize) { + candle_core::bail!( + "QtipLayer[{}]: in_features {in_features} is not a multiple of V={}", + self.tag(), + self.v() + ); + } + let want_packed = self.packed_len(in_features); + let got_packed = blocks.dims()[blocks.dims().len() - 1]; + if got_packed != want_packed { + candle_core::bail!( + "QtipLayer[{}]: blocks row is {got_packed} B but K={}/V={} with in_features \ + {in_features} needs {want_packed} B", + self.tag(), + self.k(), + self.v() + ); + } + if lut.dtype() != self.lut_dtype() { + candle_core::bail!( + "QtipLayer[{}]: reproduction table is {:?} but this geometry stores {:?}. The \ + table's element type is part of the format, not an implementation detail.", + self.tag(), + lut.dtype(), + self.lut_dtype() + ); + } + if lut.elem_count() != self.lut_values() { + candle_core::bail!( + "QtipLayer[{}]: reproduction table has {} values, expected 2^{} × {} = {}", + self.tag(), + lut.elem_count(), + self.l(), + self.v(), + self.lut_values() + ); + } + Ok(()) + } + + /// Codebooks this geometry can legally carry. + /// + /// The computed `sum2` codebook is V=2-specific — `mcg_codeword_v2` folds + /// exactly two chained MCG products, and the CUDA twin + /// (`qtip_codebook.cuh::qtip_cb_pair_from_x0`) returns a `float2`. There is + /// no V=4 form of it, so pairing it with any V≠2 geometry is not a thing + /// that can be decoded, and it is refused on both the write and the read + /// side rather than left to fail somewhere less obvious. + fn check_codebook(self, codebook: QtipCodebook) -> Result<()> { + match codebook { + QtipCodebook::Mcg { .. } if self.v() != 2 => candle_core::bail!( + "QtipLayer[{}]: the computed `sum2` codebook is V=2-only (it produces a PAIR of \ + reproduction values per state) and cannot describe a V={} geometry. This rung \ + is table-only.", + self.tag(), + self.v() + ), + _ => Ok(()), + } + } +} + /// QTIP 2-bit weight layer. /// /// # Storage layout — two modes @@ -799,6 +1135,16 @@ pub struct QtipLayer { /// computes, so the table stays a correct fallback (and the CPU search /// still wants a flat one) while every GPU path skips it entirely. codebook: QtipCodebook, + /// Which trellis geometry `blocks` was baked at, and therefore how its + /// bytes decompose into symbols and how `lut` is indexed. + /// + /// [`QtipGeometry::K4V2L16`] for every artifact written before the + /// discriminator existed — which is what they are. Deliberately has no + /// per-site default: both supported geometries are 2 bits per weight, so a + /// mislabelled layer has a correctly-sized `blocks` tensor and decodes to + /// plausible garbage rather than faulting. Every construction site states + /// it, and the compiler is what enforces that. + geometry: QtipGeometry, } /// Borrowed, dequantization-free view of a [`QtipLayer`]'s packed trellis @@ -1954,6 +2300,10 @@ impl QtipLayer { search_weights.is_some(), ), codebook, + // This crate's trellis search is K=4/V=2/L=16 throughout + // (`viterbi.rs` reaches 16 successors per step and groups by + // 2^(L-K)); there is no producer of any other geometry yet. + geometry: QtipGeometry::K4V2L16, }) } @@ -2082,6 +2432,11 @@ impl QtipLayer { // objective bit is false by construction, not by omission. search_detail: QtipSearchDetail::for_bake(mode, search, false), codebook, + // Same geometry as the CPU sibling above: the CUDA bake kernels are + // the K=4/V=2/L=16 trellis too, so a GPU-baked artifact must carry + // the identical wire tag. Omitting it here compiled fine without + // the `cuda` feature and only broke the CUDA lane. + geometry: QtipGeometry::K4V2L16, })) } @@ -2220,6 +2575,7 @@ impl QtipLayer { let mut shared_rotation_block: usize = 0; let mut shared_search_detail = QtipSearchDetail::Unknown; let mut shared_codebook = QtipCodebook::Gaussian; + let mut shared_geometry = QtipGeometry::K4V2L16; // Per-expert streaming. The full stack is kept on CPU (caller passes // device=CPU for the 3-D experts to avoid the ~4GB dense BF16 transient @@ -2341,6 +2697,7 @@ impl QtipLayer { shared_rotation_block = layer.rotation_block; shared_search_detail = layer.search_detail; shared_codebook = layer.codebook; + shared_geometry = layer.geometry; } else { debug_assert_eq!(layer.lut.dims(), shared_lut.as_ref().unwrap().dims()); debug_assert_eq!(layer.rotation_block, shared_rotation_block); @@ -2369,6 +2726,18 @@ impl QtipLayer { shared_codebook.tag() ); } + // And the geometry, for the same reason and with the same + // teeth: at 2 bpw a mixed-geometry stack has consistently + // sized `blocks`, so nothing downstream would notice. + if layer.geometry != shared_geometry { + candle_core::bail!( + "QTIP 3-D quantize: expert chunk at {expert_idx} used geometry {} but \ + the stack already carries {} — refusing to build a stack whose experts \ + decode at different geometries.", + layer.geometry.tag(), + shared_geometry.tag() + ); + } } expert_idx += this_b; } @@ -2398,6 +2767,7 @@ impl QtipLayer { search: QtipSearchStamp::for_mode(mode), search_detail: shared_search_detail, codebook: shared_codebook, + geometry: shared_geometry, }) } @@ -2409,6 +2779,9 @@ impl QtipLayer { // path (GPU-validated numerical parity); see arc-tools/CI_HYGIENE.md. #[allow(clippy::needless_range_loop)] fn dequantize_weights_rotated_f32(&self) -> Result> { + if let Some(rung) = self.v4l12_rung() { + return self.dequantize_weights_rotated_f32_v4l12(rung); + } let blocks_cpu = self.blocks.to_device(&Device::Cpu)?; let scales_cpu = self.row_scales.to_device(&Device::Cpu)?; let lut_cpu = self.lut.to_device(&Device::Cpu)?; @@ -2447,6 +2820,59 @@ impl QtipLayer { Ok(out) } + /// CPU decode for the V=4/L=12 family, in the rotated frame and scaled. + /// + /// Separate from the K=4/V=2 body rather than a branch inside it: the + /// symbol extraction, the packed row length and the table's element type + /// all differ, and `trellis_v4l12::Rung` already owns that arithmetic and + /// is bit-exactly gated against the GPU kernel. Restating it here would be + /// a second implementation of a format. + fn dequantize_weights_rotated_f32_v4l12(&self, rung: trellis_v4l12::Rung) -> Result> { + if self.lut.dtype() != DType::BF16 { + candle_core::bail!( + "QtipLayer[{}]: table must be BF16, got {:?}", + self.geometry.tag(), + self.lut.dtype() + ); + } + let blocks_cpu = self.blocks.to_device(&Device::Cpu)?.contiguous()?; + let scales_cpu = self.row_scales.to_device(&Device::Cpu)?.contiguous()?; + let lut_cpu = self.lut.to_device(&Device::Cpu)?.contiguous()?; + + let blocks_data: Vec = blocks_cpu.flatten_all()?.to_vec1()?; + let scales_data: Vec = scales_cpu.flatten_all()?.to_vec1()?; + let lut_data: Vec = lut_cpu.flatten_all()?.to_vec1()?; + + let n = scales_data.len(); + let k_in = self.in_features; + let num_symbols = rung.num_symbols(k_in); + // The allocated STRIDE, not the data length: rows carry tail padding + // so the kernel's compile-time-width extraction needs no clamp, and + // that padding is part of what the artifact stores. + let packed_per_row = rung.row_stride(num_symbols); + if blocks_data.len() != n * packed_per_row { + candle_core::bail!( + "QtipLayer[{}]: blocks hold {} B but {n} rows × {packed_per_row} B = {}", + self.geometry.tag(), + blocks_data.len(), + n * packed_per_row + ); + } + + let mut out = vec![0f32; n * k_in]; + for row in 0..n { + let packed = &blocks_data[row * packed_per_row..(row + 1) * packed_per_row]; + let dst = &mut out[row * k_in..(row + 1) * k_in]; + rung.decode_row_into( + trellis_v4l12::Row::new(packed, num_symbols, scales_data[row]), + &lut_data, + dst, + ) + .map_err(candle_core::Error::Msg)?; + } + Ok(out) + } + fn dequantize_weights_f32(&self) -> Result> { let mut out = self.dequantize_weights_rotated_f32()?; @@ -2473,6 +2899,27 @@ impl QtipLayer { } pub fn dequantize_weights(&self) -> Result { + // The V=4/L=12 family decodes on CPU (`dequantize_weights_rotated_f32` + // routes to it). There is no GPU dequantize kernel for the family and + // no 3-D path, so those two stay refused rather than falling through to + // a K=4 decoder. + if let Some(rung) = self.v4l12_rung() { + if self.num_experts.is_some() { + candle_core::bail!( + "QtipLayer::dequantize_weights: geometry {} has no stacked-expert (3-D) \ + decode. Refusing rather than reading these symbols with the K=4/V=2 \ + expert path.", + self.geometry.tag() + ); + } + let _ = rung; + let out = self.dequantize_weights_f32()?; + let n = self.row_scales.dim(0)?; + return Tensor::from_vec(out, (n, self.in_features), &Device::Cpu)? + .to_device(self.blocks.device())? + .to_dtype(DType::BF16); + } + self.require_k4v2l16("dequantize_weights")?; // 3-D stacked-expert path: iterate over experts and stack `[N, K_in]` // slices into `[E, N, K_in]`. We reuse the 2-D dequant per expert // rather than writing a 3-D CUDA kernel — the existing @@ -2676,7 +3123,11 @@ impl QtipLayer { // CPU path below. #[cfg(feature = "cuda")] { - if cuda_ops::can_use_qtip_cuda(&self.blocks) { + // `forward_dequantize_cuda` is the K=4/V=2 CUDA path end to end — + // its dequantize kernel and its fused gemv both index a + // `[2^16, 2]` F32 table. A family layer must never reach it, and + // gets its own dispatch below. + if self.v4l12_rung().is_none() && cuda_ops::can_use_qtip_cuda(&self.blocks) { if let Ok(mut result) = self.forward_dequantize_cuda(&x_2d, x.dtype(), n, k_in) { if let Some(bias) = &self.bias { result = result.broadcast_add(bias)?; @@ -2691,6 +3142,29 @@ impl QtipLayer { } } + // V=4/L=12 CUDA single-token decode: the one GPU path the family has. + // Anything else — multi-token, or a shape the launcher refuses — falls + // through to the CPU decode below, which is correct and slow, rather + // than to a K=4 kernel, which would be fast and wrong. + #[cfg(feature = "cuda")] + { + if let Some(rung) = self.v4l12_rung() { + if cuda_ops::can_use_qtip_cuda(&self.blocks) && x_2d.dim(0)? == 1 { + if let Ok(mut result) = self.forward_v4l12_gemv_cuda(&x_2d, x.dtype(), rung) { + if let Some(bias) = &self.bias { + result = result.broadcast_add(bias)?; + } + if orig_dims.len() > 2 { + let mut new_dims = orig_dims[..orig_dims.len() - 1].to_vec(); + new_dims.push(result.dim(1)?); + result = result.reshape(new_dims)?; + } + return Ok(result); + } + } + } + } + // Materializing fallback: full weight dequantize + matmul. At decode // shapes the fused GEMV above should have handled this. warn_dequant_materialize_at_decode( @@ -2760,6 +3234,76 @@ impl QtipLayer { /// /// Returns `Err` if any precondition fails (e.g. unsupported dtype or /// block size) — the caller falls back to the CPU path. + /// V=4/L=12 single-token fused decode + gemv on CUDA. + /// + /// The family's only GPU path. Deliberately narrow: `n_tokens == 1`, the + /// shape the fused kernel exists for. Multi-token falls back to the CPU + /// decode rather than to a K=4 kernel — slow and right beats fast and + /// wrong, and there is no V=4 dequantize kernel to matmul against. + /// + /// The row-scale hoist is **off**: it is the one lever that costs + /// bit-exactness against the CPU reference, and nothing about this family + /// has been measured on a device, so there is no evidence to trade + /// exactness for. + #[cfg(feature = "cuda")] + fn forward_v4l12_gemv_cuda( + &self, + x_2d: &Tensor, + out_dtype: DType, + rung: trellis_v4l12::Rung, + ) -> Result { + if !matches!(out_dtype, DType::BF16 | DType::F16 | DType::F32) { + candle_core::bail!("QTIP v4l12 forward CUDA: unsupported out dtype {out_dtype:?}"); + } + if self.num_experts.is_some() { + candle_core::bail!("QTIP v4l12 forward CUDA: no stacked-expert path"); + } + if tune::spec_pin_gemm() { + candle_core::bail!("QTIP v4l12 forward CUDA: spec_pin_gemm pins off the fused path"); + } + let k_in = self.in_features; + if x_2d.dim(1)? != k_in { + candle_core::bail!("QTIP v4l12 forward CUDA: x.dim(1) != in_features"); + } + + let x_rot_dtype = if matches!(x_2d.dtype(), DType::BF16 | DType::F16 | DType::F32) { + x_2d.dtype() + } else { + DType::BF16 + }; + let x_for_rot = if x_2d.dtype() == x_rot_dtype { + x_2d.contiguous()? + } else { + x_2d.to_dtype(x_rot_dtype)?.contiguous()? + }; + let x_rotated = if self.rotation_block >= 2 { + let signs = match &self.rotation_signs { + Some(t) => t, + None => candle_core::bail!( + "QtipLayer: rotation_block={} but rotation_signs is None", + self.rotation_block + ), + }; + cuda_ops::rotate_x_cuda(&x_for_rot, signs, self.rotation_block)? + } else { + x_for_rot + }; + + let y = cuda_ops::fused_gemv_v4l12_cuda( + &self.blocks, + &self.row_scales, + &self.lut, + &x_rotated, + k_in, + rung, + trellis_v4l12::RowScaleHoist::Off, + )?; + if y.dtype() != out_dtype { + return y.to_dtype(out_dtype); + } + Ok(y) + } + #[cfg(feature = "cuda")] fn forward_dequantize_cuda( &self, @@ -3084,6 +3628,7 @@ impl QtipLayer { search: QtipSearchStamp, search_detail: QtipSearchDetail, codebook: QtipCodebook, + geometry: QtipGeometry, ) -> Result { if blocks.dims().len() != 3 { candle_core::bail!( @@ -3112,6 +3657,12 @@ impl QtipLayer { ); } let e = blocks.dim(0)?; + // The caller asserts a geometry; make it true here rather than at the + // first decode. Both supported geometries are 2 bits per weight, so + // `blocks` alone cannot tell them apart — the table's dtype and length + // are what discriminate. + geometry.validate_shapes(&blocks, &lut, in_features)?; + geometry.check_codebook(codebook)?; Ok(QtipLayer { search_detail, blocks, @@ -3124,6 +3675,7 @@ impl QtipLayer { rotation_block, search, codebook, + geometry, }) } @@ -3201,6 +3753,15 @@ impl QtipLayer { head.codebook.tag() ); } + // Only `head.lut` survives the stack here too, and a table is + // meaningful only at the geometry it was built for. + if layer.geometry != head.geometry { + candle_core::bail!( + "QtipLayer::stack_experts: layer {i} geometry={} != head {}", + layer.geometry.tag(), + head.geometry.tag() + ); + } } let blocks_refs: Vec<&Tensor> = per_expert_layers.iter().map(|l| &l.blocks).collect(); @@ -3225,6 +3786,7 @@ impl QtipLayer { search: head.search, search_detail: head.search_detail, codebook: head.codebook, + geometry: head.geometry, }) } @@ -3507,6 +4069,7 @@ impl QtipLayer { search: self.search, search_detail: self.search_detail, codebook: self.codebook, + geometry: self.geometry, }) } } @@ -3549,6 +4112,11 @@ impl QuantMethod for QtipLayer { // table holds the computed values. It only forgoes the // in-register decode. codebook: QtipCodebook::Gaussian, + // Same again for the geometry: the config carries no + // discriminator, and every payload that reached this factory + // before the field existed was K=4/V=2/L=16. Reading it as + // anything else would be inventing a claim. + geometry: QtipGeometry::K4V2L16, }), _ => candle_core::bail!("QtipLayer requires QuantMethodConfig::Qtip"), } @@ -3559,10 +4127,22 @@ impl QuantMethod for QtipLayer { } fn qtip_packed(&self) -> Option> { - Some(self.packed_view()) + // `QtipPackedView` carries no geometry field, so a consumer would read + // these bytes as K=4/V=2 nibbles. Hand out nothing rather than + // something that cannot describe itself. + match self.geometry { + QtipGeometry::K4V2L16 => Some(self.packed_view()), + _ => None, + } } fn forward(&self, x: &Tensor) -> Result { + // 2-D forward is wired for the V=4/L=12 family; + // `forward_dequantize` picks the family's paths and refuses the ones + // that have no kernel. 3-D stacks are refused there, as before. + if self.v4l12_rung().is_none() { + self.require_k4v2l16("forward")?; + } self.forward_dequantize(x) } @@ -3606,6 +4186,7 @@ impl QuantMethod for QtipLayer { /// dequantized weight in the rotated frame and matching it with a /// rotated activation — saves one rotation pass per expert. fn gather_forward(&self, a: &Tensor, indices: &Tensor) -> Result { + self.require_k4v2l16("gather_forward")?; if self.num_experts.is_none() { // The contract is "expert-sparse dispatch" — a 2-D layer is a // single-expert (i.e. non-MoE) layer and `gather_forward` @@ -3689,14 +4270,19 @@ impl QuantMethod for QtipLayer { // RUN-161 floor, not a performance choice (see // `gather_policy`). An explicit `ARC_QTIP_ONDEVICE_MOE_MAX_TOKENS` // override also still wins, so a harness can pin the GEMV arm. + // 🔴 The `grouped_gemm_tiles_amortize` conjunct is GONE — see + // `gather_policy` §4a. It demanded `6n/256 >= 16`, i.e. + // n >= 683, so on this rung too the grouped GEMM was + // unreachable in serving and every decode step ran the per-pair + // GEMV. Measured on THIS rung (qtip2 LUT, H200, V4-Flash), + // clean rows only: forcing the grouped path is 1.10x at B=48, + // 1.14x at 64, 1.22x at 80, 1.29x at 96, 1.66x at 128 and + // 1.51x at 512. Tile fill is not the deciding quantity; staging + // each woken expert once per m-tile instead of once per + // (token, slot) pair is. let grouped_preferred = grouped_available && n_tokens > DECODE_REGIME_MAX_TOKENS - && gather_policy::ondevice_max_tokens_override().is_none() - && gather_policy::grouped_gemm_tiles_amortize( - n_tokens, - n_experts_per_tok, - num_experts, - ); + && gather_policy::ondevice_max_tokens_override().is_none(); let use_ondevice = !grouped_preferred && match gather_policy::ondevice_max_tokens_override() { @@ -3910,10 +4496,54 @@ impl QtipLayer { self.codebook } + /// Which trellis geometry this layer's symbols were baked at, and + /// therefore which decoder can read them. See [`QtipGeometry`]. + pub fn geometry(&self) -> QtipGeometry { + self.geometry + } + + /// Refuse any decode path that only knows K=4/V=2/L=16. + /// + /// **This is what keeps stage 2 from being a half-applied change.** The + /// geometry discriminator makes a K=8/V=4/L=12 artifact *loadable*; the + /// decoders that would then read it — `dequantize_weights_rotated_f32`, + /// `dequantize_single_expert`, `gather_forward_cpu`, and every CUDA + /// launcher in `cuda_ops` other than `fused_gemv_k8v4l12_cuda` — all + /// unpack nibbles and index a `[2^16, 2]` table. Handed a K=8 row they + /// would not fault: at 2 bits per weight the byte count is identical and + /// every index lands in bounds. They would serve garbage. + /// + /// So every entry point into decode states the geometries it can handle. + /// A layer this build cannot serve is refused loudly at the door rather + /// than quietly at the logits. + fn require_k4v2l16(&self, op: &str) -> Result<()> { + match self.geometry { + QtipGeometry::K4V2L16 => Ok(()), + other => candle_core::bail!( + "QtipLayer::{op}: this layer was baked at geometry {}, and {op} implements only \ + K=4/V=2/L=16. There is no V=4/L=12 kernel behind this entry point — the family \ + has a fused decode+gemv kernel (`kernels/qtip/qtip_gemv_v4l12.cu`) and a CPU \ + decode, and nothing else. Refusing rather than decoding these symbols as \ + K=4/V=2: at 2 bits per weight the packed byte count is identical, so the \ + misread would not fault and would silently serve wrong weights.", + other.tag() + ), + } + } + + /// The decode rung when this layer belongs to the V=4/L=12 family. + fn v4l12_rung(&self) -> Option { + match self.geometry { + QtipGeometry::TrellisV4L12 { rung } => Some(rung), + QtipGeometry::K4V2L16 => None, + } + } + /// Dequantize the i-th expert's `[N, K_in]` BF16 weight matrix (3-D mode /// only). Internal use by `gather_forward` and friends; bails when called /// on a 2-D layer or with `expert_idx >= num_experts`. pub fn dequantize_expert(&self, expert_idx: usize) -> Result { + self.require_k4v2l16("dequantize_expert")?; let e = self.num_experts.ok_or_else(|| { candle_core::Error::Msg("QtipLayer::dequantize_expert called on a 2-D layer".into()) })?; @@ -4020,6 +4650,9 @@ impl QtipLayer { // Gaussian default above) — correct for any table, and never // assumes a computed codebook we cannot verify. codebook: QtipCodebook::Gaussian, + // Safetensors QTIP checkpoints predate every geometry but this + // one, and carry no field that could say otherwise. + geometry: QtipGeometry::K4V2L16, })) } @@ -4175,18 +4808,57 @@ impl QtipLayer { // meant. An unknown tag value, on the other hand, IS refused: that // means a newer Arc wrote it, and decoding its symbols against a // codebook we guessed would be silent corruption. - let codebook = match buffer.read_u8() { - Ok(tag) => QtipCodebook::from_wire(tag, || { - buffer.read_u32::().map_err(|_| { - candle_core::Error::Msg( - "QtipLayer: codebook tag claims a computed codebook but the multiplier \ - is missing (truncated payload)." - .into(), - ) - }) - })?, - Err(_) => QtipCodebook::Gaussian, - }; + // The trailing region is a sequence of tagged sections, read until EOF. + // Each tag may appear at most once: a repeat means a corrupt or + // concatenated payload, and silently letting the last one win is how a + // reader ends up decoding against something nobody wrote. + let mut codebook: Option = None; + let mut geometry: Option = None; + while let Ok(tag) = buffer.read_u8() { + if tag == GEOMETRY_WIRE_TAG { + if geometry.is_some() { + candle_core::bail!( + "QtipLayer: the UQFF payload carries two geometry sections. Refusing a \ + payload that describes itself twice." + ); + } + let mut klv = [0u8; 3]; + for (i, slot) in klv.iter_mut().enumerate() { + *slot = buffer.read_u8().map_err(|_| { + candle_core::Error::Msg(format!( + "QtipLayer: geometry section is truncated (got {i} of 3 K/L/V bytes)." + )) + })?; + } + geometry = Some(QtipGeometry::from_wire_body(klv[0], klv[1], klv[2])?); + } else { + if codebook.is_some() { + candle_core::bail!( + "QtipLayer: the UQFF payload carries two codebook sections. Refusing a \ + payload that describes itself twice." + ); + } + codebook = Some(QtipCodebook::from_wire(tag, || { + buffer.read_u32::().map_err(|_| { + candle_core::Error::Msg( + "QtipLayer: codebook tag claims a computed codebook but the \ + multiplier is missing (truncated payload)." + .into(), + ) + }) + })?); + } + } + // Absence is the answer for both, and it is the same answer it has + // always been: the historical geometry, gathering from the stored table. + let codebook = codebook.unwrap_or(QtipCodebook::Gaussian); + let geometry = geometry.unwrap_or(QtipGeometry::K4V2L16); + + // The tag is a claim; these make it true. Both supported geometries are + // 2 bits per weight, so `blocks` is the same size either way — the + // table's dtype and length are what actually discriminate. + geometry.validate_shapes(&blocks, &lut, in_features)?; + geometry.check_codebook(codebook)?; Ok(( Self { @@ -4201,6 +4873,7 @@ impl QtipLayer { search, search_detail, codebook, + geometry, }, ext_bias, )) @@ -4276,10 +4949,25 @@ impl QuantizedSerde for QtipLayer { if let Some(w) = beam_width { buffer.extend(&w.to_le_bytes()); } + // The trellis geometry, written FIRST in the trailing-section region + // and only when it is not the historical K=4/V=2/L=16 — see + // `QtipGeometry::to_wire` and `GEOMETRY_WIRE_TAG` for why the order is + // load-bearing. In short: a build that predates this field reads one + // trailing byte and hands it to the codebook parser, which refuses + // every tag it does not know. Tag 3 first means an old Arc fails + // closed on a geometry it cannot decode; tag 3 last would let it + // decode K=8/V=4 symbols as K=4/V=2 without faulting, because at 2 bpw + // the packed byte count is the same. + self.geometry.check_codebook(self.codebook)?; + if let Some((geo_tag, klv)) = self.geometry.to_wire() { + buffer.push(geo_tag); + buffer.extend(&klv); + } // wave24-AU: the codebook discriminator, appended LAST and only when // the codebook is not the historical Gaussian one — see - // `QtipCodebook::to_wire`. A Gaussian artifact is therefore byte-for- - // byte what this writer produced before the field existed. + // `QtipCodebook::to_wire`. A Gaussian artifact at the default geometry + // is therefore byte-for-byte what this writer produced before either + // field existed. // // A build that predates the field, handed a computed-codebook payload, // parses identically up to here and then ignores a trailing section it @@ -5818,6 +6506,7 @@ mod tests { let mut shared_rotation_block: usize = 0; let mut shared_search_detail = QtipSearchDetail::Unknown; let mut shared_codebook = QtipCodebook::Gaussian; + let mut shared_geometry = QtipGeometry::K4V2L16; for expert_idx in 0..e { let expert_w = w.narrow(0, expert_idx, 1)?.squeeze(0)?; let layer = QtipLayer::quantize_with_options_concrete( @@ -5835,6 +6524,7 @@ mod tests { shared_rotation_block = layer.rotation_block; shared_search_detail = layer.search_detail; shared_codebook = layer.codebook; + shared_geometry = layer.geometry; } } let blocks_3d = Tensor::stack(&blocks_slices, 0)?; @@ -5851,6 +6541,7 @@ mod tests { search: QtipSearchStamp::for_mode(mode), search_detail: shared_search_detail, codebook: shared_codebook, + geometry: shared_geometry, }) .inspect(|l| { debug_assert_eq!(l.blocks.dims(), &[e, n, k_in / 4]); @@ -6737,11 +7428,921 @@ mod tests { let bb: Vec = bv.iter().map(|x| x.to_bits()).collect(); assert_eq!(ab, bb, "{what}: f32 bit patterns differ"); } + DType::BF16 => { + // The K=8/V=4/L=12 reproduction table. Compared as raw bits, + // not as floats: a table is format, and `==` on bf16 would + // still call two NaN payloads unequal and two zeros equal. + let av: Vec = a.flatten_all()?.to_vec1()?; + let bv: Vec = b.flatten_all()?.to_vec1()?; + let ab: Vec = av.iter().map(|x| x.to_bits()).collect(); + let bb: Vec = bv.iter().map(|x| x.to_bits()).collect(); + assert_eq!(ab, bb, "{what}: bf16 bit patterns differ"); + } other => panic!("assert_tensor_bits_eq: unhandled dtype {other:?}"), } Ok(()) } + // =================================================================== + // Geometry discriminator (UQFF) + // =================================================================== + + /// The V=4/L=12 family at the byte-aligned control K, as a geometry value. + fn k8() -> QtipGeometry { + QtipGeometry::trellis_v4l12(8).unwrap() + } + + /// A cheap 2-D layer whose fields we can then rewrite for format tests. + fn geometry_fixture(device: &Device) -> Result { + let (n, k_in) = (4usize, 64usize); + let wdata: Vec = (0..(n * k_in)) + .map(|i| ((i as f32) * 0.37).sin() * 1.1) + .collect(); + let w = Tensor::from_vec(wdata, (n, k_in), device)?; + QtipLayer::quantize_with_options_concrete(&w, None, device, QtipMode::Greedy, false) + } + + /// The default geometry must be **purely additive**: it writes nothing, so + /// every artifact Arc has already produced is byte-identical and no + /// checksum moves. Stated as an exact suffix rather than as a vague + /// "unchanged", so the tag value, its position, and the K byte are all + /// pinned — **at every K in the family**. + /// + /// This is the property that makes a K change cheap: K is an explicit wire + /// field, so moving from the K=8 control to K=9 changes one byte of the + /// artifact and nothing else about the format. + #[test] + fn default_geometry_writes_nothing_and_the_family_appends_exactly_four_bytes() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + assert_eq!(layer.geometry, QtipGeometry::K4V2L16); + let base = layer.serialize()?.into_owned(); + + for k in trellis_v4l12::K_SUPPORTED { + // Flip ONLY the discriminator; every tensor stays byte-identical. + layer.geometry = QtipGeometry::trellis_v4l12(k)?; + let tagged = layer.serialize()?.into_owned(); + + let mut want = base.clone(); + want.extend_from_slice(&[GEOMETRY_WIRE_TAG, k as u8, 12, 4]); + assert_eq!( + tagged, want, + "K={k}: the geometry section must be exactly [tag=3, K, L=12, V=4] appended to \ + an otherwise unchanged payload" + ); + // The K byte is the ONLY thing that differs between family members. + assert_eq!(tagged.len(), base.len() + 4); + } + Ok(()) + } + + /// Two family members differ by exactly one byte on the wire. + /// + /// The concrete statement of "K is a parameter, not a variant": going from + /// the K=8 control to the K=9 quality winner is a single byte of artifact + /// change, given identical tensors. + #[test] + fn family_members_differ_by_one_wire_byte() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + layer.geometry = QtipGeometry::trellis_v4l12(8)?; + let a = layer.serialize()?.into_owned(); + layer.geometry = QtipGeometry::trellis_v4l12(9)?; + let b = layer.serialize()?.into_owned(); + assert_eq!(a.len(), b.len()); + let diffs: Vec = (0..a.len()).filter(|&i| a[i] != b[i]).collect(); + assert_eq!( + diffs.len(), + 1, + "K=8 and K=9 payloads must differ in exactly one byte, differed in {diffs:?}" + ); + assert_eq!((a[diffs[0]], b[diffs[0]]), (8, 9)); + Ok(()) + } + + /// The section must come BEFORE the codebook section, because that is what + /// makes an older Arc fail closed instead of mis-decoding. + #[test] + fn geometry_section_precedes_the_codebook_section() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + let base = layer.serialize()?.into_owned(); + + // Codebook alone (K=4 geometry, so Mcg is legal here). + layer.codebook = QtipCodebook::COMPUTED; + let cb_only = layer.serialize()?.into_owned(); + assert_eq!( + cb_only.len(), + base.len() + 5, + "codebook section is tag + u32" + ); + assert_eq!(cb_only[base.len()], 2, "codebook tag is 2"); + + // Both sections. Geometry (4 bytes) must land first. + // + // K8V4L12 + Mcg is refused as a pair, so use a hypothetical third + // geometry is not available — instead assert the writer's ORDER + // directly by checking that the geometry bytes occupy the slot the + // codebook otherwise would. + layer.codebook = QtipCodebook::Gaussian; + layer.geometry = k8(); + let geo_only = layer.serialize()?.into_owned(); + assert_eq!( + geo_only[base.len()], + GEOMETRY_WIRE_TAG, + "the geometry section must start immediately after the search-detail section — \ + any byte between them would be read as a codebook tag by an older Arc" + ); + Ok(()) + } + + /// An Arc that predates this field parses the trailing region by handing + /// its first byte to `QtipCodebook::from_wire`. That must REFUSE tag 3. + /// + /// If it ever accepted it, an old build would decode K=8/V=4 symbols as + /// K=4/V=2 — and would not fault, because at 2 bits per weight the packed + /// byte count is identical. + #[test] + fn an_older_reader_refuses_the_geometry_tag_rather_than_ignoring_it() { + let err = QtipCodebook::from_wire(GEOMETRY_WIRE_TAG, || Ok(1)) + .expect_err("tag 3 must not parse as a codebook"); + let msg = err.to_string(); + assert!( + msg.contains("geometry") || msg.contains("unknown codebook tag"), + "refusal must name the cause: {msg}" + ); + // And the general property the whole scheme rests on: every tag this + // build does not know is refused, never skipped. + for tag in [1u8, 4, 5, 200, 255] { + assert!( + QtipCodebook::from_wire(tag, || Ok(1)).is_err(), + "tag {tag} must fail closed" + ); + } + } + + /// A K=8/V=4 tag on a K=4/V=2 payload must be caught by the SHAPES, not + /// merely trusted. This is the case a hex-editor edit or a truncated + /// concatenation produces. + #[test] + fn a_geometry_tag_that_contradicts_the_tensors_is_refused_at_load() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + layer.geometry = k8(); + let data = layer.serialize()?.into_owned(); + + let err = QtipLayer::deserialize_concrete_unchecked( + Cow::Owned(data), + &device, + crate::QuantizeOntoGuard::new(), + ) + .expect_err("a K8V4L12 tag over a [65536, 2] F32 table must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("k8v4l12") && (msg.contains("F32") || msg.contains("values")), + "refusal must name the geometry and the mismatch: {msg}" + ); + Ok(()) + } + + /// The table's ELEMENT TYPE discriminates, not only its length. + /// + /// A `[4096, 4]` F32 table has exactly the right number of values and the + /// wrong dtype — 65,536 B instead of 32,768 B, so it would not even fit + /// the shared-memory budget the geometry exists for. Without this case the + /// element-count check alone passes and the dtype check is dead + /// (measured: mutation F8 disabled it and every test stayed green). + #[test] + fn a_right_sized_table_of_the_wrong_dtype_is_refused() -> Result<()> { + let device = Device::Cpu; + let geo = k8(); + let k_in = 64usize; + let blocks = Tensor::zeros((2, geo.packed_len(k_in)), DType::U8, &device)?; + let f32_table = Tensor::zeros( + (trellis_v4l12::LUT_STATES, trellis_v4l12::V as usize), + DType::F32, + &device, + )?; + assert_eq!(f32_table.elem_count(), geo.lut_values(), "same count"); + let err = geo + .validate_shapes(&blocks, &f32_table, k_in) + .expect_err("an F32 table at this geometry must be refused"); + assert!( + err.to_string().contains("F32") && err.to_string().contains("BF16"), + "refusal must name both dtypes: {err}" + ); + Ok(()) + } + + /// `from_stacked_parts` is a public constructor that takes the geometry on + /// trust from its caller; it must validate, not trust. + #[test] + fn from_stacked_parts_refuses_a_geometry_its_tensors_contradict() -> Result<()> { + let device = Device::Cpu; + let (e, n, k_in) = (2usize, 4usize, 64usize); + let blocks = Tensor::zeros((e, n, k_in / 4), DType::U8, &device)?; + let scales = Tensor::zeros((e, n), DType::F32, &device)?; + let k4_table = Tensor::zeros((1usize << L, V as usize), DType::F32, &device)?; + + // Truthful call: K=4 tensors, K=4 tag. + QtipLayer::from_stacked_parts( + blocks.clone(), + scales.clone(), + k4_table.clone(), + None, + k_in, + None, + 0, + QtipSearchStamp::Unstamped, + QtipSearchDetail::Unknown, + QtipCodebook::Gaussian, + QtipGeometry::K4V2L16, + )?; + + // Same tensors, K=8 tag. + let err = QtipLayer::from_stacked_parts( + blocks, + scales, + k4_table, + None, + k_in, + None, + 0, + QtipSearchStamp::Unstamped, + QtipSearchDetail::Unknown, + QtipCodebook::Gaussian, + k8(), + ) + .expect_err("a K8V4L12 tag over a K4V2L16 table must be refused at construction"); + assert!(err.to_string().contains("k8v4l12"), "{err}"); + Ok(()) + } + + /// A stack whose experts were baked at different geometries keeps only one + /// table, so at most one expert could ever decode correctly. + #[test] + fn stack_experts_refuses_mixed_geometries() -> Result<()> { + let device = Device::Cpu; + let a = geometry_fixture(&device)?; + let mut b = geometry_fixture(&device)?; + // Identical in every respect except the discriminator, so nothing but + // the geometry check can be what fires. + b.geometry = k8(); + let err = QtipLayer::stack_experts(vec![a, b]) + .expect_err("mixed-geometry stacks must be refused"); + assert!( + err.to_string().contains("geometry"), + "refusal must name the geometry: {err}" + ); + + // ...and a uniform stack still works, so the guard is not just + // rejecting everything. + let c = geometry_fixture(&device)?; + let d = geometry_fixture(&device)?; + QtipLayer::stack_experts(vec![c, d])?; + Ok(()) + } + + /// Every decode entry point that has NO kernel for a geometry must refuse + /// it, rather than falling through to the K=4/V=2 decoder. + /// + /// The decoders behind these unpack nibbles and index a `[2^16, 2]` table, + /// and handed a V=4 row they would NOT fault — at 2 bpw the packed byte + /// count is identical and every index lands in bounds. They would serve + /// garbage. + /// + /// **Stage 3 narrowed this test on purpose.** `dequantize_weights` and + /// `forward` are now wired for the family (CPU decode, plus a CUDA + /// single-token fused gemv), so they are covered by + /// `a_k9_layer_forwards_instead_of_refusing` instead. What is left here is + /// what still has no kernel — and a half-wired dispatch that falls through + /// is worse than one that refuses, so this stays. + #[test] + fn unwired_entry_points_refuse_an_unsupported_geometry() -> Result<()> { + let device = Device::Cpu; + let k_in = 64usize; + + // Sanity at the geometry this build implements everywhere, so the + // refusals below are about the geometry and nothing else. + let two_d = geometry_fixture(&device)?; + two_d.dequantize_weights()?; + assert!(two_d.qtip_packed().is_some()); + + let mut stack = + QtipLayer::stack_experts(vec![geometry_fixture(&device)?, geometry_fixture(&device)?])?; + stack.dequantize_expert(0)?; + let a = Tensor::zeros((1, 1, k_in), DType::F32, &device)?; + let idx = Tensor::zeros((1, 1), DType::U32, &device)?; + stack.gather_forward(&a, &idx)?; + + // Now the family, on paths with no kernel. + stack.geometry = QtipGeometry::trellis_v4l12(8)?; + for (what, res) in [ + ("dequantize_expert", stack.dequantize_expert(0).err()), + ("gather_forward", stack.gather_forward(&a, &idx).err()), + ] { + let err = res.unwrap_or_else(|| panic!("{what} must refuse a k8v4l12 layer")); + let msg = err.to_string(); + assert!(msg.contains("k8v4l12"), "{what}: {msg}"); + assert!( + msg.contains(&format!("QtipLayer::{what}:")), + "{what}: refusal came from elsewhere — an entry point that inherits a callee's \ + guard breaks the moment the callee gains a fast path: {msg}" + ); + } + + // `qtip_packed` carries no geometry, so it must hand out nothing. + let mut packed_view = geometry_fixture(&device)?; + packed_view.geometry = QtipGeometry::trellis_v4l12(8)?; + assert!( + packed_view.qtip_packed().is_none(), + "QtipPackedView has no geometry field; a consumer would read the bytes as K=4/V=2 \ + nibbles" + ); + Ok(()) + } + + /// Build a K=9 layer whose weights are known, so a forward can be checked + /// against an independent decode rather than against itself. + fn k9_layer(device: &Device, n: usize, k_in: usize) -> Result<(QtipLayer, Vec, Vec)> { + let geo = QtipGeometry::trellis_v4l12(9)?; + let rung = trellis_v4l12::Rung::new(9).map_err(candle_core::Error::Msg)?; + let n_sym = rung.num_symbols(k_in); + let row_bytes = rung.row_stride(n_sym); + + let mut packed = Vec::with_capacity(n * row_bytes); + for row in 0..n { + let mut z = (row as u64 + 1).wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + let syms: Vec = (0..n_sym) + .map(|_| { + z ^= z << 13; + z ^= z >> 7; + z ^= z << 17; + ((z >> 24) as u32) & 0x1FF + }) + .collect(); + packed.extend(rung.pack(&syms)); + } + let scales: Vec = (0..n).map(|i| 0.013 + i as f32 * 0.002).collect(); + + let layer = QtipLayer { + blocks: Tensor::from_vec(packed.clone(), (n, row_bytes), device)?, + row_scales: Tensor::from_vec(scales.clone(), (n,), device)?, + lut: Tensor::from_slice( + &trellis_v4l12::gaussian_lut_bf16(), + (trellis_v4l12::LUT_STATES, trellis_v4l12::V as usize), + device, + )?, + bias: None, + in_features: k_in, + num_experts: None, + rotation_signs: None, + rotation_block: 0, + search: QtipSearchStamp::Trellis, + search_detail: QtipSearchDetail::Known { + beam_width: None, + hessian: false, + }, + codebook: QtipCodebook::Gaussian, + geometry: geo, + }; + Ok((layer, packed, scales)) + } + + /// **Stage 3, the point of it: a K=9 layer is SERVED, not refused.** + /// + /// Checked against an independent decode through `trellis_v4l12::Rung`, + /// not against the layer's own dequantize — otherwise a wrong decoder + /// would agree with itself. + #[test] + fn a_k9_layer_forwards_instead_of_refusing() -> Result<()> { + let device = Device::Cpu; + let (n, k_in) = (6usize, 128usize); + let (layer, packed, scales) = k9_layer(&device, n, k_in)?; + + let w = layer.dequantize_weights()?.to_dtype(DType::F32)?; + assert_eq!(w.dims(), &[n, k_in]); + let got: Vec = w.flatten_all()?.to_vec1()?; + + let rung = trellis_v4l12::Rung::new(9).unwrap(); + let n_sym = rung.num_symbols(k_in); + let row_bytes = rung.row_stride(n_sym); + let lut = trellis_v4l12::gaussian_lut_bf16(); + for row in 0..n { + let mut want = vec![0f32; k_in]; + rung.decode_row_into( + trellis_v4l12::Row::new( + &packed[row * row_bytes..(row + 1) * row_bytes], + n_sym, + scales[row], + ), + &lut, + &mut want, + ) + .unwrap(); + for i in 0..k_in { + // dequantize_weights returns BF16, so compare at BF16 precision. + let g = got[row * k_in + i]; + let e = half::bf16::from_f32(want[i]).to_f32(); + assert_eq!(g, e, "row {row} col {i}"); + } + } + + // ...and a real forward runs and matches x @ W^T. F32 on CPU: candle's + // CPU matmul has no BF16 kernel, which is a property of the shipped + // K=4 path too, not of this geometry. + let x: Vec = (0..k_in).map(|i| ((i as f32) * 0.07).sin()).collect(); + let x_t = Tensor::from_vec(x, (1, k_in), &device)?; + let y = layer.forward(&x_t)?.to_dtype(DType::F32)?; + assert_eq!(y.dims(), &[1, n]); + let y_ref = x_t.matmul(&w.t()?)?; + let (yv, rv): (Vec, Vec) = + (y.flatten_all()?.to_vec1()?, y_ref.flatten_all()?.to_vec1()?); + for (a, b) in yv.iter().zip(rv.iter()) { + let denom = a.abs().max(b.abs()).max(1e-3); + assert!((a - b).abs() / denom < 1e-2, "forward {a} vs reference {b}"); + } + Ok(()) + } + + /// A K=9 forward must NOT be the K=4 decoder quietly succeeding. + /// + /// Both geometries decode the same bytes without faulting, so "it returned + /// numbers" proves nothing. Decode the same packed bytes as K=4/V=2 and + /// require the answers to differ. + #[test] + fn the_k9_path_is_not_the_k4_decoder_in_disguise() -> Result<()> { + let device = Device::Cpu; + let (n, k_in) = (4usize, 128usize); + let (layer, _, _) = k9_layer(&device, n, k_in)?; + let k9_w: Vec = layer + .dequantize_weights()? + .to_dtype(DType::F32)? + .flatten_all()? + .to_vec1()?; + + // Same bytes, same scales, read as the shipped rung. A K=9 row is + // longer than a K=4 row of the same in_features, so trim to what the + // K=4 decoder expects — the point is that the SYMBOLS differ, not the + // lengths. + let k4_bytes = QtipGeometry::K4V2L16.packed_len(k_in); + let blocks: Vec = layer.blocks.flatten_all()?.to_vec1()?; + let row_bytes = blocks.len() / n; + let mut trimmed = Vec::with_capacity(n * k4_bytes); + for row in 0..n { + trimmed.extend_from_slice(&blocks[row * row_bytes..row * row_bytes + k4_bytes]); + } + let mut as_k4 = layer; + as_k4.geometry = QtipGeometry::K4V2L16; + as_k4.blocks = Tensor::from_vec(trimmed, (n, k4_bytes), &device)?; + as_k4.lut = Tensor::from_vec(gaussian_lut(), (1usize << L, V as usize), &device)?; + let k4_w: Vec = as_k4 + .dequantize_weights()? + .to_dtype(DType::F32)? + .flatten_all()? + .to_vec1()?; + + assert_eq!(k9_w.len(), k4_w.len()); + let differing = k9_w + .iter() + .zip(k4_w.iter()) + .filter(|(a, b)| a.to_bits() != b.to_bits()) + .count(); + assert!( + differing > k9_w.len() / 2, + "only {differing}/{} weights differ between the K=9 and K=4 readings of the same \ + bytes — the K=9 path may be the K=4 decoder", + k9_w.len() + ); + Ok(()) + } + + /// Paths the family has no kernel for must still refuse loudly, on a + /// WELL-FORMED family layer. A half-wired dispatch that falls through to a + /// K=4 path is worse than one that refuses. + #[test] + fn unwired_paths_still_refuse_for_a_well_formed_k9_layer() -> Result<()> { + let device = Device::Cpu; + let (layer, _, _) = k9_layer(&device, 4, 128)?; + + assert!( + layer.qtip_packed().is_none(), + "QtipPackedView carries no geometry; a consumer would read K=9 bytes as nibbles" + ); + + let mut three_d = layer; + three_d.num_experts = Some(2); + for (what, err) in [ + ("dequantize_expert", three_d.dequantize_expert(0).err()), + ( + "gather_forward", + three_d + .gather_forward( + &Tensor::zeros((1, 1, 128), DType::F32, &device)?, + &Tensor::zeros((1, 1), DType::U32, &device)?, + ) + .err(), + ), + ( + "dequantize_weights(3-D)", + three_d.dequantize_weights().err(), + ), + ] { + let e = err.unwrap_or_else(|| panic!("{what} must refuse a k9v4l12 layer")); + let msg = e.to_string(); + assert!( + msg.contains("k9v4l12"), + "{what}: refusal must name the geometry: {msg}" + ); + assert!( + msg.to_lowercase().contains("refus"), + "{what}: refusal must say it is refusing: {msg}" + ); + } + Ok(()) + } + + /// The two dispatch invariants that live behind `#[cfg(feature = "cuda")]` + /// and therefore cannot be exercised in a no-GPU lane. + /// + /// Both were mutated and both passed (S2, S5): a CUDA build is the only + /// thing that compiles them, and this crate's CI CUDA lane does not run + /// tests. So assert them over the source, the same way the kernel's + /// constants are asserted — a guard that only exists on hardware nobody + /// has is not a guard. + #[test] + fn the_cuda_dispatch_invariants_are_pinned_over_the_source() { + const SRC: &str = include_str!("mod.rs"); + fn squeeze(s: &str) -> String { + s.split_whitespace().collect::>().join(" ") + } + // SERVING CODE ONLY. `include_str!` pulls in this test module too, and + // the negative assertion below names the very string it forbids — so + // searching the whole file makes the guard fire on itself. (Second time + // this class of bug appeared here; the first was a barrier guard + // counting the word "__syncthreads" in a comment. A guard that matches + // its own text is not a guard.) + // Normalise line endings first: on a Windows checkout `include_str!` + // yields CRLF, so a split pattern carrying a bare `\n` never matches + // and the guard dies in `expect` instead of asserting anything. Every + // assertion below goes through `squeeze`, which is already + // newline-agnostic — this split was the one place that was not. + let src_lf = SRC.replace("\r\n", "\n"); + let serving = src_lf + .split_once("#[cfg(test)]\nmod tests {") + .map(|(before, _)| before) + .expect("mod.rs must contain its test module"); + let src = squeeze(serving); + + // A family layer must never reach the K=4/V=2 CUDA path, whose + // dequantize kernel and fused gemv both index a [2^16, 2] F32 table. + assert!( + src.contains(&squeeze( + "if self.v4l12_rung().is_none() && cuda_ops::can_use_qtip_cuda(&self.blocks) {" + )), + "the K=4/V=2 CUDA forward path lost its geometry gate. A V=4 layer reaching it \ + would be decoded as K=4/V=2 — in bounds, no fault, wrong weights." + ); + + // The family's CUDA path must ship with the row-scale hoist OFF: it is + // the one lever that costs bit-exactness against the CPU reference, and + // nothing here has been measured on a device. + assert!( + src.contains(&squeeze("trellis_v4l12::RowScaleHoist::Off,")), + "the family's CUDA gemv no longer passes RowScaleHoist::Off" + ); + assert!( + !src.contains(&squeeze("trellis_v4l12::RowScaleHoist::On")), + "something in the serving path now asks for the row-scale hoist. It reassociates \ + the sum and breaks bit-exactness against the CPU reference; there is no measured \ + device-time evidence to trade that for." + ); + } + + /// **The new geometry must not be the default.** Nothing in the family has + /// touched a GPU; the kernel has never been compiled at any K. + #[test] + fn the_new_geometry_is_not_the_default() { + assert_eq!(QtipGeometry::DEFAULT, QtipGeometry::K4V2L16); + assert_eq!(QtipGeometry::default(), QtipGeometry::K4V2L16); + let device = Device::Cpu; + let layer = geometry_fixture(&device).unwrap(); + assert_eq!(layer.geometry, QtipGeometry::K4V2L16); + } + + /// The selection surface: `ARC_QTIP_GEOMETRY`, mirroring + /// `ARC_QTIP_CODEBOOK`. Refuses rather than guessing. + #[test] + fn the_geometry_env_surface_parses_and_refuses() { + // Parsing is exercised through the same rule the reader uses rather + // than by mutating the process environment, which races other tests in + // this binary. + // THE REAL RULE, not a replica of it. An earlier version of this test + // reimplemented the parse inline and passed while the real reader + // silently resolved an unsupported K to the default (mutation S7). + let parse = QtipGeometry::from_spec; + assert_eq!(parse("").unwrap(), QtipGeometry::K4V2L16); + assert_eq!(parse("k4v2l16").unwrap(), QtipGeometry::K4V2L16); + assert_eq!(parse("k9v4l12").unwrap().k(), 9); + assert_eq!(parse("k8v4l12").unwrap().k(), 8); + // A K with no decoder is refused, not clamped to a neighbour. + assert!(parse("k11v4l12").is_err()); + assert!(parse("qtip9").is_err()); + assert!(parse("k9").is_err()); + } + + /// A genuine V=4/L=12 payload round-trips at every K, tag and tensors both. + #[test] + fn trellis_v4l12_payloads_round_trip_exactly_at_every_k() -> Result<()> { + let device = Device::Cpu; + let (n, k_in) = (4usize, 128usize); + let lut = Tensor::from_slice( + &trellis_v4l12::gaussian_lut_bf16(), + (trellis_v4l12::LUT_STATES, trellis_v4l12::V as usize), + &device, + )?; + + for k in trellis_v4l12::K_SUPPORTED { + let geo = QtipGeometry::trellis_v4l12(k)?; + let packed_per_row = geo.packed_len(k_in); + // Sanity: the DATA length tracks K, and the allocated stride is + // that plus the padding the unclamped extraction relies on. + assert_eq!(geo.data_bytes(k_in), (k_in / 4 * k as usize).div_ceil(8)); + assert!(packed_per_row >= geo.data_bytes(k_in)); + assert!(packed_per_row - geo.data_bytes(k_in) <= 4); + + let blocks_data: Vec = (0..(n * packed_per_row)) + .map(|i| (i * 37 % 251) as u8) + .collect(); + let blocks = Tensor::from_vec(blocks_data, (n, packed_per_row), &device)?; + let scales = Tensor::from_vec( + (0..n).map(|i| 0.01 + i as f32 * 0.003).collect::>(), + (n,), + &device, + )?; + + let layer = QtipLayer { + blocks, + row_scales: scales, + lut: lut.clone(), + bias: None, + in_features: k_in, + num_experts: None, + rotation_signs: None, + rotation_block: 0, + search: QtipSearchStamp::Trellis, + search_detail: QtipSearchDetail::Known { + beam_width: None, + hessian: false, + }, + codebook: QtipCodebook::Gaussian, + geometry: geo, + }; + + let data = layer.serialize()?.into_owned(); + let (restored, _) = QtipLayer::deserialize_concrete_unchecked( + Cow::Owned(data), + &device, + crate::QuantizeOntoGuard::new(), + )?; + + assert_eq!(restored.geometry, geo, "K={k}"); + assert_eq!(restored.geometry.k(), k); + assert_eq!(restored.in_features, k_in); + assert_eq!(restored.lut.dtype(), DType::BF16); + assert_eq!(restored.lut.elem_count(), trellis_v4l12::LUT_ENTRIES); + assert_tensor_bits_eq(&layer.blocks, &restored.blocks, "blocks")?; + assert_tensor_bits_eq(&layer.row_scales, &restored.row_scales, "row_scales")?; + assert_tensor_bits_eq(&layer.lut, &restored.lut, "lut")?; + } + Ok(()) + } + + /// A K the wire names but this build has no decoder for is refused, and the + /// refusal quotes the triple. + #[test] + fn an_in_family_but_unsupported_k_is_refused_with_a_diagnosis() -> Result<()> { + let device = Device::Cpu; + let layer = geometry_fixture(&device)?; + let base = layer.serialize()?.into_owned(); + // L=12/V=4 is this family, but K=11 has no decoder. + let mut bytes = base.clone(); + bytes.extend_from_slice(&[GEOMETRY_WIRE_TAG, 11, 12, 4]); + let err = QtipLayer::deserialize_concrete_unchecked( + Cow::Owned(bytes), + &device, + crate::QuantizeOntoGuard::new(), + ) + .expect_err("K=11 must be refused"); + assert!( + err.to_string().contains("K=11") || err.to_string().contains("not implemented"), + "refusal must name the K it could not decode: {err}" + ); + Ok(()) + } + + /// The computed `sum2` codebook produces a PAIR of values per state and so + /// cannot describe V=4. Refused on write and on read, not left to fail at + /// the first decode. + #[test] + fn the_computed_codebook_is_refused_at_a_v4_geometry() -> Result<()> { + let device = Device::Cpu; + let mut layer = geometry_fixture(&device)?; + layer.geometry = k8(); + layer.codebook = QtipCodebook::COMPUTED; + let err = layer + .serialize() + .expect_err("V=4 + sum2 must not serialize"); + assert!( + err.to_string().contains("V=2-only"), + "refusal must say why: {err}" + ); + // And it is legal at the geometry it was built for, so the guard is + // about the pairing and not about the codebook being blocked outright. + layer.geometry = QtipGeometry::K4V2L16; + layer.serialize()?; + Ok(()) + } + + /// Malformed trailing regions are refused rather than partially read. + #[test] + fn malformed_geometry_sections_are_refused() -> Result<()> { + let device = Device::Cpu; + let layer = geometry_fixture(&device)?; + let base = layer.serialize()?.into_owned(); + let load = |bytes: Vec| { + QtipLayer::deserialize_concrete_unchecked( + Cow::Owned(bytes), + &device, + crate::QuantizeOntoGuard::new(), + ) + }; + + // Truncated: tag present, K/L/V missing. + let mut truncated = base.clone(); + truncated.push(GEOMETRY_WIRE_TAG); + truncated.push(8); + let err = load(truncated).expect_err("truncated geometry section must be refused"); + assert!(err.to_string().contains("truncated"), "{err}"); + + // Well-formed but unsupported triple: diagnosed, not "unknown tag". + let mut unsupported = base.clone(); + unsupported.extend_from_slice(&[GEOMETRY_WIRE_TAG, 6, 18, 3]); + let err = load(unsupported).expect_err("unsupported geometry must be refused"); + let msg = err.to_string(); + assert!( + msg.contains("K=6") && msg.contains("L=18") && msg.contains("V=3"), + "refusal must quote the triple it could not decode: {msg}" + ); + + // Two geometry sections: a payload that describes itself twice. + let mut doubled = base.clone(); + doubled.extend_from_slice(&[GEOMETRY_WIRE_TAG, 4, 16, 2]); + doubled.extend_from_slice(&[GEOMETRY_WIRE_TAG, 4, 16, 2]); + let err = load(doubled).expect_err("duplicate geometry sections must be refused"); + assert!(err.to_string().contains("twice"), "{err}"); + + // An explicit default-geometry section is legal (a future writer may + // choose to be explicit) and must load as K4V2L16. + let mut explicit = base.clone(); + explicit.extend_from_slice(&[GEOMETRY_WIRE_TAG, 4, 16, 2]); + let (restored, _) = load(explicit)?; + assert_eq!(restored.geometry, QtipGeometry::K4V2L16); + Ok(()) + } + + /// The shipped rung and the K=8 control are the SAME bit rate and the SAME + /// packed size. That is the fact that makes the discriminator necessary — + /// a mislabelled artifact has a correctly-sized `blocks` tensor — so it + /// gets an assertion rather than a comment. + #[test] + fn k4v2l16_and_k8v4l12_are_indistinguishable_by_packed_size() { + let k8 = k8(); + for k_in in [64usize, 512, 4096, 7168] { + assert_eq!( + QtipGeometry::K4V2L16.packed_len(k_in), + k8.packed_len(k_in), + "k_in={k_in}: identical packed size is why the tag is load-bearing" + ); + } + assert_eq!(QtipGeometry::K4V2L16.bpw_x100(), 200); + assert_eq!(k8.bpw_x100(), 200); + // ...and the tables are what actually differ. + assert_eq!(QtipGeometry::K4V2L16.lut_values(), 65_536 * 2); + assert_eq!(k8.lut_values(), 4_096 * 4); + assert_eq!(QtipGeometry::K4V2L16.lut_dtype(), DType::F32); + assert_eq!(k8.lut_dtype(), DType::BF16); + } + + /// **Not every geometry in the enum is 2 bits per weight any more.** + /// + /// K=9/V=4 is 2.25 bpw, so its rows are 12.5% larger than K=8/V=4's and it + /// is *not* size-confusable with the shipped rung. Worth pinning: code that + /// assumed "all QTIP geometries are 2 bpw" — a true statement when this + /// family was K=8-only — is now wrong, and `packed_len` is the general + /// `ceil(n·K/8)` rather than anything that divides by `8/K`. + #[test] + fn the_family_spans_more_than_one_bit_rate() { + let k8 = k8(); + let k9 = QtipGeometry::trellis_v4l12(9).unwrap(); + let k10 = QtipGeometry::trellis_v4l12(10).unwrap(); + assert_eq!(k8.bpw_x100(), 200); + assert_eq!(k9.bpw_x100(), 225); + assert_eq!(k10.bpw_x100(), 250); + for k_in in [512usize, 4096, 7168] { + assert!(k9.packed_len(k_in) > k8.packed_len(k_in)); + // Exactly the bit-rate ratio, to the byte — on DATA bytes. The + // stride is only approximately in that ratio, because K=9 pads for + // its multi-byte extraction and K=8 does not. + assert_eq!(k9.data_bytes(k_in) * 8, k8.data_bytes(k_in) * 9); + // The whole family shares one table regardless of K. + assert_eq!(k9.lut_values(), k8.lut_values()); + assert_eq!(k9.lut_dtype(), k8.lut_dtype()); + } + } + + /// The packed length must use the CEIL, and that is only visible at an + /// `in_features` which is not a multiple of 32. + /// + /// Every plausible layer width is a multiple of 32, which makes + /// `num_symbols · K` a multiple of 8 and floor equal to ceil at K=9. A + /// floored formula therefore passes every realistic fixture (measured: + /// mutation W3). These widths are deliberately unrealistic. + #[test] + fn packed_len_uses_the_ceiling_where_it_is_observable() { + let k9 = QtipGeometry::trellis_v4l12(9).unwrap(); + // in_features=36 -> 9 symbols -> 81 bits -> 11 DATA bytes, not 10. + assert_eq!(k9.num_symbols(36), 9); + assert_eq!(k9.data_bytes(36), 11); + // in_features=4 -> 1 symbol -> 9 bits -> 2 data bytes, not 1. + assert_eq!(k9.data_bytes(4), 2); + let k10 = QtipGeometry::trellis_v4l12(10).unwrap(); + // 3 symbols -> 30 bits -> 4 data bytes, not 3. + assert_eq!(k10.data_bytes(12), 4); + // The stride is the data plus tail padding, rounded up to 4. + assert_eq!(k9.packed_len(36), 12); + assert_eq!(k9.packed_len(4), 4); + assert_eq!(k10.packed_len(12), 8); + // ...and both agree with the rung that owns the formulas. + for k in trellis_v4l12::K_SUPPORTED { + let g = QtipGeometry::trellis_v4l12(k).unwrap(); + let r = trellis_v4l12::Rung::new(k).unwrap(); + for k_in in [4usize, 12, 36, 100, 260, 4096] { + let n = r.num_symbols(k_in); + assert_eq!( + g.packed_len(k_in), + r.row_stride(n), + "K={k} k_in={k_in} stride" + ); + assert_eq!( + g.data_bytes(k_in), + r.data_bytes(n), + "K={k} k_in={k_in} data" + ); + } + } + } + + /// A K with no decoder is refused at construction, not stored and used. + #[test] + fn an_unsupported_k_cannot_become_a_geometry() { + for k in [0u32, 4, 7, 11, 16] { + assert!( + QtipGeometry::trellis_v4l12(k).is_err(), + "K={k} must not be constructible" + ); + } + for k in trellis_v4l12::K_SUPPORTED { + assert!(QtipGeometry::trellis_v4l12(k).is_ok(), "K={k} must be"); + } + } + + /// The geometry accessors must agree with the rung they describe. + #[test] + fn geometry_accessors_match_the_rung_module() { + assert_eq!(QtipGeometry::K4V2L16.k(), K); + assert_eq!(QtipGeometry::K4V2L16.l(), L); + assert_eq!(QtipGeometry::K4V2L16.v(), V); + for k in trellis_v4l12::K_SUPPORTED { + let g = QtipGeometry::trellis_v4l12(k).unwrap(); + let r = trellis_v4l12::Rung::new(k).unwrap(); + assert_eq!(g.k(), r.k()); + assert_eq!(g.l(), trellis_v4l12::L); + assert_eq!(g.v(), trellis_v4l12::V); + assert_eq!(g.bpw_x100(), r.bpw_x100()); + assert_eq!( + g.lut_values(), + trellis_v4l12::LUT_ENTRIES, + "the discriminator and the rung must agree on the table size" + ); + for k_in in [64usize, 512, 4096] { + assert_eq!(g.packed_len(k_in), r.packed_len(k_in), "K={k} k_in={k_in}"); + assert_eq!(g.num_symbols(k_in), r.num_symbols(k_in)); + } + assert_eq!(g.tag(), format!("k{k}v4l12")); + } + } + /// UQFF round-trip for the classic 2-D layout must stay lossless and /// byte-compatible (bias branch + no-rotation branch covered). #[test] diff --git a/mistralrs-quant/src/turboquant/cuda_tables.rs b/mistralrs-quant/src/turboquant/cuda_tables.rs new file mode 100644 index 000000000..e50c0c05b --- /dev/null +++ b/mistralrs-quant/src/turboquant/cuda_tables.rs @@ -0,0 +1,225 @@ +//! The contract between the Rust TurboQuant compressor and the CUDA kernels. +//! +//! The kernels cannot call [`super::wht::generate_signs`] or +//! [`super::codebook::get_codebook`] at run time, so the sign diagonals and +//! Lloyd–Max codebooks are checked into +//! `mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh` as `__constant__` +//! arrays. Two copies of the same numbers is a drift hazard: if the CUDA sign +//! table stops matching the Rust one, every compressed head decodes to noise, +//! and it decodes to noise *quickly* — a broken kernel reports a better tok/s, +//! not a worse one. +//! +//! [`turboquant_cuda_tables_match_the_generator`] closes that. It parses the +//! header and re-derives every table from the same functions the compressor +//! uses. It is a CPU test about a *format*, not a claim about kernel behaviour +//! on a GPU. + +/// Head dimensions the CUDA kernels are instantiated for. +/// +/// Mirrors the `case` arms of `TQ_ATTN_BY_HS` / `TQ_CACHE_BY_HEAD_DIM` in +/// `turbo_paged_attention.cu`, and [`cuda_head_dims_match_the_kernel_dispatch`] +/// pins the two together. Callers gate on this before dispatching so an +/// unsupported width becomes a refusal rather than an untouched output buffer. +/// +/// This is narrower than what the *format* supports: +/// [`super::layout::TurboQuantLayout`] handles any width `>= MIN_BLOCK_DIM` by +/// block decomposition. The kernels take only the widths that decompose to a +/// single power-of-two block, which is every entry here. +pub const TURBOQUANT_CUDA_HEAD_DIMS: [usize; 4] = [64, 128, 256, 512]; + +/// Whether the CUDA kernels have an instantiation for `head_dim`. +pub fn cuda_supports_head_dim(head_dim: usize) -> bool { + TURBOQUANT_CUDA_HEAD_DIMS.contains(&head_dim) +} + +/// The seed the kernels' sign tables were generated with. Matches +/// [`super::TurboQuantConfig`]'s default. +pub const CUDA_TABLE_SEED: u64 = 42; + +#[cfg(test)] +mod tests { + use super::*; + use crate::turboquant::codebook::get_codebook; + use crate::turboquant::wht::generate_signs; + + const HEADER: &str = + include_str!("../../../mistralrs-paged-attn/src/cuda/turbo_paged_attention.cuh"); + const KERNEL: &str = + include_str!("../../../mistralrs-paged-attn/src/cuda/turbo_paged_attention.cu"); + /// The FFI wrapper's own copy of the list. `mistralrs-paged-attn` has no + /// dependency on this crate (and should not grow one just to share a + /// four-element array), so the two are kept equal by this test instead. + const FFI_WRAPPER: &str = + include_str!("../../../mistralrs-paged-attn/src/cuda/backend/paged_attention.rs"); + + /// Pull `static __constant__ float NAME[N] = { ... };` out of the header. + fn table(name: &str) -> Vec { + let needle = format!("float {name}["); + let start = HEADER + .find(&needle) + .unwrap_or_else(|| panic!("{name} is not declared in turbo_paged_attention.cuh")); + let open = HEADER[start..].find('{').expect("no opening brace") + start; + let close = HEADER[open..].find("};").expect("no closing brace") + open; + HEADER[open + 1..close] + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(|s| { + s.trim_end_matches('f') + .parse::() + .unwrap_or_else(|e| panic!("{name}: cannot parse {s:?}: {e}")) + }) + .collect() + } + + fn assert_same(name: &str, cuda: &[f32], rust: &[f32]) { + assert_eq!( + cuda.len(), + rust.len(), + "{name}: CUDA table has {} entries, the generator produces {}", + cuda.len(), + rust.len() + ); + for (i, (c, r)) in cuda.iter().zip(rust).enumerate() { + assert!( + (c - r).abs() < 1e-9, + "{name}[{i}]: CUDA has {c}, the generator produces {r}. \ + The tables in turbo_paged_attention.cuh are generated — re-run \ + `cargo run -p mistralrs-quant --example emit_turboquant_cuda_tables -- \ + 64 128 256 512` rather than editing them by hand." + ); + } + } + + /// Every `__constant__` table in the CUDA header equals what the Rust + /// compressor would produce at that dimension. + /// + /// This is the guard that makes the two implementations one format. Without + /// it, the four hand-copied sign tables in the tree (`.cu`, `.cuh`, + /// `turbo_wht.cu`, `.metal`) had nothing asserting they agreed with the + /// generator or with each other. + #[test] + fn turboquant_cuda_tables_match_the_generator() { + for dim in TURBOQUANT_CUDA_HEAD_DIMS { + assert_same( + &format!("TQ_SGN_{dim}"), + &table(&format!("TQ_SGN_{dim}")), + &generate_signs(CUDA_TABLE_SEED, dim), + ); + for bits in [3u32, 4u32] { + let cb = get_codebook(dim, bits); + assert_same( + &format!("TQ_CB{bits}_{dim}"), + &table(&format!("TQ_CB{bits}_{dim}")), + cb.centroids, + ); + assert_same( + &format!("TQ_BD{bits}_{dim}"), + &table(&format!("TQ_BD{bits}_{dim}")), + cb.boundaries, + ); + } + } + } + + /// A sign table that is all `+1` is not a rotation at all — it would make + /// `D·H·D` collapse to a plain Hadamard transform and quietly change the + /// coordinate distribution the codebooks assume. Guards against a table + /// being zeroed or stubbed rather than merely wrong. + #[test] + fn cuda_sign_tables_are_actually_mixed() { + for dim in TURBOQUANT_CUDA_HEAD_DIMS { + let s = table(&format!("TQ_SGN_{dim}")); + let neg = s.iter().filter(|&&x| x < 0.0).count(); + assert!( + neg > dim / 8 && neg < dim - dim / 8, + "TQ_SGN_{dim} has {neg}/{dim} negative signs — that is not a random diagonal" + ); + assert!( + s.iter().all(|&x| x == 1.0 || x == -1.0), + "TQ_SGN_{dim} contains a value that is not ±1" + ); + } + } + + /// The kernel dispatch and [`TURBOQUANT_CUDA_HEAD_DIMS`] must agree. + /// + /// If a `case` is added to the CUDA switch without being added here the + /// Rust gate rejects a width the kernel could serve; if it is removed from + /// CUDA but left here, a launch falls through the switch and leaves the + /// output buffer untouched. That second failure mode is silent, which is + /// why this test exists rather than a comment. + #[test] + fn cuda_head_dims_match_the_kernel_dispatch() { + for dim in TURBOQUANT_CUDA_HEAD_DIMS { + assert!( + KERNEL.contains(&format!("TQ_ATTN_BY_BS({dim}, ")), + "head dim {dim} is in TURBOQUANT_CUDA_HEAD_DIMS but TQ_ATTN_BY_HS \ + in turbo_paged_attention.cu has no case for it" + ); + assert!( + KERNEL.contains(&format!("TQ_CACHE_BY_DTYPE(FN, {dim},")), + "head dim {dim} is in TURBOQUANT_CUDA_HEAD_DIMS but TQ_CACHE_BY_HEAD_DIM \ + in turbo_paged_attention.cu has no case for it" + ); + } + // And nothing the kernel dispatches is missing from the Rust list. + for dim in [64usize, 128, 256, 512, 1024] { + let in_kernel = KERNEL.contains(&format!("TQ_ATTN_BY_BS({dim}, ")); + assert_eq!( + in_kernel, + cuda_supports_head_dim(dim), + "head dim {dim}: kernel dispatch says {in_kernel}, \ + TURBOQUANT_CUDA_HEAD_DIMS says {}", + cuda_supports_head_dim(dim) + ); + } + } + + /// The Rust FFI wrapper refuses exactly the widths the kernel cannot serve. + /// + /// The wrapper is the last gate before the launch, so if its list is wider + /// than the kernel's the caller reads an uninitialized output buffer. + #[test] + fn ffi_wrapper_head_dims_match() { + let start = FFI_WRAPPER + .find("pub const TURBO_HEAD_DIMS") + .expect("TURBO_HEAD_DIMS is not declared in the paged-attn FFI wrapper"); + let open = FFI_WRAPPER[start..].find('[').unwrap() + start; + let open = FFI_WRAPPER[open + 1..].find('[').unwrap() + open + 1; // skip the [usize; N] type + let close = FFI_WRAPPER[open..].find(']').unwrap() + open; + let declared: Vec = FFI_WRAPPER[open + 1..close] + .split(',') + .map(str::trim) + .filter(|s| !s.is_empty()) + .map(|s| s.parse().expect("head dim is not an integer")) + .collect(); + assert_eq!( + declared, + TURBOQUANT_CUDA_HEAD_DIMS.to_vec(), + "mistralrs-paged-attn's TURBO_HEAD_DIMS has drifted from \ + TURBOQUANT_CUDA_HEAD_DIMS" + ); + } + + /// The packed byte counts the CUDA cache layout assumes must be divisible + /// by the x=16 swizzle group, or a decode lane's run straddles two groups + /// and the vectorized load in `tq_load_klane` reads the wrong bytes. + #[test] + fn packed_k_bytes_are_swizzle_aligned() { + for dim in TURBOQUANT_CUDA_HEAD_DIMS { + let k_bytes = super::super::packed_size(dim, 4); + assert_eq!( + k_bytes % 16, + 0, + "head dim {dim}: {k_bytes} packed K bytes is not a multiple of the x=16 group" + ); + // Each of the 32 lanes owns a whole, power-of-two-sized run. + let per_lane = k_bytes / 32; + assert!( + per_lane > 0 && per_lane.is_power_of_two() && 16 % per_lane == 0, + "head dim {dim}: {per_lane} bytes per lane does not divide the 16-byte group" + ); + } + } +} diff --git a/mistralrs-quant/src/turboquant/mod.rs b/mistralrs-quant/src/turboquant/mod.rs index f28f9eca4..2037de258 100644 --- a/mistralrs-quant/src/turboquant/mod.rs +++ b/mistralrs-quant/src/turboquant/mod.rs @@ -23,11 +23,13 @@ #[allow(clippy::excessive_precision)] pub mod codebook; +pub mod cuda_tables; pub mod generate; pub mod layout; pub mod wht; use codebook::Codebook; +pub use cuda_tables::{cuda_supports_head_dim, TURBOQUANT_CUDA_HEAD_DIMS}; pub use generate::{MAX_BLOCK_DIM, MIN_BLOCK_DIM}; pub use layout::{block_plan, BlockSpec, TurboQuantLayout}; use serde::{Deserialize, Serialize};