From e4df9c25af08d5846831b583e32eca343a4be3bf Mon Sep 17 00:00:00 2001 From: Noah Gift Date: Thu, 24 Sep 2026 14:14:39 +0200 Subject: [PATCH 1/4] GPU argmax returned the LAST of two exactly-equal maxima; the CPU argmax it replaces returns the FIRST (Closes #4259) The block and final tree reductions kept `other` only on strict `>`. After the strided load, thread 0 holds index 768 and thread 1 holds index 1, so on an exact tie the reduction keeps 768; in the final pass a subtree that has absorbed block 128 beats block 1 the same way. `gguf::ops::argmax` (strict `>` in index order) returns the first maximum, so a temp-0 decode moved onto the GPU argmax (#4215 (b)) would emit a different token on every exact tie. `other_wins` = greater OR (equal AND lower index), applied at all three reduction sites. Both comparisons are ordered, so NaN still never wins. New builder ops: `setp_eq_f32`, `or_pred`. Top-k is untouched. tests/argmax_falsification_f114/tie_break.rs: within-block tie, cross-block tie, all-equal, NaN-at-low-index, and a no-tie regression at Qwen3.5's 248320 vocab, each asserted against the CPU first-max contract. APR_REQUIRE_CUDA=1 turns a missing device into a failure instead of a skip. Closes #4259. Cherry-picked from perf/4215 (c8ce988c3). Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/aprender-gpu/src/kernels/argmax.rs | 31 +++- .../src/ptx/builder/bitwise_ops.rs | 12 ++ .../src/ptx/builder/comparison.rs | 18 ++- .../tests/argmax_falsification_f114/main.rs | 2 + .../argmax_falsification_f114/tie_break.rs | 149 ++++++++++++++++++ docs/roadmaps/entries/PMAT-4259.yaml | 18 +++ docs/roadmaps/roadmap.yaml | 18 +++ 7 files changed, 242 insertions(+), 6 deletions(-) create mode 100644 crates/aprender-gpu/tests/argmax_falsification_f114/tie_break.rs create mode 100644 docs/roadmaps/entries/PMAT-4259.yaml diff --git a/crates/aprender-gpu/src/kernels/argmax.rs b/crates/aprender-gpu/src/kernels/argmax.rs index 57d9bd8964..29fb6a1a6d 100644 --- a/crates/aprender-gpu/src/kernels/argmax.rs +++ b/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); diff --git a/crates/aprender-gpu/src/ptx/builder/bitwise_ops.rs b/crates/aprender-gpu/src/ptx/builder/bitwise_ops.rs index dc2bc1bf5f..eb6a8cd9d4 100644 --- a/crates/aprender-gpu/src/ptx/builder/bitwise_ops.rs +++ b/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 diff --git a/crates/aprender-gpu/src/ptx/builder/comparison.rs b/crates/aprender-gpu/src/ptx/builder/comparison.rs index a55b861fa5..a63d46725d 100644 --- a/crates/aprender-gpu/src/ptx/builder/comparison.rs +++ b/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] diff --git a/crates/aprender-gpu/tests/argmax_falsification_f114/main.rs b/crates/aprender-gpu/tests/argmax_falsification_f114/main.rs index d9a6035208..f08b52b2e1 100644 --- a/crates/aprender-gpu/tests/argmax_falsification_f114/main.rs +++ b/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; diff --git a/crates/aprender-gpu/tests/argmax_falsification_f114/tie_break.rs b/crates/aprender-gpu/tests/argmax_falsification_f114/tie_break.rs new file mode 100644 index 0000000000..fb7aa7d7b4 --- /dev/null +++ b/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 { + 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 = GpuBuffer::new(&ctx, input.len()).expect("alloc"); + let mut vals: GpuBuffer = GpuBuffer::new(&ctx, num_blocks as usize).expect("alloc"); + let mut idxs: GpuBuffer = GpuBuffer::new(&ctx, num_blocks as usize).expect("alloc"); + let mut out: GpuBuffer = 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 = (0..VOCAB) + .map(|i| ((i as u64 * 2_654_435_761) % 1_000_003) as f32 * 1e-6) + .collect(); + case("distinct", &v); +} diff --git a/docs/roadmaps/entries/PMAT-4259.yaml b/docs/roadmaps/entries/PMAT-4259.yaml new file mode 100644 index 0000000000..33f18399bc --- /dev/null +++ b/docs/roadmaps/entries/PMAT-4259.yaml @@ -0,0 +1,18 @@ +- id: PMAT-4259 + github_issue: 4259 + item_type: task + title: 'GPU argmax kernel keeps the later index on an exact tie; CPU argmax keeps the first — breaks greedy token identity once #4215(b) moves argmax to GPU' + status: inprogress + priority: medium + assigned_to: null + created: 2026-09-24T12:01:41Z + updated: 2026-09-24T12:01:41Z + spec: null + acceptance_criteria: + - 'The GPU argmax returns the FIRST index of an exact maximum, as the CPU argmax does; a red-first tie-break falsifier covers ties within a thread, across threads, and across blocks.' + phases: [] + subtasks: [] + estimated_effort: null + labels: + - kind:code + notes: 'Ported from c8ce988c3 (perf/4215, based on 0.69.3) onto main; prerequisite for #4215(b), greedy GPU argmax.' diff --git a/docs/roadmaps/roadmap.yaml b/docs/roadmaps/roadmap.yaml index 90671942ca..b187fff708 100644 --- a/docs/roadmaps/roadmap.yaml +++ b/docs/roadmaps/roadmap.yaml @@ -20610,3 +20610,21 @@ roadmap: labels: - kind:code notes: null +- id: PMAT-4259 + github_issue: 4259 + item_type: task + title: 'GPU argmax kernel keeps the later index on an exact tie; CPU argmax keeps the first — breaks greedy token identity once #4215(b) moves argmax to GPU' + status: inprogress + priority: medium + assigned_to: null + created: 2026-09-24T12:01:41Z + updated: 2026-09-24T12:01:41Z + spec: null + acceptance_criteria: + - 'The GPU argmax returns the FIRST index of an exact maximum, as the CPU argmax does; a red-first tie-break falsifier covers ties within a thread, across threads, and across blocks.' + phases: [] + subtasks: [] + estimated_effort: null + labels: + - kind:code + notes: 'Ported from c8ce988c3 (perf/4215, based on 0.69.3) onto main; prerequisite for #4215(b), greedy GPU argmax.' From 7e08584ca664658b6f485acb67064f333b23af2f Mon Sep 17 00:00:00 2001 From: Noah Gift Date: Thu, 24 Sep 2026 14:14:45 +0200 Subject: [PATCH 2/4] [WIP] perf(qwen35): allocation-free greedy decode token (Refs #4215) Co-Authored-By: Claude Opus 5.5 (1M context) --- .../aprender-gpu/src/driver/memory/buffer.rs | 26 +++ crates/aprender-gpu/src/driver/memory/mod.rs | 2 +- .../src/driver/memory/transfer.rs | 5 +- crates/aprender-gpu/src/driver/mod.rs | 3 +- .../aprender-serve/src/cuda/executor/core.rs | 1 + .../src/cuda/executor/gdn_ops.rs | 120 ++++++++--- .../src/cuda/executor/layer_norm_gpu.rs | 24 +-- .../src/cuda/executor/layers/par-062.rs | 42 ++-- .../aprender-serve/src/cuda/executor/mod.rs | 53 +++++ .../src/cuda/executor/q4k_mwv_gemv.rs | 102 +++++----- .../src/cuda/executor/q6k_gemv_indexed.rs | 36 ++-- .../src/cuda/executor/residual.rs | 10 +- .../aprender-serve/src/cuda/executor/silu.rs | 10 +- .../src/cuda/executor/weight.rs | 96 ++++----- .../src/gguf/cuda/forward_qwen35_cuda.rs | 105 ++++++++-- .../forward_qwen35_decode_overhead_tests.rs | 189 ++++++++++++++++++ .../gguf/inference/forward/forward_qwen35.rs | 51 ++++- 17 files changed, 656 insertions(+), 219 deletions(-) create mode 100644 crates/aprender-serve/src/gguf/cuda/forward_qwen35_decode_overhead_tests.rs diff --git a/crates/aprender-gpu/src/driver/memory/buffer.rs b/crates/aprender-gpu/src/driver/memory/buffer.rs index 2590cfbf61..01b67e648a 100644 --- a/crates/aprender-gpu/src/driver/memory/buffer.rs +++ b/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 = const { std::cell::Cell::new(0) }; + /// Device allocations this thread has made through `GpuBuffer`, ever (#4215). + static DEVICE_ALLOCS: std::cell::Cell = 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 = 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. diff --git a/crates/aprender-gpu/src/driver/memory/mod.rs b/crates/aprender-gpu/src/driver/memory/mod.rs index 8aae9bbbeb..14fd3b7e2b 100644 --- a/crates/aprender-gpu/src/driver/memory/mod.rs +++ b/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, }; diff --git a/crates/aprender-gpu/src/driver/memory/transfer.rs b/crates/aprender-gpu/src/driver/memory/transfer.rs index 111d6cbc9c..7cd6a9f18e 100644 --- a/crates/aprender-gpu/src/driver/memory/transfer.rs +++ b/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 GpuBuffer { // 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 GpuBuffer { stream.raw(), ) }; + record_device_to_host(size); CudaDriver::check(result).map_err(|e| GpuError::Transfer(e.to_string())) } @@ -382,6 +384,7 @@ impl GpuBuffer { // 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())) } diff --git a/crates/aprender-gpu/src/driver/mod.rs b/crates/aprender-gpu/src/driver/mod.rs index 8df3944a87..e87433fd72 100644 --- a/crates/aprender-gpu/src/driver/mod.rs +++ b/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")] diff --git a/crates/aprender-serve/src/cuda/executor/core.rs b/crates/aprender-serve/src/cuda/executor/core.rs index 58aff3d9bd..becef41b1b 100644 --- a/crates/aprender-serve/src/cuda/executor/core.rs +++ b/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 diff --git a/crates/aprender-serve/src/cuda/executor/gdn_ops.rs b/crates/aprender-serve/src/cuda/executor/gdn_ops.rs index 5446ac9d6b..63840905f0 100644 --- a/crates/aprender-serve/src/cuda/executor/gdn_ops.rs +++ b/crates/aprender-serve/src/cuda/executor/gdn_ops.rs @@ -16,6 +16,27 @@ use super::*; +/// The most argument slots any Gated `DeltaNet` kernel takes. +const GDN_MAX_ARGS: usize = 16; + +/// Pack `ptrs` then `scalars` into `args`; returns the slot count. +fn gdn_pack_args( + kernel_name: &str, + args: &mut [u64; GDN_MAX_ARGS], + ptrs: &[u64], + scalars: &[u64], +) -> Result { + let n = ptrs.len() + scalars.len(); + if n > GDN_MAX_ARGS { + return Err(GpuError::InvalidParameter(format!( + "{kernel_name}: {n} arguments exceed {GDN_MAX_ARGS}" + ))); + } + args[..ptrs.len()].copy_from_slice(ptrs); + args[ptrs.len()..n].copy_from_slice(scalars); + Ok(n) +} + impl CudaExecutor { /// Compile (once) and fetch the module for a Gated `DeltaNet` kernel. /// @@ -49,13 +70,16 @@ impl CudaExecutor { ptrs: &[u64], ) -> Result<(), GpuError> { for (i, &p) in ptrs.iter().enumerate() { - validate_device_ptr(p, &format!("{kernel_name} arg {i}"))?; + validate_kernel_arg(p, kernel_name, i)?; + } + // #4215: argument slots live on the stack; the `Vec` is built only when + // a graph is being recorded. + let mut args = [0u64; GDN_MAX_ARGS]; + let n = gdn_pack_args(kernel_name, &mut args, ptrs, &[])?; + let mut raw = [std::ptr::null_mut::(); GDN_MAX_ARGS]; + for (r, a) in raw.iter_mut().zip(args.iter_mut()) { + *r = std::ptr::from_mut(a).cast::(); } - let mut args: Vec = ptrs.to_vec(); - let mut raw: Vec<*mut std::ffi::c_void> = args - .iter_mut() - .map(|a| std::ptr::from_mut(a).cast::()) - .collect(); let module = self .modules @@ -66,7 +90,7 @@ impl CudaExecutor { // argument order matches the kernel's `.param` declarations. unsafe { self.stream - .launch_kernel(module, kernel_name, &config, &mut raw)?; + .launch_kernel(module, kernel_name, &config, &mut raw[..n])?; } // trueno#243 / #3413: the manual decode graph is rebuilt ONLY from @@ -78,7 +102,7 @@ impl CudaExecutor { self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), config, - arg_data: args, + arg_data: args[..n].to_vec(), }); } Ok(()) @@ -102,14 +126,16 @@ impl CudaExecutor { scalars: &[u64], ) -> Result<(), GpuError> { for (i, &p) in ptrs.iter().enumerate() { - validate_device_ptr(p, &format!("{kernel_name} arg {i}"))?; + validate_kernel_arg(p, kernel_name, i)?; + } + // #4215: argument slots live on the stack; the `Vec` is built only when + // a graph is being recorded. + let mut args = [0u64; GDN_MAX_ARGS]; + let n = gdn_pack_args(kernel_name, &mut args, ptrs, scalars)?; + let mut raw = [std::ptr::null_mut::(); GDN_MAX_ARGS]; + for (r, a) in raw.iter_mut().zip(args.iter_mut()) { + *r = std::ptr::from_mut(a).cast::(); } - let mut args: Vec = ptrs.to_vec(); - args.extend_from_slice(scalars); - let mut raw: Vec<*mut std::ffi::c_void> = args - .iter_mut() - .map(|a| std::ptr::from_mut(a).cast::()) - .collect(); let module = self .modules @@ -121,7 +147,7 @@ impl CudaExecutor { // kernel's `.param` declarations. unsafe { self.stream - .launch_kernel(module, kernel_name, &config, &mut raw)?; + .launch_kernel(module, kernel_name, &config, &mut raw[..n])?; } // trueno#243 / #3413: the manual decode graph is rebuilt ONLY from @@ -133,7 +159,7 @@ impl CudaExecutor { self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), config, - arg_data: args, + arg_data: args[..n].to_vec(), }); } Ok(()) @@ -161,7 +187,7 @@ impl CudaExecutor { num_heads, head_dim, }; - let cache_key = format!("gdn_split_interleaved_{num_heads}_{head_dim}"); + let cache_key = module_key!(self, "gdn_split_interleaved_{}_{}", num_heads, head_dim); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); @@ -201,7 +227,13 @@ impl CudaExecutor { head_dim, n_rot, }; - let cache_key = format!("gdn_partial_neox_rope_{num_heads}_{head_dim}_{n_rot}"); + let cache_key = module_key!( + self, + "gdn_partial_neox_rope_{}_{}_{}", + num_heads, + head_dim, + n_rot + ); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); @@ -246,7 +278,13 @@ impl CudaExecutor { num_kv_heads, head_dim, }; - let cache_key = format!("gdn_decode_attention_{num_heads}_{num_kv_heads}_{head_dim}"); + let cache_key = module_key!( + self, + "gdn_decode_attention_{}_{}_{}", + num_heads, + num_kv_heads, + head_dim + ); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); @@ -288,7 +326,7 @@ impl CudaExecutor { channels, kernel_size, }; - let cache_key = format!("gdn_causal_conv1d_silu_{channels}_{kernel_size}"); + let cache_key = module_key!(self, "gdn_causal_conv1d_silu_{}_{}", channels, kernel_size); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); @@ -328,7 +366,18 @@ impl CudaExecutor { epsilon, }; // epsilon is an immediate in the PTX, so it belongs in the cache key. - let cache_key = format!("gdn_per_head_l2_norm_{head_dim}_{num_heads}_{epsilon:e}"); + // `{epsilon:e}` is not an integer: key it by its bits, format it once. + let cache_key = self.module_key( + ( + "gdn_per_head_l2_norm_{}_{}_{:e}", + key_dims(&[ + u64::from(head_dim), + u64::from(num_heads), + u64::from(epsilon.to_bits()), + ]), + ), + || format!("gdn_per_head_l2_norm_{head_dim}_{num_heads}_{epsilon:e}"), + ); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); @@ -354,7 +403,7 @@ impl CudaExecutor { ) -> Result<(), GpuError> { let kernel = trueno_gpu::kernels::gdn::GdnGatesKernel::new(num_heads); let kernel_type = KernelType::GdnGates { num_heads }; - let cache_key = format!("gdn_gates_{num_heads}"); + let cache_key = module_key!(self, "gdn_gates_{}", num_heads); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); @@ -412,8 +461,14 @@ impl CudaExecutor { num_k_heads, head_k_dim, }; - let cache_key = - format!("gdn_delta_rule_{num_v_heads}_{head_v_dim}_{num_k_heads}_{head_k_dim}"); + let cache_key = module_key!( + self, + "gdn_delta_rule_{}_{}_{}_{}", + num_v_heads, + head_v_dim, + num_k_heads, + head_k_dim + ); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); @@ -457,7 +512,18 @@ impl CudaExecutor { num_heads, epsilon, }; - let cache_key = format!("gdn_gated_rmsnorm_{head_dim}_{num_heads}_{epsilon:e}"); + // `{epsilon:e}` is not an integer: key it by its bits, format it once. + let cache_key = self.module_key( + ( + "gdn_gated_rmsnorm_{}_{}_{:e}", + key_dims(&[ + u64::from(head_dim), + u64::from(num_heads), + u64::from(epsilon.to_bits()), + ]), + ), + || format!("gdn_gated_rmsnorm_{head_dim}_{num_heads}_{epsilon:e}"), + ); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); @@ -488,7 +554,7 @@ impl CudaExecutor { ) -> Result<(), GpuError> { let kernel = trueno_gpu::kernels::gdn::SigmoidGateKernel::new(n); let kernel_type = KernelType::GdnSigmoidGate { n }; - let cache_key = format!("gdn_sigmoid_gate_{n}"); + let cache_key = module_key!(self, "gdn_sigmoid_gate_{}", n); let kernel_name = self.gdn_prepare(&kernel_type, &cache_key)?; let (gx, _, _) = kernel.grid(); let (bx, _, _) = kernel.block(); diff --git a/crates/aprender-serve/src/cuda/executor/layer_norm_gpu.rs b/crates/aprender-serve/src/cuda/executor/layer_norm_gpu.rs index 0164695b2e..d0ef5af7d2 100644 --- a/crates/aprender-serve/src/cuda/executor/layer_norm_gpu.rs +++ b/crates/aprender-serve/src/cuda/executor/layer_norm_gpu.rs @@ -208,7 +208,7 @@ impl CudaExecutor { hidden_size, epsilon, }, - format!("rmsnorm_simple_{}", hidden_size), + module_key!(self, "rmsnorm_simple_{}", hidden_size), ) } else if use_precise { ( @@ -216,7 +216,7 @@ impl CudaExecutor { hidden_size, epsilon, }, - format!("rmsnorm_precise_{}", hidden_size), + module_key!(self, "rmsnorm_precise_{}", hidden_size), ) } else { ( @@ -224,24 +224,24 @@ impl CudaExecutor { hidden_size, epsilon, }, - format!("rmsnorm_vectorized_{}", hidden_size), + module_key!(self, "rmsnorm_vectorized_{}", hidden_size), ) }; let kernel_name = self.kernels.kernel_name(&kernel_type); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); if use_simple { eprintln!("[GH-559] RmsNorm PTX ({} bytes)", ptx.len()); } let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // PAR-081: 256 threads for vectorized, 32 threads for simple @@ -268,7 +268,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -310,17 +310,17 @@ impl CudaExecutor { epsilon, }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("per_head_rmsnorm_{}_{}", head_dim, num_heads); + let cache_key = module_key!(self, "per_head_rmsnorm_{}_{}", head_dim, num_heads); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // One warp (32 threads) per head, one block per head @@ -350,7 +350,7 @@ impl CudaExecutor { // diverged at position 1 on sm_89 and GB10 alike (position 0 is immune: // attention over one key is V, and RoPE(0) is the identity). if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), diff --git a/crates/aprender-serve/src/cuda/executor/layers/par-062.rs b/crates/aprender-serve/src/cuda/executor/layers/par-062.rs index cd5e148db4..4c757e3a9d 100644 --- a/crates/aprender-serve/src/cuda/executor/layers/par-062.rs +++ b/crates/aprender-serve/src/cuda/executor/layers/par-062.rs @@ -141,6 +141,26 @@ impl CudaExecutor { self.argmax_num_blocks = num_blocks; } + // #4215: resolve both modules before borrowing the argmax buffers, so the + // cached keys can take `&mut self`. + // Load first-pass kernel module (cached after first use) + let argmax_kernel_type = KernelType::ArgMax { length: vocab_size }; + let argmax_key = module_key!(self, "argmax_{}", vocab_size); + if !self.modules.contains_key(&*argmax_key) { + let ptx = self.kernels.generate_ptx(&argmax_kernel_type); + let module = self.compile_ptx(&ptx)?; + self.modules.insert(argmax_key.to_string(), module); + } + + // Load second-pass kernel module (cached after first use) + let final_kernel_type = KernelType::ArgMaxFinal { num_blocks }; + let final_key = module_key!(self, "argmax_final_{}", num_blocks); + if !self.modules.contains_key(&*final_key) { + let ptx = self.kernels.generate_ptx(&final_kernel_type); + let module = self.compile_ptx(&ptx)?; + self.modules.insert(final_key.to_string(), module); + } + let block_max_vals = self .argmax_block_vals .as_ref() @@ -154,24 +174,6 @@ impl CudaExecutor { .as_ref() .expect("argmax_result must be initialized"); - // Load first-pass kernel module (cached after first use) - let argmax_kernel_type = KernelType::ArgMax { length: vocab_size }; - let argmax_key = format!("argmax_{}", vocab_size); - if !self.modules.contains_key(&argmax_key) { - let ptx = self.kernels.generate_ptx(&argmax_kernel_type); - let module = self.compile_ptx(&ptx)?; - self.modules.insert(argmax_key.clone(), module); - } - - // Load second-pass kernel module (cached after first use) - let final_kernel_type = KernelType::ArgMaxFinal { num_blocks }; - let final_key = format!("argmax_final_{}", num_blocks); - if !self.modules.contains_key(&final_key) { - let ptx = self.kernels.generate_ptx(&final_kernel_type); - let module = self.compile_ptx(&ptx)?; - self.modules.insert(final_key.clone(), module); - } - // Prepare kernel arguments let kernel_name = self.kernels.kernel_name(&argmax_kernel_type); // PAR-068-FIX: Do NOT use .with_shared_mem() - PTX declares static shared memory via .shared directive @@ -187,7 +189,7 @@ impl CudaExecutor { unsafe { let module = self .modules - .get_mut(&argmax_key) + .get_mut(&*argmax_key) .expect("argmax module just inserted"); self.stream.launch_kernel( module, @@ -216,7 +218,7 @@ impl CudaExecutor { unsafe { let final_module = self .modules - .get_mut(&final_key) + .get_mut(&*final_key) .expect("argmax_final module just inserted"); self.stream.launch_kernel( final_module, diff --git a/crates/aprender-serve/src/cuda/executor/mod.rs b/crates/aprender-serve/src/cuda/executor/mod.rs index ba746ce993..82d9b9322e 100644 --- a/crates/aprender-serve/src/cuda/executor/mod.rs +++ b/crates/aprender-serve/src/cuda/executor/mod.rs @@ -247,6 +247,55 @@ use crate::cuda::types::{ /// can recover — only a process restart fixes it. /// /// This check prevents kernel launch and returns a clean error instead. +/// #4215: the `modules` cache key `format!($fmt, $args…)`, formatted the first +/// time these arguments are seen and served from `module_keys` afterwards, so a +/// hot launch path neither allocates nor formats. Every argument must convert +/// to `u64` losslessly (pass an `f32` as `.to_bits()`, and format it yourself if +/// the key spells it differently); at most four. +macro_rules! module_key { + ($ex:expr, $fmt:literal $(, $arg:expr)* $(,)?) => {{ + let dims = $crate::cuda::executor::key_dims(&[$(u64::from($arg)),*]); + $ex.module_key(($fmt, dims), || format!($fmt $(, $arg)*)) + }}; +} +pub(crate) use module_key; + +/// Pad a cache key's integer arguments to the fixed width `module_keys` stores. +#[inline] +pub(crate) fn key_dims(args: &[u64]) -> [u64; 4] { + assert!(args.len() <= 4, "module_key!: at most four key arguments"); + let mut dims = [0u64; 4]; + dims[..args.len()].copy_from_slice(args); + dims +} + +impl CudaExecutor { + /// The formatted `modules` key for `id`, calling `make` only on the first + /// request for it (#4215). Use through [`module_key!`]. + pub(crate) fn module_key( + &mut self, + id: (&'static str, [u64; 4]), + make: impl FnOnce() -> String, + ) -> std::sync::Arc { + if let Some(key) = self.module_keys.get(&id) { + return std::sync::Arc::clone(key); + } + let key: std::sync::Arc = make().into(); + self.module_keys.insert(id, std::sync::Arc::clone(&key)); + key + } +} + +/// [`validate_device_ptr`] for argument `index` of `kernel`, naming it only on +/// failure — the name costs a `format!`, and the check runs on every launch. +#[inline] +fn validate_kernel_arg(ptr: u64, kernel: &str, index: usize) -> Result<(), GpuError> { + if ptr == 0 { + return validate_device_ptr(ptr, &format!("{kernel} arg {index}")); + } + Ok(()) +} + #[inline] fn validate_device_ptr(ptr: u64, name: &str) -> Result<(), GpuError> { if ptr == 0 { @@ -353,6 +402,10 @@ pub struct CudaExecutor { // Thousands of cuModuleUnload cycles exhaust the CUDA driver. // Leaked modules (~KB each) are cleaned up at process exit. modules: std::mem::ManuallyDrop>, + /// #4215: `modules` keys already formatted, by (format string, integer + /// arguments) — see [`module_key!`]. A decode token looks its keys up here + /// instead of `format!`-ing ~600 of them. + module_keys: HashMap<(&'static str, [u64; 4]), std::sync::Arc>, // Persistent weight buffers on GPU (PARITY-037) // These are loaded once at startup and reused for all forward passes weight_cache: HashMap>, diff --git a/crates/aprender-serve/src/cuda/executor/q4k_mwv_gemv.rs b/crates/aprender-serve/src/cuda/executor/q4k_mwv_gemv.rs index c481080c9c..bcb28cc570 100644 --- a/crates/aprender-serve/src/cuda/executor/q4k_mwv_gemv.rs +++ b/crates/aprender-serve/src/cuda/executor/q4k_mwv_gemv.rs @@ -17,17 +17,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::MwvQ4KGemv { k, n, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("mwv_q4k_gemv_{}_{}_{}", k, n, num_warps); + let cache_key = module_key!(self, "mwv_q4k_gemv_{}_{}_{}", k, n, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // num_warps * 32 threads per output element, one block per output @@ -75,7 +75,7 @@ impl CudaExecutor { // was unaffected because it routes Q4K to the (recorded) DP4A variant. // Fix: record this GEMV too, mirroring mwv_dp4a_q4k_gemv_into. if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -135,17 +135,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::MwvDp4aQ4KGemv { k, n, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("mwv_dp4a_q4k_gemv_{}_{}_{}", k, n, num_warps); + let cache_key = module_key!(self, "mwv_dp4a_q4k_gemv_{}_{}_{}", k, n, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -182,7 +182,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction // BUG FIX (realizr#198): MWV DP4A variant was missing recording. if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -240,17 +240,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::HwDp4aQ4KGemv { k, n, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("hw_dp4a_q4k_gemv_{}_{}_{}", k, n, num_warps); + let cache_key = module_key!(self, "hw_dp4a_q4k_gemv_{}_{}_{}", k, n, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -281,7 +281,7 @@ impl CudaExecutor { // trueno#243: Record kernel AFTER launch (avoids borrow conflict with module) if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -319,17 +319,17 @@ impl CudaExecutor { validate_device_ptr(weight_ptr, "dp4a_q4k_gemv_into")?; let kernel_type = KernelType::Dp4aQ4KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("dp4a_q4k_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "dp4a_q4k_gemv_{}_{}", k, n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // One warp (32 threads) per output element @@ -399,15 +399,15 @@ impl CudaExecutor { let kernel_name = self.kernels.kernel_name(&kernel_type); let cache_key = format!("fused_rmsnorm_q4k_gemv_{}_{}_{:.0e}", k, n, epsilon); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // One block per output element, 256 threads per block @@ -483,17 +483,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::FusedGateUpSwigluHwDp4aQ4KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("fused_gate_up_swiglu_hw_dp4a_q4k_{}_{}", k, n); + let cache_key = module_key!(self, "fused_gate_up_swiglu_hw_dp4a_q4k_{}_{}", k, n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -529,7 +529,7 @@ impl CudaExecutor { // kernels per forward pass were never captured. Graph replay used stale // ffn_act_buf from the eager pass, causing hidden_buf2 divergence (81.0). if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -625,17 +625,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::FusedKVHwDp4aQ4KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("fused_kv_hw_dp4a_q4k_{}_{}", k, n); + let cache_key = module_key!(self, "fused_kv_hw_dp4a_q4k_{}_{}", k, n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -736,17 +736,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::BatchedHwDp4aQ4KGemv { k, n, m, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("batched_hw_dp4a_q4k_gemv_{}_{}_{}_{}", k, n, m, num_warps); + let cache_key = module_key!(self, "batched_hw_dp4a_q4k_gemv_{}_{}_{}_{}", k, n, m, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -798,17 +798,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::BatchedHwDp4aQ4KGemv { k, n, m, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("batched_hw_dp4a_q4k_gemv_{}_{}_{}_{}", k, n, m, num_warps); + let cache_key = module_key!(self, "batched_hw_dp4a_q4k_gemv_{}_{}_{}_{}", k, n, m, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -860,17 +860,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::InlineQ8Dp4aQ4KGemv { k, n, m, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("inline_q8_dp4a_q4k_gemv_{}_{}_{}_{}", k, n, m, num_warps); + let cache_key = module_key!(self, "inline_q8_dp4a_q4k_gemv_{}_{}_{}_{}", k, n, m, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -922,17 +922,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::FusedFp32Q4KGemv { k, n, m, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("fused_fp32_q4k_gemv_{}_{}_{}_{}", k, n, m, num_warps); + let cache_key = module_key!(self, "fused_fp32_q4k_gemv_{}_{}_{}_{}", k, n, m, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -1145,17 +1145,17 @@ impl CudaExecutor { validate_device_ptr(up_weight_ptr, "fused_gate_up_q4k_gemv_into(up)")?; let kernel_type = KernelType::FusedGateUpQ4KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("fused_gate_up_q4k_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "fused_gate_up_q4k_gemv_{}_{}", k, n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // One block per output element, 256 threads per block diff --git a/crates/aprender-serve/src/cuda/executor/q6k_gemv_indexed.rs b/crates/aprender-serve/src/cuda/executor/q6k_gemv_indexed.rs index 060c0a1310..a22b2c3e13 100644 --- a/crates/aprender-serve/src/cuda/executor/q6k_gemv_indexed.rs +++ b/crates/aprender-serve/src/cuda/executor/q6k_gemv_indexed.rs @@ -57,15 +57,15 @@ impl CudaExecutor { }; let kernel_name = self.kernels.kernel_name(&kernel_type); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // Allocate output buffer @@ -166,15 +166,15 @@ impl CudaExecutor { }; let kernel_name = self.kernels.kernel_name(&kernel_type); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -228,17 +228,17 @@ impl CudaExecutor { validate_device_ptr(weight_ptr, "coalesced_q4k_gemv_into")?; let kernel_type = KernelType::CoalescedQ4KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("coalesced_q4k_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "coalesced_q4k_gemv_{}_{}", k, n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // One warp (32 threads) per output element @@ -289,17 +289,17 @@ impl CudaExecutor { validate_device_ptr(weight_ptr, "wide_q4k_gemv_into")?; let kernel_type = KernelType::WideQ4KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("wide_q4k_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "wide_q4k_gemv_{}_{}", k, n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // PAR-132: 8 warps (256 threads) per output element @@ -357,17 +357,17 @@ impl CudaExecutor { validate_device_ptr(weight_ptr, "vectorized_q4k_gemv_into")?; let kernel_type = KernelType::VectorizedQ4KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("vectorized_q4k_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "vectorized_q4k_gemv_{}_{}", k, n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // One warp (32 threads) per output element diff --git a/crates/aprender-serve/src/cuda/executor/residual.rs b/crates/aprender-serve/src/cuda/executor/residual.rs index 31ded422d6..39ed884c23 100644 --- a/crates/aprender-serve/src/cuda/executor/residual.rs +++ b/crates/aprender-serve/src/cuda/executor/residual.rs @@ -312,17 +312,17 @@ impl CudaExecutor { let kernel_type = KernelType::ResidualAdd { n }; let kernel_name = self.kernels.kernel_name(&kernel_type); // GH-129: PTX is n-independent (n is a runtime param), so use constant cache key. - let cache_key = "residual_add".to_string(); + let cache_key = "residual_add"; - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads_per_block = 256u32; @@ -351,7 +351,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), diff --git a/crates/aprender-serve/src/cuda/executor/silu.rs b/crates/aprender-serve/src/cuda/executor/silu.rs index db1bce1953..ed7dc8df99 100644 --- a/crates/aprender-serve/src/cuda/executor/silu.rs +++ b/crates/aprender-serve/src/cuda/executor/silu.rs @@ -237,17 +237,17 @@ impl CudaExecutor { ) -> Result<(), GpuError> { let kernel_type = KernelType::FusedSwiglu { n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("fused_swiglu_{}", n); + let cache_key = module_key!(self, "fused_swiglu_{}", n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = 256; @@ -276,7 +276,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), diff --git a/crates/aprender-serve/src/cuda/executor/weight.rs b/crates/aprender-serve/src/cuda/executor/weight.rs index c84c9b5cd4..4a696dda1c 100644 --- a/crates/aprender-serve/src/cuda/executor/weight.rs +++ b/crates/aprender-serve/src/cuda/executor/weight.rs @@ -19,17 +19,17 @@ impl CudaExecutor { let kernel_type = KernelType::BatchedQ6KGemv { k, n, m }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("batched_q6k_gemv_{}_{}_{}", m, k, n); + let cache_key = module_key!(self, "batched_q6k_gemv_{}_{}_{}", m, k, n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // Grid: N blocks (one per output row), 32 threads per block @@ -107,18 +107,18 @@ impl CudaExecutor { // Original Q6K kernel (CoalescedQ6K disabled due to CORRECTNESS-006) let kernel_type = KernelType::Q6KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("q6k_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "q6k_gemv_{}_{}", k, n); let config = LaunchConfig::grid_2d(n, 1, 32, 1); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -145,7 +145,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -183,17 +183,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::MwvQ6KGemv { k, n, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("mwv_q6k_gemv_{}_{}_{}", k, n, num_warps); + let cache_key = module_key!(self, "mwv_q6k_gemv_{}_{}_{}", k, n, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -224,7 +224,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -280,17 +280,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::Dp4aQ6KGemv { k, n, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("dp4a_q6k_gemv_{}_{}_{}", k, n, num_warps); + let cache_key = module_key!(self, "dp4a_q6k_gemv_{}_{}_{}", k, n, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -323,7 +323,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -373,17 +373,17 @@ impl CudaExecutor { let num_warps = self.gpu_profile.mwv_warps; let kernel_type = KernelType::HwDp4aQ6KGemv { k, n, num_warps }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("hw_dp4a_q6k_gemv_{}_{}_{}", k, n, num_warps); + let cache_key = module_key!(self, "hw_dp4a_q6k_gemv_{}_{}_{}", k, n, num_warps); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let threads = num_warps * 32; @@ -414,7 +414,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -453,18 +453,18 @@ impl CudaExecutor { validate_device_ptr(weight_ptr, "coalesced_q6k_gemv_into")?; let kernel_type = KernelType::CoalescedQ6KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("coalesced_q6k_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "coalesced_q6k_gemv_{}_{}", k, n); let config = LaunchConfig::grid_2d(n, 1, 32, 1); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -519,18 +519,18 @@ impl CudaExecutor { // PAR-058: Zero allocation Q8_0 GEMV for mixed-quantization models let kernel_type = KernelType::Q8_0Gemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("q8_0_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "q8_0_gemv_{}_{}", k, n); let config = LaunchConfig::grid_2d(n, 1, 32, 1); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -557,7 +557,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -596,18 +596,18 @@ impl CudaExecutor { // PAR-058: Zero allocation Q5_0 GEMV for Qwen 0.5B Q/K weights let kernel_type = KernelType::Q5_0Gemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("q5_0_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "q5_0_gemv_{}_{}", k, n); let config = LaunchConfig::grid_2d(n, 1, 32, 1); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -634,7 +634,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -673,18 +673,18 @@ impl CudaExecutor { // PAR-058: Zero allocation Q4_0 GEMV for GGUF qtype mismatch let kernel_type = KernelType::Q4_0Gemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("q4_0_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "q4_0_gemv_{}_{}", k, n); let config = LaunchConfig::grid_2d(n, 1, 32, 1); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -711,7 +711,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -742,18 +742,18 @@ impl CudaExecutor { validate_device_ptr(weight_ptr, "f32_gemv_into")?; let kernel_type = KernelType::Gemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("f32_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "f32_gemv_{}_{}", k, n); let config = LaunchConfig::grid_2d(n, 1, 32, 1); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -780,7 +780,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), diff --git a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs index 8ff1f85732..bf9f77b22b 100644 --- a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs +++ b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_cuda.rs @@ -257,6 +257,10 @@ pub struct Qwen35CudaModel<'a> { out_normed: GpuBuffer, /// `[vocab_size]`, the logits — the ONE buffer a token's forward downloads. logits_buf: GpuBuffer, + /// `[hidden_dim]`, the token's residual stream (#4215): the embedding row is + /// copied in, every layer updates it in place. `None` only while a forward + /// has it out. + hidden_buf: Option>, dims: Qwen35CudaDims, /// Positions the device KV caches hold. max_seq_len: usize, @@ -589,6 +593,7 @@ impl<'a> Qwen35CudaModel<'a> { let attn_scratch = Self::build_attn_scratch(&executor, dims)?; let out_normed = Self::zeros(&executor, dims.hidden_dim as usize)?; let logits_buf = Self::zeros(&executor, dims.vocab_size as usize)?; + let hidden_buf = Some(Self::zeros(&executor, dims.hidden_dim as usize)?); let state = Self::build_state(&executor, &layers, dims, max_seq_len)?; Ok(Self { model, @@ -601,6 +606,7 @@ impl<'a> Qwen35CudaModel<'a> { lm_head, out_normed, logits_buf, + hidden_buf, dims, max_seq_len, }) @@ -1388,6 +1394,51 @@ impl<'a> Qwen35CudaModel<'a> { state: &mut Qwen35CudaState, position: usize, ) -> Result> { + self.forward_to_logits(token, state, position)?; + // The ONE sync of the whole token, in front of the ONE download. + self.executor + .sync_stream() + .map_err(|e| gpu_err("qwen35_cuda_forward", &e))?; + let mut logits = vec![0.0f32; self.dims.vocab_size as usize]; + self.logits_buf + .copy_to_host(&mut logits) + .map_err(|e| gpu_err("qwen35_cuda_forward", &e))?; + + state.kv_len = state.kv_len.max(position + 1); + Ok(logits) + } + + /// The token temperature 0 picks after `token` at `position` (#4215): the + /// argmax runs on the device over `logits_buf`, so the token downloads its + /// 4-byte id instead of the whole logits vector. Ties go to the LOWEST + /// index, as [`crate::gguf::ops::argmax`] does. + /// + /// # Errors + /// As [`Self::forward_single`]. + pub fn forward_single_greedy( + &mut self, + token: u32, + state: &mut Qwen35CudaState, + position: usize, + ) -> Result { + self.forward_to_logits(token, state, position)?; + let next = self + .executor + .gpu_argmax(self.logits_buf.as_ptr(), self.dims.vocab_size) + .map_err(|e| gpu_err("qwen35_cuda_greedy", &e))?; + state.kv_len = state.kv_len.max(position + 1); + Ok(next) + } + + /// Enqueue one token at `position` through every layer, the output norm and + /// the `lm_head` into `logits_buf`. Nothing is synced or downloaded, and the + /// caller advances `state.kv_len` once it has read the result. + fn forward_to_logits( + &mut self, + token: u32, + state: &mut Qwen35CudaState, + position: usize, + ) -> Result<()> { let hidden_dim = self.dims.hidden_dim as usize; let embedding = self.model.base.token_embedding(); let start = (token as usize) * hidden_dim; @@ -1408,15 +1459,37 @@ impl<'a> Qwen35CudaModel<'a> { }); } - let dev = GpuBuffer::from_host( - self.executor.context(), + // #4215: the residual stream lives in one persistent buffer; taking it + // out of `self` lets the layers borrow `self` mutably beside it. + let mut dev = self + .hidden_buf + .take() + .ok_or_else(|| RealizarError::InvalidShape { + reason: "qwen35_cuda: the hidden buffer is already in use".to_string(), + })?; + let run = self.run_layers_and_head( + &mut dev, &embedding[start..start + hidden_dim], - ) - .map_err(|e| gpu_err("qwen35_cuda_forward", &e))?; + state, + position, + ); + self.hidden_buf = Some(dev); + run + } + + fn run_layers_and_head( + &mut self, + dev: &mut GpuBuffer, + embedding_row: &[f32], + state: &mut Qwen35CudaState, + position: usize, + ) -> Result<()> { + dev.copy_from_host(embedding_row) + .map_err(|e| gpu_err("qwen35_cuda_forward", &e))?; for il in 0..self.layers.len() { match self.layers[il] { - CudaLayer::DeltaNet(_) => self.deltanet_layer(state, il, &dev)?, - CudaLayer::Attention(_) => self.attention_layer(state, il, &dev, position)?, + CudaLayer::DeltaNet(_) => self.deltanet_layer(state, il, dev)?, + CudaLayer::Attention(_) => self.attention_layer(state, il, dev, position)?, } } // The tail stays on the device (#3090 review). `hidden_to_logits` would @@ -1429,7 +1502,7 @@ impl<'a> Qwen35CudaModel<'a> { let d = self.dims; self.executor .rmsnorm_into( - &dev, + dev, &self.output_norm, &self.out_normed, d.hidden_dim, @@ -1446,18 +1519,7 @@ impl<'a> Qwen35CudaModel<'a> { self.lm_head.k, ) .map_err(|e| gpu_err("qwen35_cuda_lm_head", &e))?; - - // The ONE sync of the whole token, in front of the ONE download. - self.executor - .sync_stream() - .map_err(|e| gpu_err("qwen35_cuda_forward", &e))?; - let mut logits = vec![0.0f32; d.vocab_size as usize]; - self.logits_buf - .copy_to_host(&mut logits) - .map_err(|e| gpu_err("qwen35_cuda_forward", &e))?; - - state.kv_len = state.kv_len.max(position + 1); - Ok(logits) + Ok(()) } /// Run every Gated `DeltaNet` layer over a host hidden state and read the @@ -1499,3 +1561,8 @@ impl<'a> Qwen35CudaModel<'a> { #[cfg(test)] #[path = "forward_qwen35_cuda_tests.rs"] mod qwen35_cuda_tests; + +/// #4215: host allocations, device allocations and D2H bytes per decode token. +#[cfg(test)] +#[path = "forward_qwen35_decode_overhead_tests.rs"] +mod qwen35_decode_overhead_tests; diff --git a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_decode_overhead_tests.rs b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_decode_overhead_tests.rs new file mode 100644 index 0000000000..ab7460adf7 --- /dev/null +++ b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_decode_overhead_tests.rs @@ -0,0 +1,189 @@ +//! #4215: the host cost of one steady-state Qwen3.5 greedy decode token. +//! +//! Decode is close to launch-bound (~634 launches/token), so the host work +//! around each launch is on the critical path. Before #4215 every token +//! allocated a device buffer for its embedding row, downloaded the whole +//! 248,320-entry logits vector (~1 MB) to run the argmax on the CPU, and +//! allocated on the host around every GDN launch (a `format!` per pointer, the +//! argument `Vec`s, a `format!` cache key per wrapper). +//! +//! These tests read three per-thread counters across steady-state tokens (after +//! warm-up has compiled every module and sized every lazy buffer): host heap +//! allocations (a counting global allocator), device allocations and +//! device→host bytes (`trueno_gpu::driver`'s `GpuBuffer` counters). + +use super::Qwen35CudaModel; +use crate::gguf::forward_qwen35::Qwen35Model; +use std::alloc::{GlobalAlloc, Layout, System}; +use std::cell::Cell; +use trueno_gpu::driver::{device_allocs_total, device_to_host_bytes_total}; + +const MODEL_PATH: &str = "/home/noah/models/Qwen3.5-0.8B-Q4_K_M.gguf"; + +/// Prompt ids (all well inside the 248,320-row vocabulary). +const PROMPT: [u32; 6] = [9707, 11, 1879, 0, 3555, 374]; + +/// Steady-state tokens measured, after the prompt and `WARMUP` decode tokens. +const MEASURED: u64 = 4; +const WARMUP: usize = 2; + +/// Counts every heap allocation the current thread makes; forwards to `System`. +struct CountingAlloc; + +thread_local! { + static HOST_ALLOCS: Cell = const { Cell::new(0) }; +} + +fn count_host_alloc() { + // `try_with`: an allocation during thread-local teardown is simply not counted. + let _ = HOST_ALLOCS.try_with(|c| c.set(c.get() + 1)); +} + +// SAFETY: every method forwards to `System` with the caller's arguments unchanged; +// the only addition is a thread-local counter increment, which never allocates. +unsafe impl GlobalAlloc for CountingAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + count_host_alloc(); + // SAFETY: forwarded verbatim; the caller upholds `GlobalAlloc::alloc`'s contract. + unsafe { System.alloc(layout) } + } + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + count_host_alloc(); + // SAFETY: forwarded verbatim. + unsafe { System.alloc_zeroed(layout) } + } + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + count_host_alloc(); + // SAFETY: forwarded verbatim. + unsafe { System.realloc(ptr, layout, new_size) } + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + // SAFETY: forwarded verbatim. + unsafe { System.dealloc(ptr, layout) } + } +} + +#[global_allocator] +static COUNTING_ALLOC: CountingAlloc = CountingAlloc; + +fn host_allocs_total() -> u64 { + HOST_ALLOCS.with(Cell::get) +} + +/// One greedy decode step: the token temperature 0 picks after `token` at `position`. +fn greedy_step( + gpu: &mut Qwen35CudaModel<'_>, + token: u32, + state: &mut super::Qwen35CudaState, + position: usize, +) -> u32 { + gpu.forward_single_greedy(token, state, position) + .expect("greedy forward") +} + +macro_rules! model_or_skip { + () => {{ + if !std::path::Path::new(MODEL_PATH).exists() { + eprintln!("SKIP: {MODEL_PATH} is absent"); + return; + } + crate::cuda_executor_or_skip!(0) + }}; +} + +/// A steady-state greedy token makes no host or device allocation and downloads +/// exactly the 4-byte token id. +#[test] +#[serial_test::serial] +fn a_steady_state_greedy_decode_token_allocates_nothing_and_downloads_only_the_token_id() { + let executor = model_or_skip!(); + let mapped = crate::gguf::MappedGGUFModel::from_path(MODEL_PATH).expect("map the GGUF"); + let base = Qwen35Model::create_base_model(&mapped.model, mapped.data()).expect("base"); + let qwen = + Qwen35Model::from_model_and_layers(&base, &mapped.model, mapped.data()).expect("qwen35"); + let mut gpu = Qwen35CudaModel::new(&qwen, executor).expect("build the CUDA model"); + let mut state = gpu.new_state().expect("device state"); + + let mut pos = 0; + let mut next = 0; + for &t in &PROMPT { + next = greedy_step(&mut gpu, t, &mut state, pos); + pos += 1; + } + for _ in 0..WARMUP { + next = greedy_step(&mut gpu, next, &mut state, pos); + pos += 1; + } + + let (h0, g0, d0) = ( + host_allocs_total(), + device_allocs_total(), + device_to_host_bytes_total(), + ); + for _ in 0..MEASURED { + next = greedy_step(&mut gpu, next, &mut state, pos); + pos += 1; + } + let (h1, g1, d1) = ( + host_allocs_total(), + device_allocs_total(), + device_to_host_bytes_total(), + ); + let (host, dev, d2h) = (h1 - h0, g1 - g0, d1 - d0); + eprintln!( + "[4215] per steady-state token: host allocs {:.1}, device allocs {:.1}, D2H bytes {:.1}", + host as f64 / MEASURED as f64, + dev as f64 / MEASURED as f64, + d2h as f64 / MEASURED as f64, + ); + assert_eq!( + dev, 0, + "device allocations across {MEASURED} steady-state tokens" + ); + assert_eq!( + d2h, + 4 * MEASURED, + "device→host bytes across {MEASURED} tokens: one u32 token id each" + ); + assert_eq!( + host, 0, + "host heap allocations across {MEASURED} steady-state tokens" + ); +} + +/// Temperature 0 on the device picks exactly the token the CPU argmax of the full +/// logits picks, token after token. +#[test] +#[serial_test::serial] +fn greedy_on_the_device_is_token_identical_to_the_cpu_argmax_of_the_logits() { + let executor = model_or_skip!(); + let mapped = crate::gguf::MappedGGUFModel::from_path(MODEL_PATH).expect("map the GGUF"); + let base = Qwen35Model::create_base_model(&mapped.model, mapped.data()).expect("base"); + let qwen = + Qwen35Model::from_model_and_layers(&base, &mapped.model, mapped.data()).expect("qwen35"); + let mut gpu = Qwen35CudaModel::new(&qwen, executor).expect("build the CUDA model"); + let mut logits_state = gpu.new_state().expect("device state"); + let mut greedy_state = gpu.new_state().expect("device state"); + + let steps = 32; + let (mut via_logits, mut via_greedy) = (Vec::new(), Vec::new()); + let (mut a, mut b) = (PROMPT[0], PROMPT[0]); + for pos in 0..PROMPT.len() + steps { + let (ta, tb) = match PROMPT.get(pos) { + Some(&t) => (t, t), + None => (a, b), + }; + let logits = gpu + .forward_single(ta, &mut logits_state, pos) + .expect("logits forward"); + a = crate::gguf::ops::argmax(&logits); + b = greedy_step(&mut gpu, tb, &mut greedy_state, pos); + via_logits.push(a); + via_greedy.push(b); + } + eprintln!("[4215] greedy tokens: {via_greedy:?}"); + assert_eq!( + via_greedy, via_logits, + "device greedy diverged from the CPU argmax of the downloaded logits" + ); +} diff --git a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs index ec870db618..7766430d73 100644 --- a/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs +++ b/crates/aprender-serve/src/gguf/inference/forward/forward_qwen35.rs @@ -1454,6 +1454,9 @@ fn qwen35_gpu_decode( let mut state = gpu .new_state() .map_err(|e| format!("the decode state would not allocate: {e}"))?; + if gen_config.temperature == 0.0 || gen_config.top_k == 1 { + return qwen35_gpu_decode_greedy(gpu, &mut state, input_tokens, gen_config, max_seq_len); + } let mut rng = rand::rngs::StdRng::seed_from_u64(gen_config.seed); let mut logits = Vec::new(); @@ -1464,17 +1467,13 @@ fn qwen35_gpu_decode( } let mut tokens = input_tokens.to_vec(); for _ in 0..gen_config.max_tokens { - let next = if gen_config.temperature == 0.0 || gen_config.top_k == 1 { - crate::gguf::ops::argmax(&logits) - } else { - OwnedQuantizedModel::sample_topk_seeded( - &logits, - gen_config.temperature, - gen_config.top_k, - gen_config.top_p, - &mut rng, - ) - }; + let next = OwnedQuantizedModel::sample_topk_seeded( + &logits, + gen_config.temperature, + gen_config.top_k, + gen_config.top_p, + &mut rng, + ); tokens.push(next); if gen_config.stop_tokens.contains(&next) || tokens.len() >= max_seq_len { break; @@ -1487,6 +1486,36 @@ fn qwen35_gpu_decode( Ok(tokens) } +/// Temperature 0 (#4215): the argmax runs on the device, so each token +/// downloads its 4-byte id instead of the 248,320-entry logits vector. +#[cfg(feature = "cuda")] +fn qwen35_gpu_decode_greedy( + gpu: &mut crate::gguf::cuda::Qwen35CudaModel<'_>, + state: &mut crate::gguf::cuda::Qwen35CudaState, + input_tokens: &[u32], + gen_config: &crate::gguf::QuantizedGenerateConfig, + max_seq_len: usize, +) -> std::result::Result, String> { + let mut next = 0; + for (pos, &token) in input_tokens.iter().enumerate() { + next = gpu + .forward_single_greedy(token, state, pos) + .map_err(|e| format!("the GPU forward failed at prompt position {pos}: {e}"))?; + } + let mut tokens = input_tokens.to_vec(); + for _ in 0..gen_config.max_tokens { + tokens.push(next); + if gen_config.stop_tokens.contains(&next) || tokens.len() >= max_seq_len { + break; + } + let pos = tokens.len() - 1; + next = gpu + .forward_single_greedy(next, state, pos) + .map_err(|e| format!("the GPU forward failed at decode position {pos}: {e}"))?; + } + Ok(tokens) +} + /// The F2 runtime guard for the hybrid: forward the real prompt through BOTH /// backends and accept the GPU only if every real position agrees. /// From f08fce510fc2ed5ace5e8735e6118f544a3d4470 Mon Sep 17 00:00:00 2001 From: Noah Gift Date: Thu, 24 Sep 2026 14:16:29 +0200 Subject: [PATCH 3/4] test(qwen35): decode tok/s at 850 bench, ignored (Refs #4215) Co-Authored-By: Claude Opus 5.5 (1M context) --- .../forward_qwen35_decode_overhead_tests.rs | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_decode_overhead_tests.rs b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_decode_overhead_tests.rs index ab7460adf7..ebefa71bb1 100644 --- a/crates/aprender-serve/src/gguf/cuda/forward_qwen35_decode_overhead_tests.rs +++ b/crates/aprender-serve/src/gguf/cuda/forward_qwen35_decode_overhead_tests.rs @@ -187,3 +187,42 @@ fn greedy_on_the_device_is_token_identical_to_the_cpu_argmax_of_the_logits() { "device greedy diverged from the CPU argmax of the downloaded logits" ); } + +/// Decode tokens/s after an 850-token prompt (#4215's reported number). Timing, +/// not a gate: run in release with `--ignored`; `APR_4215_MODEL` picks the GGUF. +#[test] +#[ignore = "timing: run in release with --ignored"] +#[serial_test::serial] +fn bench_greedy_decode_tok_s_at_850() { + let path = std::env::var("APR_4215_MODEL").unwrap_or_else(|_| MODEL_PATH.to_string()); + if !std::path::Path::new(&path).exists() { + eprintln!("SKIP: {path} is absent"); + return; + } + let executor = crate::cuda_executor_or_skip!(0); + let mapped = crate::gguf::MappedGGUFModel::from_path(&path).expect("map the GGUF"); + let base = Qwen35Model::create_base_model(&mapped.model, mapped.data()).expect("base"); + let qwen = + Qwen35Model::from_model_and_layers(&base, &mapped.model, mapped.data()).expect("qwen35"); + let mut gpu = + Qwen35CudaModel::with_max_seq_len(&qwen, executor, 1024).expect("build the CUDA model"); + let mut state = gpu.new_state().expect("device state"); + + const PROMPT_LEN: usize = 850; + const DECODE: usize = 128; + let mut next = 0; + for pos in 0..PROMPT_LEN { + let t = PROMPT[pos % PROMPT.len()] + (pos as u32 % 97) * 13; + next = greedy_step(&mut gpu, t, &mut state, pos); + } + let t0 = std::time::Instant::now(); + for i in 0..DECODE { + next = greedy_step(&mut gpu, next, &mut state, PROMPT_LEN + i); + } + let secs = t0.elapsed().as_secs_f64(); + eprintln!( + "[4215] {path}: decode {:.1} tok/s at {PROMPT_LEN} ({DECODE} tokens, {:.3} ms/token)", + DECODE as f64 / secs, + secs * 1e3 / DECODE as f64 + ); +} From f80b3f82e2fe862eea432a8f477d529f7d4d7933 Mon Sep 17 00:00:00 2001 From: Noah Gift Date: Thu, 24 Sep 2026 21:30:22 +0200 Subject: [PATCH 4/4] perf(qwen35): a greedy decode token allocates nothing on host or device and downloads 4 bytes (Refs #4215) Remaining host allocations were the format! module keys in q5k_gemv_into (two DeltaNet projections per layer, 36/token), q4_1_gemv_into and q8_quantize_{async,into}; all now use module_key!. Measured on the 4090: host allocs 1713 -> 0, device allocs 1 -> 0, D2H 993,280 -> 4 bytes/token. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../src/cuda/executor/device.rs | 20 +++++++++---------- .../src/cuda/executor/tiled_gemv.rs | 18 ++++++++--------- docs/roadmaps/entries/PMAT-4215.yaml | 19 ++++++++++++++++++ docs/roadmaps/roadmap.yaml | 19 ++++++++++++++++++ 4 files changed, 57 insertions(+), 19 deletions(-) create mode 100644 docs/roadmaps/entries/PMAT-4215.yaml diff --git a/crates/aprender-serve/src/cuda/executor/device.rs b/crates/aprender-serve/src/cuda/executor/device.rs index ec12078736..c96307536f 100644 --- a/crates/aprender-serve/src/cuda/executor/device.rs +++ b/crates/aprender-serve/src/cuda/executor/device.rs @@ -29,18 +29,18 @@ impl CudaExecutor { // PAR-058: Zero allocation Q4_1 GEMV for Qwen2.5-0.5B FFN down let kernel_type = KernelType::Q4_1Gemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("q4_1_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "q4_1_gemv_{}_{}", k, n); let config = LaunchConfig::grid_2d(n, 1, 32, 1); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -67,7 +67,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), @@ -106,18 +106,18 @@ impl CudaExecutor { // PAR-058: Zero allocation Q5K GEMV for mixed-quantization models let kernel_type = KernelType::Q5KGemv { k, n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("q5k_gemv_{}_{}", k, n); + let cache_key = module_key!(self, "q5k_gemv_{}_{}", k, n); let config = LaunchConfig::grid_2d(n, 1, 32, 1); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let mut ptr_output = output.as_ptr(); @@ -144,7 +144,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), diff --git a/crates/aprender-serve/src/cuda/executor/tiled_gemv.rs b/crates/aprender-serve/src/cuda/executor/tiled_gemv.rs index 5f0f362390..3928d3fcec 100644 --- a/crates/aprender-serve/src/cuda/executor/tiled_gemv.rs +++ b/crates/aprender-serve/src/cuda/executor/tiled_gemv.rs @@ -290,17 +290,17 @@ impl CudaExecutor { // Load kernel module let kernel_type = KernelType::Q8Quantize { n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("q8_quantize_{}", n); + let cache_key = module_key!(self, "q8_quantize_{}", n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); // Q8_1 format: 36 bytes per 32 values @@ -350,17 +350,17 @@ impl CudaExecutor { ) -> Result<(), GpuError> { let kernel_type = KernelType::Q8Quantize { n }; let kernel_name = self.kernels.kernel_name(&kernel_type); - let cache_key = format!("q8_quantize_{}", n); + let cache_key = module_key!(self, "q8_quantize_{}", n); - if !self.modules.contains_key(&cache_key) { + if !self.modules.contains_key(&*cache_key) { let ptx = self.kernels.generate_ptx(&kernel_type); let module = self.compile_ptx(&ptx)?; - self.modules.insert(cache_key.clone(), module); + self.modules.insert(cache_key.to_string(), module); } let module = self .modules - .get_mut(&cache_key) + .get_mut(&*cache_key) .expect("module just inserted"); let num_blocks = (n + 31) / 32; @@ -386,7 +386,7 @@ impl CudaExecutor { // trueno#243: Record kernel for manual graph construction if self.graph_recording { - let module = self.modules.get_mut(&cache_key).expect("module exists"); + let module = self.modules.get_mut(&*cache_key).expect("module exists"); let func = module.get_function(kernel_name)?; self.graph_recorded_kernels.push(RecordedKernel { func: SendCUfunction(func), diff --git a/docs/roadmaps/entries/PMAT-4215.yaml b/docs/roadmaps/entries/PMAT-4215.yaml new file mode 100644 index 0000000000..3af47a7c7e --- /dev/null +++ b/docs/roadmaps/entries/PMAT-4215.yaml @@ -0,0 +1,19 @@ +- id: PMAT-4215 + github_issue: 4215 + item_type: task + title: 'Qwen3.5 CUDA greedy decode: persistent embedding buffer, GPU argmax, no format!/Vec on the launch path' + status: inprogress + priority: high + assigned_to: null + created: 2026-09-24T08:28:19Z + updated: 2026-09-24T20:45:00Z + spec: null + acceptance_criteria: + - 'A steady-state greedy decode token makes 0 host heap allocations, 0 device allocations and downloads exactly the 4-byte token id (counted, not asserted by reading).' + - 'Greedy on the device is token-identical to the CPU argmax of the downloaded logits at temperature 0.' + phases: [] + subtasks: [] + estimated_effort: null + labels: + - kind:code + notes: 'Scope is #4215 (a)(b)(c) from /mnt/nvme-raid0/tmp/alt0693/REPORT.md §A. Measured +2% at 850 on 4B (78.3 -> 79.9 tok/s); the 0.58x gap to ollama is NOT closed by this slice.' diff --git a/docs/roadmaps/roadmap.yaml b/docs/roadmaps/roadmap.yaml index b187fff708..33f075867a 100644 --- a/docs/roadmaps/roadmap.yaml +++ b/docs/roadmaps/roadmap.yaml @@ -20628,3 +20628,22 @@ roadmap: labels: - kind:code notes: 'Ported from c8ce988c3 (perf/4215, based on 0.69.3) onto main; prerequisite for #4215(b), greedy GPU argmax.' +- id: PMAT-4215 + github_issue: 4215 + item_type: task + title: 'Qwen3.5 CUDA greedy decode: persistent embedding buffer, GPU argmax, no format!/Vec on the launch path' + status: inprogress + priority: high + assigned_to: null + created: 2026-09-24T08:28:19Z + updated: 2026-09-24T20:45:00Z + spec: null + acceptance_criteria: + - 'A steady-state greedy decode token makes 0 host heap allocations, 0 device allocations and downloads exactly the 4-byte token id (counted, not asserted by reading).' + - 'Greedy on the device is token-identical to the CPU argmax of the downloaded logits at temperature 0.' + phases: [] + subtasks: [] + estimated_effort: null + labels: + - kind:code + notes: 'Scope is #4215 (a)(b)(c) from /mnt/nvme-raid0/tmp/alt0693/REPORT.md §A. Measured +2% at 850 on 4B (78.3 -> 79.9 tok/s); the 0.58x gap to ollama is NOT closed by this slice.'