[heydryft]

Parent system: ArcInfer / ArcSample. Four changes, ordered by the brief's priority; the instrument merges on its own merits even if the flag work is debated.

1. STEP_us instrument on the arm V4 actually runs (port, ~20 lines)

The only STEP_us TOTAL/fwd/sample/other split lived on the PagedAttention step arm (pipeline/mod.rs:1461-1482 post-change), which V4 can never reach (DeepSeekV4Loader::supports_paged_attention is false — the arm's own comment says so). The DefaultInstructions arm — the path V4 runs — had never once reported how a step divides between GPU and host. Ported verbatim: same skip-4/every-25 cadence, byte-identical log format, so tooling that parses one parses both. fwd wraps forward_inputs; sample wraps the sample-and-dispatch region; cache pre_op/post_op, the logits D2H and the split land in other — which is exactly the bucket that ranks the host candidates below.

2. ARC_SAMPLE_ON_DEVICE — device-resident sampling for B > 1 (default OFF)

Premise, verified on this checkout: at any n_seqs > 1, host_copy_batched_result (pipeline/mod.rs, the CausalGeneration => n_seqs > 1 arm) copies the ENTIRE [B, 1, vocab] logits tensor to the host — 256 × 129,280 × 2 B = 66.2 MB/step at B=256, serial, on the critical path. Sampler::sample's fast path then disqualifies every row because the logits are now on CPU (!logits.device().is_cpu(), the trivial gate), sending all 256 rows into the CPU pipeline the code itself prices at ~5.9 ms/token. The function's own doc states the trap for B=1 ("moving a B=1 batch to the host would silently disable it") and never follows it to B>1. The GPU sampler is structurally unreachable at every batch above 1. No premise in the brief was refuted; one nuance: the O(B²) split-then-copy regression is already fixed on this base — the remaining defect is the O(B) full copy itself plus the disqualification.

Change (flag OFF = bit-for-bit shipped behaviour):

• Sampler::device_sampling_terminates (sampler.rs) — the single-source mirror of sample's dispatch: trivial gate minus the device term, then greedy / no-sort / shared-memory-sort / CUDA-radix geometry. Lives next to the gates it mirrors so it cannot drift silently.
• split_batched_logits_for_sampling (pipeline/mod.rs) — eligible rows stay device views (zero bytes moved); CPU-bound rows are gathered via index_select into fresh row-sized storage and shipped in ONE batched D2H. The gather order is load-bearing: to_device/to_vec1 copy the whole backing storage of a view (CudaStorage::to_cpu_storage, candle 9586979, tensor.rs:1949/2379), so gather-then-copy is what keeps the B² term dead. All-eligible → 0 copies; none-eligible → exactly the flag-off single copy.
• Wiring in the DefaultInstructions arm only. The speculative/MTP pipelines override step entirely and never see it; the PagedAttention arm is untouched (V4 cannot reach it).
• Registered in capability_reachability.rs (Status::Live); flag read by value via env_flag_is_set, latched in a OnceLock.

Correctness bar: seeded sampling with the flag on produces IDENTICAL tokens to flag-off on CPU-checkable fixtures — asserted for every mask shape (all-device, all-host, both interleavings) against the literal flag-off path (host_copy_batched_result + index_bs): tokens [23, 42, 31, 22], pinned so both-sides drift cannot pass as parity.

Interaction with #217: is_greedy_trivial (new on this base) reads has_penalties() and therefore inherits fix 3a below — a disabled-default DRY no longer blocks the ArcGraph device-loop either. Same argument applies: apply_dry_penalty is a proven no-op at multiplier == 0.0.

3. Three measured aggravators

(a) has_penalties() counted a disabled DRY. dry_params.is_some() answered true for the disabled default (multiplier: 0.0, DrySamplingParams::default()), though apply_dry_penalty returns without touching a logit at exactly 0.0 — disqualifying the fast path and is_raw_argmax for requests that merely carried the struct. Now is_some_and(|p| p.multiplier != 0.0).

(b) Server defaults ran a 129,280-element sort nothing read. top_k = -1 (the add_request.rs default), top_p = 1.0, min_p = 0.0: sample_top_kp_min_p still ran the full-vocab sort_unstable_by via partial_sort_top_k(probs, n, true) — whose three consumers (zero-rest, top-p walk, min-p scan) are all inert under those params. Skipped straight to the multinomial; probs untouched either way, so same rng stream ⇒ same token.

(c) The global rng MutexGuard outlived the draw. Temporary lifetime extension through &mut *rng.lock() kept the guard alive across get_top_logprobs (a partial sort plus up to top_n tokenizer decodes) and tokenizer.decode — serialising every concurrently-sampling sequence behind work that never touches the rng. Narrowed to the draw.

$0 CPU probe — HOST-cost evidence ONLY, not GPU validation (D14)

Ignored test host_cost_probe::time_cpu_sample_at_v4_vocab (release, Apple Silicon host, vocab 129,280, CPU tensor, T=0.7; 3 warmup + 30 timed iters). The parity fixture's VOCAB stays 4096 — bumping it breaks that fixture's nucleus-containment premise — so the probe carries its own 129,280 constant.

params	before (mean / min)	after (mean / min)
bench: top_k=32	575 µs / 558 µs	664 µs / 578 µs (unchanged, within noise)
server: top_k=-1, top_p=1.0	10,338 µs / 1,971 µs	403 µs / 367 µs

The server-default row is fix 3b: the full-vocab sort (plus its 1 MB pair-Vec) was ~80-96% of per-row CPU sampling cost at V4's vocabulary. This prices the host pipeline only; it says nothing about GPU behaviour.

Tests + mutation checks (assertion values both ways)

• 3a disabled_dry_is_not_a_penalty_and_enabled_dry_is: disabled DRY ⇒ is_raw_argmax() == true and sampled token equals no-DRY token; enabled (multiplier 0.8) ⇒ false. Mutated (revert to dry_params.is_some()): fails at "multiplier = 0.0 … must not disqualify" (is_raw_argmax = false).
• 3b server_default_top_k_skips_the_unconsumed_sort: skip path sorts 0 times (thread-local PARTIAL_SORT_CALLS delta), pre-fix path spelled out literally sorts 1, both draw token 506 on seed 42 (pinned). Mutated (skip removed): fails left: 1, right: 0 ("server defaults must skip the sort entirely"). each_active_filter_keeps_the_sort_alive: top_p=0.9 ⇒ 1 sort, min_p=0.1 ⇒ 1, top_k=32 ⇒ 1.
• 3c no direct lock-window assertion (needs loom); evidence: the full seeded sampler suite (34 tests) passes unchanged and both pinned draws (506, [23,42,31,22]) are unmoved, i.e. the draw's position in the rng stream did not shift.
• Partition device_sampling_split_tests: mixed batch = 1 host copy and the gathered storage holds exactly n_cpu × vocab elements (21 = 3×7) while eligible rows stay views of the full batch (42); all-eligible = 0 copies; none-eligible = 1 copy (flag-off parity). Ramp fixture makes any host-position mapping error a visibly wrong slice. Counter-delta tests serialise on a test-only lock (the counter is process-global and was racy under the parallel test runner).
• capability_reachability (12) and doc_citations (11) suites green on the rebased head.

Known adjacent facts (report only — outside my file fence)

• The per-sequence prologue seq.get_toks().to_vec() runs twice on the engine thread (sampling.rs, sample_sequence + its llguidance re-sample arm), and the response phase is a serial for-await. Owned by the device-loop/completion lane; not touched here.
• Residual bounded risk with the flag ON: a device-kept row whose CUDA radix sample declines at runtime (Ok(None): pure-top-p nucleus not contained in the candidates) falls back to the CPU path and pays a per-row D2H — and if the batch logits are already F32, to_vec1 on the shared-storage view prices that at a full-batch copy for that row. Statically predictable declines are already excluded by device_sampling_terminates; the runtime case is counted by gpu_sampling_health::DECLINED. If the box A/B shows DECLINED > 0 under real traffic, the one-line hardening is a row-narrowing copy in sampling.rs's prologue — which is outside this PR's fence, so it is named here instead of shipped.

What this PR does NOT do

• No GPU numbers are claimed anywhere: the flag ships OFF for the box A/B to flip; the probe is labelled host-cost-only (D14).
• No changes to pipeline/sampling.rs, engine/mod.rs, or the PagedAttention arm.
• The STEP_us statics are per-arm (the two arms never run in one process); the paged arm's instrument is untouched.

🤖 Generated with Claude Code

https://claude.ai/code/session_01WK8VgocBUrr5djjE1ZfCNf

[github-actions]

Code Metrics Report

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                  5          305          210           52           43
 CSS                       2         1181         1036           34          111
 CUDA                     81        29475        20161         6275         3039
 Dockerfile                1           39           22            8            9
 JavaScript               16         3546         2676          482          388
 Jinja2                    7          694          656            5           33
 JSON                     75         4896         4893            0            3
 Makefile                  1            6            5            0            1
 Metal Shading Lan|       33        12224         9431         1142         1651
 PowerShell                1          300          227           30           43
 Python                  147        15371        12677          824         1870
 Shell                    42        10304         6858         2769          677
 Plain Text                4         3801            0         2479         1322
 TOML                     33         1498         1294           54          150
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 HTML                      4         2687         2604           43           40
 |- CSS                    2          543          479           37           27
 |- JavaScript             1         1233         1215           12            6
 (Total)                             4463         4298           92           73
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         4          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                213        48932            0        38081        10851
 |- BASH                  72         1655         1203          331          121
 |- C                      4           19           19            0            0
 |- CUDA                   2           84           56           16           12
 |- JSON                  19          779          779            0            0
 |- PowerShell             1            1            1            0            0
 |- Python                23         1008          787          113          108
 |- Rust                  68         2063         1727           78          258
 |- TOML                   6          207          164            0           43
 |- YAML                   5           41           36            5            0
 (Total)                            54789         4772        38624        11393
─────────────────────────────────────────────────────────────────────────────────
 Rust                    681       340988       293252        18096        29640
 |- Markdown             504        32562          471        28063         4028
 (Total)                           373550       293723        46159        33668
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1353       516771       363188        99077        54506
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━