perf(qwen35): greedy decode allocates nothing and downloads 4 bytes per token (Refs #4215)

crates/aprender-gpu/src/driver/memory/buffer.rs

@@ -90,6 +90,31 @@ thread_local! {
     /// Device bytes this thread has allocated through `GpuBuffer` and not yet
     /// freed.
     static OUTSTANDING_DEVICE_BYTES: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
+    /// Device allocations this thread has made through `GpuBuffer`, ever (#4215).
+    static DEVICE_ALLOCS: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
+    /// Bytes this thread has copied device→host through `GpuBuffer`, ever (#4215).
+    static DEVICE_TO_HOST_BYTES: std::cell::Cell<u64> = const { std::cell::Cell::new(0) };
+}
+
+/// Device allocations the current thread has made through `GpuBuffer` since it
+/// started — a monotonic count, so a caller measures a region by differencing
+/// two readings (#4215: a steady-state decode token must allocate nothing).
+#[must_use]
+pub fn device_allocs_total() -> u64 {
+    DEVICE_ALLOCS.with(std::cell::Cell::get)
+}
+
+/// Bytes the current thread has copied device→host through `GpuBuffer` since
+/// it started, monotonic like [`device_allocs_total`] (#4215: a greedy decode
+/// token downloads the 4-byte token id, not the logits).
+#[must_use]
+pub fn device_to_host_bytes_total() -> u64 {
+    DEVICE_TO_HOST_BYTES.with(std::cell::Cell::get)
+}
+
+/// Record a device→host copy against the current thread.
+pub(super) fn record_device_to_host(bytes: usize) {
+    DEVICE_TO_HOST_BYTES.with(|c| c.set(c.get().saturating_add(bytes as u64)));
 }
 
 /// Device bytes allocated by the current thread through `GpuBuffer` and not
@@ -124,6 +149,7 @@ pub fn device_bytes_outstanding() -> u64 {
 /// Record a successful device allocation against the current thread.
 fn record_device_alloc(bytes: usize) {
     OUTSTANDING_DEVICE_BYTES.with(|c| c.set(c.get().saturating_add(bytes as u64)));
+    DEVICE_ALLOCS.with(|c| c.set(c.get().saturating_add(1)));
 }
 
 /// Record a device free against the current thread.

crates/aprender-gpu/src/driver/memory/mod.rs

@@ -18,7 +18,7 @@ mod transfer;
 
 #[allow(unused_imports, unreachable_pub)] // Re-exported by driver/mod.rs
 pub use buffer::{
-    classify_device_memory, device_bytes_outstanding, device_memory_exclusive, DeviceMemoryClass,
+    classify_device_memory, device_allocs_total, device_bytes_outstanding, device_to_host_bytes_total, device_memory_exclusive, DeviceMemoryClass,
     DeviceMemoryExclusive, GpuBuffer, GpuBufferView,
 };
 

crates/aprender-gpu/src/driver/memory/transfer.rs

@@ -15,7 +15,7 @@ use crate::GpuError;
 // zeros when the CUDA context is not current on the calling thread
 // (see paiml/trueno#232).
 
-use super::buffer::GpuBuffer;
+use super::buffer::{record_device_to_host, GpuBuffer};
 
 // ============================================================================
 // Blocking host-to-device primitive
@@ -203,6 +203,7 @@ impl<T: Copy> GpuBuffer<T> {
         // SAFETY: data is valid for size bytes, ptr is valid device pointer
         let result =
             unsafe { (driver.cuMemcpyDtoH)(data.as_mut_ptr() as *mut c_void, self.ptr, size) };
+        record_device_to_host(size);
         CudaDriver::check(result).map_err(|e| GpuError::Transfer(e.to_string()))
     }
 
@@ -294,6 +295,7 @@ impl<T: Copy> GpuBuffer<T> {
                 stream.raw(),
             )
         };
+        record_device_to_host(size);
         CudaDriver::check(result).map_err(|e| GpuError::Transfer(e.to_string()))
     }
 
@@ -382,6 +384,7 @@ impl<T: Copy> GpuBuffer<T> {
         // SAFETY: bounds checked above, data and ptr are valid
         let result =
             unsafe { (driver.cuMemcpyDtoH)(data.as_mut_ptr() as *mut c_void, src_ptr, size) };
+        record_device_to_host(size);
         CudaDriver::check(result).map_err(|e| GpuError::Transfer(e.to_string()))
     }
 

crates/aprender-gpu/src/driver/mod.rs

@@ -128,7 +128,8 @@ pub use cublaslt::CublasLtHandle;
 pub use graph::{CaptureMode, CudaGraph, CudaGraphExec};
 #[cfg(feature = "cuda")]
 pub use memory::{
-    device_bytes_outstanding, device_memory_exclusive, DeviceMemoryExclusive, GpuBuffer,
+    device_allocs_total, device_bytes_outstanding, device_memory_exclusive,
+    device_to_host_bytes_total, DeviceMemoryExclusive, GpuBuffer,
 };
 pub use module::CudaModule;
 #[cfg(feature = "cuda")]

crates/aprender-gpu/src/kernels/argmax.rs

@@ -19,8 +19,29 @@
 //! - Expected speedup: ~1.2x overall (from ~163 tok/s to ~200 tok/s)
 
 use super::Kernel;
-use crate::ptx::builder::{PtxArithmetic, PtxComparison, PtxControl};
-use crate::ptx::{PtxKernel, PtxType};
+use crate::ptx::builder::{KernelBuilder, PtxArithmetic, PtxComparison, PtxControl};
+use crate::ptx::{PtxKernel, PtxType, VirtualReg};
+
+/// Does the candidate at `tid + stride` replace this thread's current best?
+///
+/// Larger value wins; on an EXACT tie the lower index wins, so the reduction
+/// returns the FIRST maximum — the contract of the CPU `ops::argmax` this
+/// kernel replaces (#4215). A strict `>` alone keeps whichever side the tree
+/// happened to put in `tid`, which after the strided load can be the later
+/// index. Both comparisons are ordered, so a NaN never wins.
+fn other_wins(
+    ctx: &mut KernelBuilder<'_>,
+    other_val: VirtualReg,
+    other_idx: VirtualReg,
+    my_val: VirtualReg,
+    my_idx: VirtualReg,
+) -> VirtualReg {
+    let greater = ctx.setp_gt_f32(other_val, my_val);
+    let equal = ctx.setp_eq_f32(other_val, my_val);
+    let earlier = ctx.setp_lt_u32(other_idx, my_idx);
+    let tie_to_earlier = ctx.and_pred(equal, earlier);
+    ctx.or_pred(greater, tie_to_earlier)
+}
 
 /// ArgMax kernel configuration
 ///
@@ -175,7 +196,7 @@ impl Kernel for ArgMaxKernel {
                     let my_val = ctx.ld_generic_f32(sh_val_addr);
                     let my_idx = ctx.ld_generic_u32(sh_idx_addr);
 
-                    let is_greater = ctx.setp_gt_f32(other_val, my_val);
+                    let is_greater = other_wins(ctx, other_val, other_idx, my_val, my_idx);
                     let new_val = ctx.selp_f32(is_greater, other_val, my_val);
                     let new_idx = ctx.selp_u32(is_greater, other_idx, my_idx);
 
@@ -201,7 +222,7 @@ impl Kernel for ArgMaxKernel {
                         let my_val = ctx.ld_generic_f32(sh_val_addr);
                         let my_idx = ctx.ld_generic_u32(sh_idx_addr);
 
-                        let is_greater = ctx.setp_gt_f32(other_val, my_val);
+                        let is_greater = other_wins(ctx, other_val, other_idx, my_val, my_idx);
                         let new_val = ctx.selp_f32(is_greater, other_val, my_val);
                         let new_idx = ctx.selp_u32(is_greater, other_idx, my_idx);
 
@@ -328,7 +349,7 @@ impl Kernel for ArgMaxFinalKernel {
                         let my_val = ctx.ld_generic_f32(sh_val_addr);
                         let my_idx = ctx.ld_generic_u32(sh_idx_addr);
 
-                        let is_greater = ctx.setp_gt_f32(other_val, my_val);
+                        let is_greater = other_wins(ctx, other_val, other_idx, my_val, my_idx);
                         let new_val = ctx.selp_f32(is_greater, other_val, my_val);
                         let new_idx = ctx.selp_u32(is_greater, other_idx, my_idx);
 

crates/aprender-gpu/src/ptx/builder/bitwise_ops.rs

@@ -184,6 +184,18 @@ impl<'a> KernelBuilder<'a> {
         dst
     }
 
+    /// OR two predicates: dst = a OR b
+    pub fn or_pred(&mut self, a: VirtualReg, b: VirtualReg) -> VirtualReg {
+        let dst = self.registers.allocate_virtual(PtxType::Pred);
+        self.instructions.push(
+            PtxInstruction::new(PtxOp::Or, PtxType::Pred)
+                .dst(Operand::Reg(dst))
+                .src(Operand::Reg(a))
+                .src(Operand::Reg(b)),
+        );
+        dst
+    }
+
     /// Get shared memory base pointer
     ///
     /// PAR-062: Returns base address of shared memory for this block

crates/aprender-gpu/src/ptx/builder/comparison.rs

@@ -118,6 +118,18 @@ pub trait PtxComparison: KernelBuilderCore {
         pred
     }
 
+    /// Set predicate if a == b (f32). Ordered: false when either is NaN.
+    fn setp_eq_f32(&mut self, a: VirtualReg, b: VirtualReg) -> VirtualReg {
+        let pred = self.registers_mut().allocate_virtual(PtxType::Pred);
+        let mut instr = PtxInstruction::new(PtxOp::Setp, PtxType::F32)
+            .dst(Operand::Reg(pred))
+            .src(Operand::Reg(a))
+            .src(Operand::Reg(b));
+        instr.label = Some(CmpOp::Eq.to_ptx_string().to_string());
+        self.instructions_mut().push(instr);
+        pred
+    }
+
     /// Set predicate comparing u32 with immediate
     fn setp_ge_u32_imm(&mut self, a: VirtualReg, b: u32) -> VirtualReg {
         let pred = self.registers_mut().allocate_virtual(PtxType::Pred);
@@ -226,9 +238,13 @@ mod tests {
         let _lt = builder.setp_lt_f32(a, b);
         let _gt = builder.setp_gt_f32(a, b);
 
-        assert_eq!(builder.instructions.len(), 2);
+        let _eq = builder.setp_eq_f32(a, b);
+
+        assert_eq!(builder.instructions.len(), 3);
         assert_eq!(builder.instructions[0].ty, PtxType::F32);
         assert_eq!(builder.instructions[1].ty, PtxType::F32);
+        assert_eq!(builder.instructions[2].ty, PtxType::F32);
+        assert_eq!(builder.instructions[2].label.as_deref(), Some("eq"));
     }
 
     #[test]

crates/aprender-gpu/tests/argmax_falsification_f114/main.rs

@@ -16,3 +16,5 @@ mod barrier_bounds;
 mod gpu_correctness;
 #[cfg(feature = "cuda")]
 mod ptx_analysis;
+#[cfg(feature = "cuda")]
+mod tie_break;

crates/aprender-gpu/tests/argmax_falsification_f114/tie_break.rs

@@ -0,0 +1,149 @@
+//! F114 tie-break (#4215): the GPU argmax must return the FIRST maximum, the
+//! contract of the CPU `ops::argmax` it replaces on the greedy decode path.
+//!
+//! The tree reduction compares thread `tid` against `tid + stride`, but after a
+//! strided load `tid` can already hold the LARGER index (thread 0 owns
+//! 0/256/512/768, thread 1 owns 1/257/…), so a strict `>` keeps index 768 over
+//! an exactly equal value at index 1. The same holds across blocks in the final
+//! pass once a subtree has absorbed a later block. On a temp-0 decode an exact
+//! tie then emits a different token than the CPU path — a silent divergence.
+//!
+//! Each case below states the index the CPU contract picks; the pre-fix kernel
+//! returns the other one.
+
+use std::ffi::c_void;
+use trueno_gpu::driver::{CudaContext, CudaModule, CudaStream, GpuBuffer, LaunchConfig};
+use trueno_gpu::kernels::{ArgMaxFinalKernel, ArgMaxKernel, Kernel};
+
+/// The CPU contract (`aprender-serve` `gguf::ops::argmax`): strict `>` in index
+/// order, so the first maximum wins and NaN never does.
+fn cpu_first_argmax(v: &[f32]) -> u32 {
+    let (mut best, mut best_val) = (0u32, f32::NEG_INFINITY);
+    for (i, &x) in v.iter().enumerate() {
+        if x > best_val {
+            best_val = x;
+            best = i as u32;
+        }
+    }
+    best
+}
+
+/// Both passes, exactly as `CudaExecutor::gpu_argmax` launches them. `None`
+/// when this host has no CUDA device, so the caller can say it did not measure.
+fn gpu_argmax(input: &[f32]) -> Option<u32> {
+    let ctx = CudaContext::new(0).ok()?;
+    let stream = CudaStream::new(&ctx).expect("stream");
+    let len = input.len() as u32;
+
+    let block = ArgMaxKernel::new(len);
+    let num_blocks = block.num_blocks();
+    let fin = ArgMaxFinalKernel::new(num_blocks);
+    let mut block_mod = CudaModule::from_ptx(&ctx, &block.emit_ptx()).expect("block ptx");
+    let mut final_mod = CudaModule::from_ptx(&ctx, &fin.emit_ptx()).expect("final ptx");
+
+    let mut in_buf: GpuBuffer<f32> = GpuBuffer::new(&ctx, input.len()).expect("alloc");
+    let mut vals: GpuBuffer<f32> = GpuBuffer::new(&ctx, num_blocks as usize).expect("alloc");
+    let mut idxs: GpuBuffer<u32> = GpuBuffer::new(&ctx, num_blocks as usize).expect("alloc");
+    let mut out: GpuBuffer<u32> = GpuBuffer::new(&ctx, 1).expect("alloc");
+    in_buf.copy_from_host(input).expect("upload");
+
+    let mut args: [*mut c_void; 4] = [
+        in_buf.as_kernel_arg(),
+        vals.as_kernel_arg(),
+        idxs.as_kernel_arg(),
+        &len as *const u32 as *mut c_void,
+    ];
+    let cfg = LaunchConfig {
+        grid: (num_blocks, 1, 1),
+        block: (256, 1, 1),
+        shared_mem: 0,
+    };
+    // SAFETY: buffers sized to the kernel's immediates; args match `.param` order.
+    unsafe { stream.launch_kernel(&mut block_mod, block.name(), &cfg, &mut args) }
+        .expect("block launch");
+
+    let mut fargs: [*mut c_void; 4] = [
+        vals.as_kernel_arg(),
+        idxs.as_kernel_arg(),
+        out.as_kernel_arg(),
+        &num_blocks as *const u32 as *mut c_void,
+    ];
+    let fcfg = LaunchConfig {
+        grid: (1, 1, 1),
+        block: (256, 1, 1),
+        shared_mem: 0,
+    };
+    // SAFETY: as above.
+    unsafe { stream.launch_kernel(&mut final_mod, fin.name(), &fcfg, &mut fargs) }
+        .expect("final launch");
+    stream.synchronize().expect("sync");
+
+    let mut host = [0u32; 1];
+    out.copy_to_host(&mut host).expect("download");
+    Some(host[0])
+}
+
+/// Qwen3.5's vocabulary: 243 blocks, all of the final pass's tree in play.
+const VOCAB: usize = 248_320;
+
+fn case(name: &str, input: &[f32]) {
+    let want = cpu_first_argmax(input);
+    let Some(got) = gpu_argmax(input) else {
+        // Not a pass: nothing was measured. APR_REQUIRE_CUDA=1 (the GPU host
+        // that runs this falsifier) turns the absence into a failure.
+        assert!(
+            std::env::var("APR_REQUIRE_CUDA").as_deref() != Ok("1"),
+            "{name}: APR_REQUIRE_CUDA=1 but no CUDA device"
+        );
+        eprintln!("{name}: NOT MEASURED (no CUDA device)");
+        return;
+    };
+    assert_eq!(got, want, "{name}: GPU argmax {got}, CPU first-max {want}");
+}
+
+#[test]
+fn f114_tie_within_block_picks_lowest_index() {
+    // Thread 0 reaches 768 (its 4th strided load); thread 1 holds 1. Final
+    // intra-block step compares tid 0 (idx 768) with tid 1 (idx 1).
+    let mut v = vec![-1.0f32; VOCAB];
+    v[768] = 5.0;
+    v[1] = 5.0;
+    case("within-block", &v);
+}
+
+#[test]
+fn f114_tie_across_blocks_picks_lowest_index() {
+    // After the final pass's stride-128 step, tid 0 holds block 128's winner;
+    // the last step compares it with tid 1 (block 1). Block 1's element is
+    // the lower index.
+    let mut v = vec![-1.0f32; VOCAB];
+    v[128 * 1024 + 7] = 5.0;
+    v[1024 + 3] = 5.0;
+    case("across-blocks", &v);
+}
+
+#[test]
+fn f114_tie_many_equal_maxima_picks_first() {
+    // Every element equal: the first index is the only correct answer.
+    case("all-equal", &vec![0.25f32; VOCAB]);
+}
+
+#[test]
+fn f114_nan_never_wins() {
+    // A NaN at a lower index than the real maximum must not be chosen, with or
+    // without the tie-break (setp.eq and setp.gt are ordered comparisons).
+    let mut v = vec![-1.0f32; VOCAB];
+    v[3] = f32::NAN;
+    v[1] = 2.0;
+    v[900] = 2.0;
+    case("nan", &v);
+}
+
+#[test]
+fn f114_no_tie_is_unchanged() {
+    // Regression guard: distinct values, maximum late in the vocabulary.
+    let v: Vec<f32> = (0..VOCAB)
+        .map(|i| ((i as u64 * 2_654_435_761) % 1_000_003) as f32 * 1e-6)
+        .collect();
+    case("distinct", &v);
+}

crates/aprender-serve/src/cuda/executor/core.rs

@@ -49,6 +49,7 @@ impl CudaExecutor {
             memory_pool: GpuMemoryPool::new(),
             staging_pool: StagingBufferPool::new(), // PARITY-042: pinned memory pool
             modules: std::mem::ManuallyDrop::new(HashMap::new()),
+            module_keys: HashMap::new(),
             weight_cache: HashMap::new(),
             named_fp16_weight_cache: HashMap::new(), // GH-174: SafeTensors F16
             quantized_weight_cache: HashMap::new(),  // PAR-005: quantized weight cache