73 lines
2.4 KiB
Rust
73 lines
2.4 KiB
Rust
use std::ffi::c_void;
|
|
use xserv_tensor::{DType, Device, Tensor};
|
|
|
|
unsafe extern "C" {
|
|
fn launch_argmax_bf16(
|
|
logits: *const c_void,
|
|
out_idx: *mut c_void,
|
|
rows: i32,
|
|
cols: i32,
|
|
stream: *mut c_void,
|
|
);
|
|
}
|
|
|
|
/// GPU argmax over the last dim of a [rows, cols] BF16 tensor.
|
|
///
|
|
/// Returns a host `Vec<u32>` of length `rows`. Internally:
|
|
/// - launches one kernel that writes [rows] i32 indices on device
|
|
/// - D2H copies just `rows * 4` bytes (vs `rows * cols * 2` for the
|
|
/// "copy logits to CPU then argmax" path it replaces)
|
|
///
|
|
/// This is the greedy-decode hot path: avoids touching the full
|
|
/// [B, vocab] logits buffer on the host every step.
|
|
pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
|
|
assert_eq!(logits.ndim(), 2, "argmax expects a 2D [rows, cols] tensor");
|
|
assert_eq!(logits.dtype(), DType::BF16, "argmax kernel is BF16-only");
|
|
assert!(logits.is_contiguous(), "argmax requires contiguous input");
|
|
assert!(
|
|
matches!(logits.device(), Device::Cuda(_)),
|
|
"argmax requires GPU input"
|
|
);
|
|
|
|
let rows = logits.shape()[0];
|
|
let cols = logits.shape()[1];
|
|
assert!(rows <= i32::MAX as usize);
|
|
assert!(cols <= i32::MAX as usize);
|
|
|
|
// Output buffer: rows * i32. Pooled allocator so this is essentially free
|
|
// after the first call.
|
|
let bytes = rows * std::mem::size_of::<i32>();
|
|
let mut out = xserv_cuda::allocator::cached_alloc(bytes).expect("argmax out alloc");
|
|
|
|
unsafe {
|
|
launch_argmax_bf16(
|
|
logits.data_ptr() as *const c_void,
|
|
out.as_mut_ptr() as *mut c_void,
|
|
rows as i32,
|
|
cols as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
|
|
let mut host_bytes = vec![0u8; bytes];
|
|
out.copy_to_host(&mut host_bytes).expect("argmax D2H");
|
|
drop(out); // returned to pool
|
|
|
|
let host_i32: &[i32] =
|
|
unsafe { std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows) };
|
|
host_i32.iter().map(|&v| v as u32).collect()
|
|
}
|
|
|
|
/// Convenience: argmax of a single row [1, cols] (or [cols] reshaped to [1, cols]).
|
|
pub fn argmax_bf16_single(logits: &Tensor) -> u32 {
|
|
let cols = *logits.shape().last().unwrap();
|
|
let rows = logits.numel() / cols;
|
|
assert_eq!(rows, 1, "argmax_bf16_single requires a single row");
|
|
let view = if logits.ndim() == 2 {
|
|
logits.clone()
|
|
} else {
|
|
logits.reshape(&[1, cols])
|
|
};
|
|
argmax_bf16_to_host(&view)[0]
|
|
}
|