117 lines
4.0 KiB
Rust
117 lines
4.0 KiB
Rust
//! CUDA softmax and log-softmax kernels (f32 native).
|
|
//!
|
|
//! Numerically stable implementations operating on row-major `[batch, dim]` data.
|
|
|
|
use std::sync::Arc;
|
|
|
|
use cudarc;
|
|
use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
|
|
|
|
use ml_core::MLError;
|
|
|
|
use super::tensor_util::CudaVec;
|
|
|
|
struct SoftmaxKernels {
|
|
softmax_fwd: cudarc::driver::CudaFunction,
|
|
log_softmax_fwd: cudarc::driver::CudaFunction,
|
|
}
|
|
|
|
fn compile_softmax_kernels(stream: &Arc<CudaStream>) -> Result<SoftmaxKernels, MLError> {
|
|
static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/softmax_kernels.cubin"));
|
|
let module = stream.context()
|
|
.load_cubin(CUBIN.to_vec())
|
|
.map_err(|e| MLError::InitializationError {
|
|
component: "softmax".to_owned(),
|
|
message: format!("Failed to load softmax cubin: {e}"),
|
|
})?;
|
|
Ok(SoftmaxKernels {
|
|
softmax_fwd: module.load_function("softmax_forward").map_err(|e| {
|
|
MLError::InitializationError {
|
|
component: "softmax".to_owned(),
|
|
message: format!("Failed to load softmax_forward: {e}"),
|
|
}
|
|
})?,
|
|
log_softmax_fwd: module.load_function("log_softmax_forward").map_err(|e| {
|
|
MLError::InitializationError {
|
|
component: "softmax".to_owned(),
|
|
message: format!("Failed to load log_softmax_forward: {e}"),
|
|
}
|
|
})?,
|
|
})
|
|
}
|
|
|
|
/// Compute softmax over the last dimension: `[batch_size, dim]`
|
|
pub fn cuda_softmax(
|
|
stream: &Arc<CudaStream>,
|
|
input: &CudaSlice<f32>,
|
|
batch_size: usize,
|
|
dim: usize,
|
|
) -> Result<CudaVec, MLError> {
|
|
let kernels = compile_softmax_kernels(stream)?;
|
|
let total = batch_size * dim;
|
|
let mut output = stream.alloc_zeros::<f32>(total).map_err(|e| {
|
|
MLError::ModelError(format!("Failed to alloc softmax output: {e}"))
|
|
})?;
|
|
|
|
// One block per row, warp-width threads (32 for dims <= 32, else 32)
|
|
let threads = dim.min(32) as u32;
|
|
let batch_i32 = batch_size as i32;
|
|
let dim_i32 = dim as i32;
|
|
|
|
// SAFETY: softmax_forward reads batch_size*dim f32 input elements and
|
|
// writes batch_size*dim output elements. One block processes each row.
|
|
// Both buffers are pre-allocated with `total` = batch_size*dim elements.
|
|
unsafe {
|
|
stream.launch_builder(&kernels.softmax_fwd)
|
|
.arg(&mut output)
|
|
.arg(input)
|
|
.arg(&batch_i32)
|
|
.arg(&dim_i32)
|
|
.launch(LaunchConfig {
|
|
grid_dim: (batch_size as u32, 1, 1),
|
|
block_dim: (threads, 1, 1),
|
|
shared_mem_bytes: 0,
|
|
})
|
|
.map_err(|e| MLError::ModelError(format!("softmax_forward launch failed: {e}")))?;
|
|
}
|
|
|
|
Ok(CudaVec::new(output, total))
|
|
}
|
|
|
|
/// Compute log-softmax over the last dimension: `[batch_size, dim]`
|
|
pub fn cuda_log_softmax(
|
|
stream: &Arc<CudaStream>,
|
|
input: &CudaSlice<f32>,
|
|
batch_size: usize,
|
|
dim: usize,
|
|
) -> Result<CudaVec, MLError> {
|
|
let kernels = compile_softmax_kernels(stream)?;
|
|
let total = batch_size * dim;
|
|
let mut output = stream.alloc_zeros::<f32>(total).map_err(|e| {
|
|
MLError::ModelError(format!("Failed to alloc log_softmax output: {e}"))
|
|
})?;
|
|
|
|
let threads = dim.min(32) as u32;
|
|
let batch_i32 = batch_size as i32;
|
|
let dim_i32 = dim as i32;
|
|
|
|
// SAFETY: log_softmax_forward reads batch_size*dim f32 input elements and
|
|
// writes batch_size*dim output elements. One block per row, same invariants
|
|
// as softmax_forward above.
|
|
unsafe {
|
|
stream.launch_builder(&kernels.log_softmax_fwd)
|
|
.arg(&mut output)
|
|
.arg(input)
|
|
.arg(&batch_i32)
|
|
.arg(&dim_i32)
|
|
.launch(LaunchConfig {
|
|
grid_dim: (batch_size as u32, 1, 1),
|
|
block_dim: (threads, 1, 1),
|
|
shared_mem_bytes: 0,
|
|
})
|
|
.map_err(|e| MLError::ModelError(format!("log_softmax_forward launch failed: {e}")))?;
|
|
}
|
|
|
|
Ok(CudaVec::new(output, total))
|
|
}
|