From 89ec3a7ff6d6ad2ad0112ba1c9034779bfb8baf1 Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Thu, 13 Aug 2026 21:54:59 -0700 Subject: [PATCH] feat(native-decode): make the device argmax / greedy token selection batch-aware (2b-impl-3, #750) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Select `batch` greedy tokens from a `[batch, 1, vocab]` logits tensor instead of one token from `[1, 1, vocab]`, threading a `batch` count through the whole device-argmax stack: the CUDA kernel launch, the `ExecutionProvider` trait, the `onnx-runtime-ep-cuda` provider, `DeviceIoBinding::device_argmax`, the result buffer allocation, and `read_greedy_result`. `batch` stays constructor-fixed to 1; this delivers the capability, not batch-N on (that is 2b-impl-4). Batch-1 is byte-identical by construction: the grid stays 1-deep (gridDim.y == 1 / gridDim.x == 1), the result header stays 2 words, and the scratch layout is unchanged, so the launch geometry and the 8-byte read-back are bit-for-bit the prior single-sequence path. The reduction operator is additionally order-invariant (max value / min index on ties), so tie-breaks cannot drift with the grid shape. New device test `device_argmax_batch_selects_per_row_max_with_deliberate_ties` puts the max at a different position per row, forces multi-partial reductions, and plants ties at an earlier index than the peak, asserting each batched row equals a standalone single-sequence argmax — the control that a shared reduction buffer or a wrong per-sequence stride would fail. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624 --- .../src/native_decode/cuda.rs | 55 ++++-- crates/onnx-runtime-ep-api/src/provider.rs | 10 +- .../src/kernels/device_argmax.rs | 182 +++++++++++++++--- crates/onnx-runtime-ep-cuda/src/lib.rs | 8 +- crates/onnx-runtime-ep-cuda/src/provider.rs | 3 +- crates/onnx-runtime-session/src/tensor.rs | 8 +- 6 files changed, 220 insertions(+), 46 deletions(-) diff --git a/crates/onnx-genai-engine/src/native_decode/cuda.rs b/crates/onnx-genai-engine/src/native_decode/cuda.rs index f5cfe78df5..63690eca79 100644 --- a/crates/onnx-genai-engine/src/native_decode/cuda.rs +++ b/crates/onnx-genai-engine/src/native_decode/cuda.rs @@ -2661,15 +2661,19 @@ impl DecodeCudaState { let vocab = *logits_shape .last() .context("CUDA logits shape has no vocabulary dimension")?; - 2 + onnx_runtime_ep_cuda::device_argmax_scratch_words(vocab) + // The device-argmax result buffer holds a `2 × batch` header (a token + // id + capture-error pair per sequence) plus per-sequence scratch + // (stage 2b-impl-3, #750). At `batch == 1` this is byte-identical to + // the previous `2 + scratch_words(vocab)` allocation. + 2 * batch + onnx_runtime_ep_cuda::device_argmax_scratch_words(vocab, batch) }; #[cfg(not(feature = "cuda"))] - let argmax_words = 2; + let argmax_words = 2 * batch; let greedy_result = session.allocate_device_output_binding( "__native_greedy_argmax", DataType::Uint32, vec![argmax_words], - vec![2], + vec![2 * batch], )?; // A graph records launch geometry, so replay is unsafe when a persistent @@ -3312,26 +3316,49 @@ impl DecodeCudaState { self.bindings[self.logits_binding].device_argmax_supported() } - fn read_greedy_result(&mut self) -> anyhow::Result<(TokenId, u32)> { + /// Run the batched device argmax over the `[batch, 1, vocab]` logits binding + /// and read back the `batch` selected token ids paired with the shared + /// device capture-error word (stage 2b-impl-3, #750). Row `i` of the returned + /// vector is the greedy token of sequence `i`; every row carries the same + /// latched capture-error bitmask (the kernel copies the one shared word into + /// each slot). At `batch == 1` the argmax launch geometry and the 8-byte + /// read-back are byte-identical to the previous single-sequence path. + fn read_greedy_result_batch(&mut self) -> anyhow::Result> { let vocab = *self .logits_shape .last() .context("CUDA logits shape has no vocabulary dimension")?; - self.bindings[self.logits_binding].device_argmax(vocab, &mut self.greedy_result)?; - let mut bytes = [0_u8; 2 * std::mem::size_of::()]; + self.bindings[self.logits_binding].device_argmax( + vocab, + self.batch, + &mut self.greedy_result, + )?; + let word = std::mem::size_of::(); + let mut bytes = vec![0_u8; self.batch * 2 * word]; self.greedy_result.read_bytes_into(&mut bytes)?; - Ok(( - u32::from_ne_bytes( - bytes[..4] + let mut rows = Vec::with_capacity(self.batch); + for sequence in 0..self.batch { + let base = sequence * 2 * word; + let token = u32::from_ne_bytes( + bytes[base..base + word] .try_into() .map_err(|_| anyhow::anyhow!("four token-id bytes"))?, - ), - u32::from_ne_bytes( - bytes[4..] + ); + let capture_error = u32::from_ne_bytes( + bytes[base + word..base + 2 * word] .try_into() .map_err(|_| anyhow::anyhow!("four capture-error bytes"))?, - ), - )) + ); + rows.push((token, capture_error)); + } + Ok(rows) + } + + fn read_greedy_result(&mut self) -> anyhow::Result<(TokenId, u32)> { + let rows = self.read_greedy_result_batch()?; + rows.into_iter() + .next() + .context("device argmax returned no rows") } pub(crate) fn invalidate_graph( diff --git a/crates/onnx-runtime-ep-api/src/provider.rs b/crates/onnx-runtime-ep-api/src/provider.rs index aff944a65f..d64b9b45a0 100644 --- a/crates/onnx-runtime-ep-api/src/provider.rs +++ b/crates/onnx-runtime-ep-api/src/provider.rs @@ -645,13 +645,17 @@ pub trait ExecutionProvider: Send + Sync { false } - /// Launch an allocation-free device argmax over `elements` contiguous - /// `dtype` values (Float32 or Float16). `result` receives two native-endian - /// u32 values: token id, then the latching device capture-error bitmask. + /// Launch an allocation-free device argmax over `batch` sequences of + /// `elements` contiguous `dtype` values (Float32 or Float16) each, laid out + /// as a `[batch, elements]` row-major block. `result` receives, per sequence + /// `s`, two native-endian u32 values at word offset `2*s`: the token id, then + /// the latching device capture-error bitmask. At `batch == 1` this is the + /// previous single-sequence contract byte-for-byte. fn device_argmax( &self, _logits: &DeviceBuffer, _elements: usize, + _batch: usize, _dtype: DataType, _result: &mut DeviceBuffer, ) -> Result<()> { diff --git a/crates/onnx-runtime-ep-cuda/src/kernels/device_argmax.rs b/crates/onnx-runtime-ep-cuda/src/kernels/device_argmax.rs index 94bba9cd7b..92c28f3da9 100644 --- a/crates/onnx-runtime-ep-cuda/src/kernels/device_argmax.rs +++ b/crates/onnx-runtime-ep-cuda/src/kernels/device_argmax.rs @@ -18,8 +18,16 @@ pub(crate) fn partial_count(elements: usize) -> usize { .clamp(1, MAX_PARTIALS) } -pub(crate) fn scratch_words(elements: usize) -> usize { - 2 * partial_count(elements) +/// Scratch u32 words the device argmax result buffer needs *beyond* its +/// `2 × batch` header words, for `batch` sequences of `elements` logits each. +/// +/// The scratch holds, per sequence, `partial_count(elements)` partial values +/// (f32, one u32 word each) followed by `partial_count(elements)` partial +/// indices (u32), i.e. `2 × partial_count` words per sequence. At `batch == 1` +/// this is byte-identical to the previous single-sequence `2 × partial_count` +/// scratch (stage 2b-impl-3, #750). +pub(crate) fn scratch_words(elements: usize, batch: usize) -> usize { + 2 * partial_count(elements) * batch } const SOURCE: &str = r#" @@ -67,6 +75,16 @@ __device__ __forceinline__ void greedy_argmax_partials_impl( unsigned long long elements, float* partial_values, unsigned int* partial_indices) { + // One sequence per grid row (blockIdx.y); `gridDim.x` blocks cooperatively + // reduce that sequence's `elements` contiguous logits. Sequence `s` reads + // `logits[s*elements + ..]` and writes its `gridDim.x` partials at + // `partial_{values,indices}[s*gridDim.x + ..]`. At batch 1 (gridDim.y == 1) + // the sequence offset is 0 and the launch is byte-identical to the previous + // single-sequence kernel (stage 2b-impl-3, #750). + unsigned long long sequence = blockIdx.y; + const T* seq_logits = logits + sequence * elements; + float* seq_values = partial_values + sequence * gridDim.x; + unsigned int* seq_indices = partial_indices + sequence * gridDim.x; float best = -1.0f / 0.0f; unsigned int best_index = 0; unsigned long long i = @@ -74,7 +92,7 @@ __device__ __forceinline__ void greedy_argmax_partials_impl( unsigned long long stride = static_cast(blockDim.x) * gridDim.x; for (; i < elements; i += stride) { - float value = argmax_load(logits[i]); + float value = argmax_load(seq_logits[i]); if (isnan(value)) continue; unsigned int index = static_cast(i); argmax_update(value, index, best, best_index); @@ -97,8 +115,8 @@ __device__ __forceinline__ void greedy_argmax_partials_impl( best_index = lane < warp_count ? warp_indices[lane] : 0; warp_argmax(best, best_index); if (lane == 0) { - partial_values[blockIdx.x] = best; - partial_indices[blockIdx.x] = best_index; + seq_values[blockIdx.x] = best; + seq_indices[blockIdx.x] = best_index; } } } @@ -127,10 +145,19 @@ extern "C" __global__ void greedy_argmax_finalize( unsigned int partial_count, const unsigned int* capture_error, unsigned int* result) { + // One block per sequence (blockIdx.x). Sequence `s` reduces its own + // `partial_count` partials at `partial_{values,indices}[s*partial_count + ..]` + // and writes its token id / capture-error pair at `result[2*s .. 2*s+2]`. At + // batch 1 (gridDim.x == 1) this writes result[0]/result[1] from the base + // partial region, byte-identical to the previous single-sequence finalize + // (stage 2b-impl-3, #750). + unsigned int sequence = blockIdx.x; + const float* seq_values = partial_values + sequence * partial_count; + const unsigned int* seq_indices = partial_indices + sequence * partial_count; float best = -1.0f / 0.0f; unsigned int best_index = 0; for (unsigned int i = threadIdx.x; i < partial_count; i += blockDim.x) { - argmax_update(partial_values[i], partial_indices[i], best, best_index); + argmax_update(seq_values[i], seq_indices[i], best, best_index); } warp_argmax(best, best_index); __shared__ float warp_values[32]; @@ -148,8 +175,8 @@ extern "C" __global__ void greedy_argmax_finalize( best_index = lane < warp_count ? warp_indices[lane] : 0; warp_argmax(best, best_index); if (lane == 0) { - result[0] = best_index; - result[1] = *capture_error; + result[2 * sequence] = best_index; + result[2 * sequence + 1] = *capture_error; } } } @@ -159,6 +186,7 @@ pub(crate) fn launch( runtime: &CudaRuntime, logits: &DeviceBuffer, elements: usize, + batch: usize, dtype: DataType, result: &mut DeviceBuffer, ) -> Result<()> { @@ -167,11 +195,25 @@ pub(crate) fn launch( "cuda_ep device argmax: logits must not be empty".into(), )); } + if batch == 0 { + return Err(EpError::KernelFailed( + "cuda_ep device argmax: batch must not be zero".into(), + )); + } if elements > u32::MAX as usize { return Err(EpError::KernelFailed(format!( "cuda_ep device argmax: {elements} elements exceed the u32 token-id range" ))); } + // The grid dedicates one row (blockIdx.y for partials, blockIdx.x for + // finalize) to each of the `batch` sequences, so the batch count must fit + // the u32 grid dimension. At batch 1 the grid is 1-deep and byte-identical + // to the previous single-sequence launch (stage 2b-impl-3, #750). + if batch > u32::MAX as usize { + return Err(EpError::KernelFailed(format!( + "cuda_ep device argmax: batch {batch} exceeds the u32 grid range" + ))); + } let (entry, elem_size) = match dtype { DataType::Float32 => ("greedy_argmax_partials_f32", std::mem::size_of::()), DataType::Float16 => ("greedy_argmax_partials_f16", std::mem::size_of::()), @@ -181,18 +223,28 @@ pub(crate) fn launch( ))); } }; - let logits_bytes = elements.checked_mul(elem_size).ok_or_else(|| { - EpError::KernelFailed("cuda_ep device argmax: logits byte size overflows".into()) - })?; + let logits_bytes = elements + .checked_mul(elem_size) + .and_then(|per_seq| per_seq.checked_mul(batch)) + .ok_or_else(|| { + EpError::KernelFailed("cuda_ep device argmax: logits byte size overflows".into()) + })?; if logits_bytes > logits.len() { return Err(EpError::KernelFailed(format!( - "cuda_ep device argmax: {elements} values require {logits_bytes} bytes, buffer has {}", + "cuda_ep device argmax: {batch}×{elements} values require {logits_bytes} bytes, buffer has {}", logits.len() ))); } let partial_count = partial_count(elements); - let required_result_bytes = RESULT_BYTES - + scratch_words(elements) + // Header is `2 × batch` u32 words (a token id + capture-error pair per + // sequence, laid out contiguously so sequence `s` owns `result[2*s..2*s+2]`), + // followed by `batch × 2 × partial_count` scratch words. At batch 1 this is + // `RESULT_BYTES` header + the previous single-sequence scratch — byte-identical. + let header_bytes = batch.checked_mul(RESULT_BYTES).ok_or_else(|| { + EpError::KernelFailed("cuda_ep device argmax: result header size overflows".into()) + })?; + let required_result_bytes = header_bytes + + scratch_words(elements, batch) .checked_mul(std::mem::size_of::()) .ok_or_else(|| { EpError::KernelFailed("cuda_ep device argmax: scratch byte size overflows".into()) @@ -217,10 +269,12 @@ pub(crate) fn launch( runtime.nvrtc_function("native_device_argmax", SOURCE, "greedy_argmax_finalize")?; let logits_ptr = cuptr(logits.as_ptr()); let elements = elements as u64; - let scratch_ptr = unsafe { result.as_mut_ptr().add(RESULT_BYTES) }; + // Scratch begins after the `batch`-wide header. Per-sequence partials are + // laid out values-then-indices, `batch × partial_count` of each. + let scratch_ptr = unsafe { result.as_mut_ptr().add(header_bytes) }; let partial_values_ptr = cuptr(scratch_ptr); let partial_indices_ptr = - cuptr(unsafe { scratch_ptr.add(partial_count * std::mem::size_of::()) }); + cuptr(unsafe { scratch_ptr.add(batch * partial_count * std::mem::size_of::()) }); let capture_error_ptr = runtime.capture_error_ptr(); let result_ptr = cuptr(result.as_mut_ptr()); let mut builder = runtime.stream().launch_builder(&partial_function); @@ -231,7 +285,7 @@ pub(crate) fn launch( .arg(&partial_indices_ptr); unsafe { builder.launch(LaunchConfig { - grid_dim: (partial_count as u32, 1, 1), + grid_dim: (partial_count as u32, batch as u32, 1), block_dim: (BLOCK, 1, 1), shared_mem_bytes: 0, }) @@ -248,7 +302,7 @@ pub(crate) fn launch( .arg(&result_ptr); unsafe { builder.launch(LaunchConfig { - grid_dim: (1, 1, 1), + grid_dim: (batch as u32, 1, 1), block_dim: (BLOCK, 1, 1), shared_mem_bytes: 0, }) @@ -281,8 +335,8 @@ mod tests { .unwrap_or(0) as u32 } - fn result_bytes(elements: usize) -> usize { - RESULT_BYTES + scratch_words(elements) * std::mem::size_of::() + fn result_bytes(elements: usize, batch: usize) -> usize { + (batch * RESULT_BYTES) + scratch_words(elements, batch) * std::mem::size_of::() } fn run_case(ep: &CudaExecutionProvider, logits: &[f32]) -> [u32; 2] { @@ -291,9 +345,9 @@ mod tests { .flat_map(|value| value.to_ne_bytes()) .collect::>(); let mut input = ep.allocate(bytes.len(), 256).unwrap(); - let mut output = ep.allocate(result_bytes(logits.len()), 256).unwrap(); + let mut output = ep.allocate(result_bytes(logits.len(), 1), 256).unwrap(); ep.copy_from_host(&bytes, &mut input).unwrap(); - ep.device_argmax(&input, logits.len(), DataType::Float32, &mut output) + ep.device_argmax(&input, logits.len(), 1, DataType::Float32, &mut output) .unwrap(); let mut result = [0_u8; RESULT_BYTES]; ep.copy_to_host(&output, &mut result).unwrap(); @@ -306,15 +360,48 @@ mod tests { values } + /// Run the argmax over `batch` sequences of `vocab` f32 logits laid out + /// row-major (`rows.concat()`), returning the `(token_id, capture_error)` + /// pair per sequence read back from the `2 × batch`-word header. + fn run_batch_case(ep: &CudaExecutionProvider, rows: &[Vec]) -> Vec<[u32; 2]> { + let batch = rows.len(); + let vocab = rows[0].len(); + assert!(rows.iter().all(|row| row.len() == vocab)); + let flat = rows.iter().flatten().copied().collect::>(); + let bytes = flat + .iter() + .flat_map(|value| value.to_ne_bytes()) + .collect::>(); + let mut input = ep.allocate(bytes.len(), 256).unwrap(); + let mut output = ep.allocate(result_bytes(vocab, batch), 256).unwrap(); + ep.copy_from_host(&bytes, &mut input).unwrap(); + ep.device_argmax(&input, vocab, batch, DataType::Float32, &mut output) + .unwrap(); + let mut header = vec![0_u8; batch * RESULT_BYTES]; + ep.copy_to_host(&output, &mut header).unwrap(); + let out = (0..batch) + .map(|s| { + let base = s * RESULT_BYTES; + [ + u32::from_ne_bytes(header[base..base + 4].try_into().unwrap()), + u32::from_ne_bytes(header[base + 4..base + 8].try_into().unwrap()), + ] + }) + .collect(); + ep.deallocate(input).unwrap(); + ep.deallocate(output).unwrap(); + out + } + fn run_case_f16(ep: &CudaExecutionProvider, logits: &[f32]) -> [u32; 2] { let bytes = logits .iter() .flat_map(|&value| half::f16::from_f32(value).to_bits().to_ne_bytes()) .collect::>(); let mut input = ep.allocate(bytes.len(), 256).unwrap(); - let mut output = ep.allocate(result_bytes(logits.len()), 256).unwrap(); + let mut output = ep.allocate(result_bytes(logits.len(), 1), 256).unwrap(); ep.copy_from_host(&bytes, &mut input).unwrap(); - ep.device_argmax(&input, logits.len(), DataType::Float16, &mut output) + ep.device_argmax(&input, logits.len(), 1, DataType::Float16, &mut output) .unwrap(); let mut result = [0_u8; RESULT_BYTES]; ep.copy_to_host(&output, &mut result).unwrap(); @@ -382,4 +469,51 @@ mod tests { let result = run_case_f16(&ep, &all_non_finite); assert_eq!(result, [host_argmax(&all_non_finite), 0]); } + + #[test] + fn device_argmax_batch_selects_per_row_max_with_deliberate_ties() { + let Some(ep) = gpu() else { return }; + // A control that would falsify a shared reduction buffer or a wrong + // per-sequence stride: the maximum sits at a *different* position in + // every row, each row is larger than one block's grid-stride reach so + // multiple partials are reduced, and two rows carry a deliberate tie so + // the lowest-index tie-break is exercised per sequence rather than + // globally (stage 2b-impl-3, #750). + let vocab = 4096; + let peaks = [17_usize, 4000, 128, 2049]; + let mut rows: Vec> = peaks + .iter() + .enumerate() + .map(|(row, &peak)| { + let mut logits = (0..vocab) + .map(|i| ((i + row) % 13) as f32 * 0.1) + .collect::>(); + logits[peak] = 9.0; + logits + }) + .collect(); + // Row 1 gets a tie at an *earlier* index than its peak; the lower index + // must win, proving each sequence tie-breaks against its own partials. + rows[1][2500] = 9.0; // peak was 4000 → lower index 2500 must win + rows[3][900] = 9.0; // peak was 2049 → lower index 900 must win + let expected = [17_u32, 2500, 128, 900]; + + let batched = run_batch_case(&ep, &rows); + assert_eq!(batched.len(), rows.len()); + for (row, (result, &want)) in batched.iter().zip(expected.iter()).enumerate() { + assert_eq!( + result[0], want, + "batched argmax row {row} selected token {} not {want}", + result[0] + ); + assert_eq!(result[1], 0, "row {row} unexpected capture-error flag"); + // Each batched row must equal a standalone single-sequence argmax of + // the same logits — the byte-identity contract, row by row. + let solo = run_case(&ep, &rows[row]); + assert_eq!( + *result, solo, + "batched row {row} diverged from the single-sequence argmax" + ); + } + } } diff --git a/crates/onnx-runtime-ep-cuda/src/lib.rs b/crates/onnx-runtime-ep-cuda/src/lib.rs index b64c9d2859..9f1f7c521f 100644 --- a/crates/onnx-runtime-ep-cuda/src/lib.rs +++ b/crates/onnx-runtime-ep-cuda/src/lib.rs @@ -127,8 +127,10 @@ pub use weight_paging::{ global_offload_stats, reset_global_offload_stats, zero_copy_hybrid_from_env, }; -/// Number of additional u32 words required by the CUDA device argmax result buffer. -pub fn device_argmax_scratch_words(elements: usize) -> usize { - kernels::device_argmax::scratch_words(elements) +/// Number of additional u32 words the CUDA device argmax result buffer needs +/// beyond its `2 × batch` header words, for `batch` sequences of `elements` +/// logits each. +pub fn device_argmax_scratch_words(elements: usize, batch: usize) -> usize { + kernels::device_argmax::scratch_words(elements, batch) } pub use runtime::{CudaAllocationCounts, CudaRuntime}; diff --git a/crates/onnx-runtime-ep-cuda/src/provider.rs b/crates/onnx-runtime-ep-cuda/src/provider.rs index 720aac3983..6085472ca2 100644 --- a/crates/onnx-runtime-ep-cuda/src/provider.rs +++ b/crates/onnx-runtime-ep-cuda/src/provider.rs @@ -1207,10 +1207,11 @@ impl ExecutionProvider for CudaExecutionProvider { &self, logits: &DeviceBuffer, elements: usize, + batch: usize, dtype: DataType, result: &mut DeviceBuffer, ) -> Result<()> { - crate::kernels::device_argmax::launch(&self.runtime, logits, elements, dtype, result) + crate::kernels::device_argmax::launch(&self.runtime, logits, elements, batch, dtype, result) } fn copy_from_host(&self, src: &[u8], dst: &mut DeviceBuffer) -> Result<()> { diff --git a/crates/onnx-runtime-session/src/tensor.rs b/crates/onnx-runtime-session/src/tensor.rs index 146f0f7550..1ab782e1fc 100644 --- a/crates/onnx-runtime-session/src/tensor.rs +++ b/crates/onnx-runtime-session/src/tensor.rs @@ -730,7 +730,12 @@ impl DeviceIoBinding { self.allocator.device_argmax_supported() } - pub fn device_argmax(&self, elements: usize, result: &mut DeviceIoBinding) -> Result<()> { + pub fn device_argmax( + &self, + elements: usize, + batch: usize, + result: &mut DeviceIoBinding, + ) -> Result<()> { if !matches!(self.dtype, DataType::Float32 | DataType::Float16) || result.dtype != DataType::Uint32 { @@ -747,6 +752,7 @@ impl DeviceIoBinding { Ok(self.allocator.device_argmax( self.buffer(), elements, + batch, self.dtype, result.buffer_mut(), )?)