gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(bf16): ml-core compiles clean — nvrtc removed, BF16 boundaries fixed

Browse Source 
- Deleted cuda_compile.rs (nvrtc runtime compilation)
- Removed nvrtc feature from cudarc dependency
- Added f16 feature + half crate
- Stubbed nvrtc-dependent constructors (ReductionKernels, GpuAdamW, etc.)
- Fixed all BF16 boundary conversions (f32 host ↔ bf16 GPU)
- GpuTensor.from_host/to_host now convert at boundary

ml-core: 0 errors. ml: 49 errors remaining (Phase 3 Task 14).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-28 01:13:51 +01:00
parent 6189ed7b70
commit 1d9aa6e32d
14 changed files with 167 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
crates/ml-core/src/cuda_autograd/activations.rs
View File

@@ -9,7 +9,6 @@
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use crate::MLError;
use super::gpu_tensor::GpuTensor;
@@ -43,32 +42,13 @@ fn elem_cfg(n: usize) -> LaunchConfig {
impl ActivationKernels {
/// Compile all activation kernels from a single CUDA source.
pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(ACTIVATION_CUDA_SRC, &context)
.map_err(|e| MLError::ModelError(format!("activation kernel compilation: {e}")))?;
let module = context.load_module(ptx).map_err(|e| {
MLError::ModelError(format!("activation module load: {e}"))
})?;
let load = |name: &str| -> Result<CudaFunction, MLError> {
module.load_function(name).map_err(|e| {
MLError::ModelError(format!("activation kernel '{name}' load: {e}"))
})
};
Ok(Self {
relu_forward: load("relu_forward")?,
relu_backward: load("relu_backward")?,
leaky_relu_forward: load("leaky_relu_forward")?,
leaky_relu_backward: load("leaky_relu_backward")?,
gelu_forward: load("gelu_forward")?,
gelu_backward: load("gelu_backward")?,
sigmoid_forward: load("sigmoid_forward")?,
sigmoid_backward: load("sigmoid_backward")?,
tanh_forward: load("tanh_forward")?,
tanh_backward: load("tanh_backward")?,
})
///
/// STUB: nvrtc runtime compilation removed. These kernels will be provided
/// by precompiled cubins from build.rs in a future step.
pub fn new(_stream: &Arc<CudaStream>) -> Result<Self, MLError> {
Err(MLError::ModelError(
"activation kernels not available: nvrtc removed, cubins not yet wired".to_owned(),
))
}
// ── ReLU ──────────────────────────────────────────────────────────

73
crates/ml-core/src/cuda_autograd/dropout.rs
View File

@@ -9,7 +9,6 @@
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use crate::MLError;
use super::gpu_tensor::GpuTensor;
@@ -139,70 +138,24 @@ impl GpuDropout {
) -> Result<GpuTensor, MLError> {
match mask {
None => clone_gpu_tensor(dy, stream),
Some(m) => {
// dx = dy * mask (mask already includes the 1/(1-p) scaling)
// We need an element-wise multiply kernel. Reuse a simple approach:
// allocate output, launch multiply kernel.
let n = dy.numel();
let dx = GpuTensor::zeros(&dy.shape, stream)?;
// Compile a simple multiply kernel
let src = r#"
extern "C" __global__
void elemwise_mul(const float* __restrict__ a,
const float* __restrict__ b,
float* __restrict__ c,
int n) {
int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) {
c[i] = a[i] * b[i];
}
}
"#;
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(src, &context)
.map_err(|e| MLError::ModelError(format!("dropout backward kernel: {e}")))?;
let module = context.load_module(ptx)
.map_err(|e| MLError::ModelError(format!("dropout backward module: {e}")))?;
let mul_fn = module.load_function("elemwise_mul")
.map_err(|e| MLError::ModelError(format!("dropout backward load: {e}")))?;
let n_i32 = n as i32;
let threads = 256_u32;
let blocks = ((n as u32) + threads - 1) / threads;
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
stream
.launch_builder(&mul_fn)
.arg(&dy.data)
.arg(&m.data)
.arg(&dx.data)
.arg(&n_i32)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("dropout backward: {e}")))?;
}
Ok(dx)
Some(_m) => {
// STUB: dropout backward needs element-wise mul kernel.
// nvrtc removed, cubins not yet wired. Use elementwise kernels instead.
let kernels = super::elementwise::get_or_compile(stream)?;
let out = kernels.binary(&dy.data, &_m.data, dy.numel(), 2)?; // 2 = mul
GpuTensor::new(out, dy.shape.clone())
}
}
}
/// Compile the dropout forward kernel.
fn compile_kernel(&mut self, stream: &Arc<CudaStream>) -> Result<(), MLError> {
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(DROPOUT_CUDA_SRC, &context)
.map_err(|e| MLError::ModelError(format!("dropout kernel compilation: {e}")))?;
let module = context.load_module(ptx)
.map_err(|e| MLError::ModelError(format!("dropout module load: {e}")))?;
let kernel = module.load_function("dropout_forward")
.map_err(|e| MLError::ModelError(format!("dropout_forward load: {e}")))?;
self.kernel = Some(kernel);
Ok(())
///
/// STUB: nvrtc runtime compilation removed. These kernels will be provided
/// by precompiled cubins from build.rs in a future step.
fn compile_kernel(&mut self, _stream: &Arc<CudaStream>) -> Result<(), MLError> {
Err(MLError::ModelError(
"dropout kernel not available: nvrtc removed, cubins not yet wired".to_owned(),
))
}
}

35
crates/ml-core/src/cuda_autograd/elementwise.rs
View File

@@ -15,7 +15,6 @@
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use crate::MLError;
@@ -51,33 +50,13 @@ pub struct ElementwiseKernels {
impl ElementwiseKernels {
/// Compile all element-wise kernels from a single CUDA source.
pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(ELEMENTWISE_CUDA_SRC, &context)
.map_err(|e| MLError::ModelError(format!("elementwise kernel compilation: {e}")))?;
let module = context.load_module(ptx).map_err(|e| {
MLError::ModelError(format!("elementwise module load: {e}"))
})?;
let load = |name: &str| -> Result<CudaFunction, MLError> {
module.load_function(name).map_err(|e| {
MLError::ModelError(format!("elementwise kernel '{name}' load: {e}"))
})
};
Ok(Self {
binary_fn: load("elementwise_binary")?,
unary_fn: load("elementwise_unary")?,
broadcast_binary_fn: load("broadcast_scalar_binary")?,
broadcast_row_binary_fn: load("broadcast_row_binary")?,
broadcast_col_binary_fn: load("broadcast_col_binary")?,
expand_fn: load("expand_broadcast")?,
transpose_2d_fn: load("transpose_2d")?,
gather_fn: load("gather_select")?,
gather_rows_fn: load("gather_rows")?,
gather_rows_u32_fn: load("gather_rows_u32")?,
stream: Arc::clone(stream),
})
///
/// STUB: nvrtc runtime compilation removed. These kernels will be provided
/// by precompiled cubins from build.rs in a future step.
pub fn new(_stream: &Arc<CudaStream>) -> Result<Self, MLError> {
Err(MLError::ModelError(
"elementwise kernels not available: nvrtc removed, cubins not yet wired".to_owned(),
))
}
/// Element-wise binary operation on two same-shape buffers.

13
crates/ml-core/src/cuda_autograd/gpu_tensor.rs
View File

@@ -113,6 +113,8 @@ impl GpuTensor {
}
/// Upload host data to a new GPU tensor.
///
/// Accepts `&[f32]` on the host side, converts to `bf16` before upload.
pub fn from_host(
host: &[f32],
shape: Vec<usize>,
@@ -125,23 +127,26 @@ impl GpuTensor {
actual: host.len(),
});
}
let bf16_host: Vec<half::bf16> = host.iter().map(|&x| half::bf16::from_f32(x)).collect();
let mut data = stream.alloc_zeros::<half::bf16>(expected).map_err(|e| {
MLError::ModelError(format!("GpuTensor::from_host alloc: {e}"))
})?;
stream.memcpy_htod(host, &mut data).map_err(|e| {
stream.memcpy_htod(&bf16_host, &mut data).map_err(|e| {
MLError::ModelError(format!("GpuTensor::from_host HtoD: {e}"))
})?;
Ok(Self { data, shape })
}
/// Download tensor contents to host memory.
///
/// Downloads bf16 from GPU, converts to f32 for the caller.
#[cold]
pub fn to_host(&self, stream: &Arc<CudaStream>) -> Result<Vec<f32>, MLError> {
let mut host = vec![0.0_f32; self.numel()];
stream.memcpy_dtoh(&self.data, &mut host).map_err(|e| { // gpu-exit: API boundary readback
let mut host_bf16 = vec![half::bf16::ZERO; self.numel()];
stream.memcpy_dtoh(&self.data, &mut host_bf16).map_err(|e| { // gpu-exit: API boundary readback
MLError::ModelError(format!("GpuTensor::to_host DtoH: {e}"))
})?;
Ok(host)
Ok(host_bf16.iter().map(|x| x.to_f32()).collect())
}
/// Total number of elements.

14
crates/ml-core/src/cuda_autograd/init.rs
View File

@@ -121,16 +121,18 @@ fn generate_uniform(n: usize, lo: f64, hi: f64) -> Vec<f32> {
/// Upload host f32 slice to a new CudaSlice on the given stream.
///
/// Converts f32 to bf16 before upload.
/// Public for use by downstream crates that need to initialize GPU parameters
/// from host data (e.g., RMSNorm weights initialized to ones).
pub fn upload_to_gpu(
host: &[f32],
stream: &Arc<CudaStream>,
) -> Result<CudaSlice<half::bf16>, MLError> {
let bf16_host: Vec<half::bf16> = host.iter().map(|&x| half::bf16::from_f32(x)).collect();
let mut buf = stream.alloc_zeros::<half::bf16>(host.len()).map_err(|e| {
MLError::ModelError(format!("init upload alloc: {e}"))
})?;
stream.memcpy_htod(host, &mut buf).map_err(|e| {
stream.memcpy_htod(&bf16_host, &mut buf).map_err(|e| {
MLError::ModelError(format!("init upload HtoD: {e}"))
})?;
Ok(buf)
@@ -160,8 +162,9 @@ mod tests {
let fan_in = 256;
let fan_out = 256;
let w = xavier_uniform(fan_in, fan_out, &stream).unwrap();
let mut host = vec![0.0_f32; w.len()];
stream.memcpy_dtoh(&w, &mut host).unwrap(); // test-only readback
let mut host_bf16 = vec![half::bf16::ZERO; w.len()];
stream.memcpy_dtoh(&w, &mut host_bf16).unwrap(); // test-only readback
let host: Vec<f32> = host_bf16.iter().map(|x| x.to_f32()).collect();
let mean: f32 = host.iter().sum::<f32>() / host.len() as f32;
let variance: f32 = host.iter().map(|x| (x - mean) * (x - mean)).sum::<f32>()
@@ -193,8 +196,9 @@ mod tests {
fn test_zeros() {
let stream = make_stream();
let z = zeros(100, &stream).unwrap();
let mut host = vec![1.0_f32; 100];
stream.memcpy_dtoh(&z, &mut host).unwrap(); // test-only readback
let mut host_bf16 = vec![half::bf16::ZERO; 100];
stream.memcpy_dtoh(&z, &mut host_bf16).unwrap(); // test-only readback
let host: Vec<f32> = host_bf16.iter().map(|x| x.to_f32()).collect();
assert!(host.iter().all(|&x| x == 0.0));
}
}

18
crates/ml-core/src/cuda_autograd/layer_norm.rs
View File

@@ -8,7 +8,6 @@
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use crate::MLError;
use super::gpu_tensor::GpuTensor;
@@ -139,16 +138,13 @@ impl GpuLayerNorm {
}
/// Compile the LayerNorm kernel.
fn compile_kernel(&mut self, stream: &Arc<CudaStream>) -> Result<(), MLError> {
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(LAYER_NORM_CUDA_SRC, &context)
.map_err(|e| MLError::ModelError(format!("layer_norm kernel compilation: {e}")))?;
let module = context.load_module(ptx)
.map_err(|e| MLError::ModelError(format!("layer_norm module load: {e}")))?;
let kernel = module.load_function("layer_norm_forward")
.map_err(|e| MLError::ModelError(format!("layer_norm_forward load: {e}")))?;
self.kernel = Some(kernel);
Ok(())
///
/// STUB: nvrtc runtime compilation removed. These kernels will be provided
/// by precompiled cubins from build.rs in a future step.
fn compile_kernel(&mut self, _stream: &Arc<CudaStream>) -> Result<(), MLError> {
Err(MLError::ModelError(
"layer_norm kernel not available: nvrtc removed, cubins not yet wired".to_owned(),
))
}
}

75
crates/ml-core/src/cuda_autograd/linear.rs
View File

@@ -15,7 +15,7 @@ use std::sync::Arc;
use cudarc::cublas::CudaBlas;
use cudarc::cublas::sys::cublasOperation_t;
use cudarc::driver::{CudaSlice, CudaStream, CudaFunction, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
// nvrtc removed — kernels compiled via build.rs cubins
use crate::MLError;
use super::gpu_tensor::GpuTensor;
@@ -114,6 +114,8 @@ impl GpuLinear {
}
/// Download weight tensor to host. Checkpoint use only. Do not call in hot path.
///
/// Downloads bf16 from GPU, converts to f32 for the caller.
#[cold]
pub fn weight_to_vec(
&self,
@@ -123,14 +125,16 @@ impl GpuLinear {
let param = store.get(&self.weight_name).ok_or_else(|| {
MLError::ModelError(format!("weight_to_vec: '{}' not found", self.weight_name))
})?;
let mut host = vec![0.0_f32; param.data.len()];
stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| { // gpu-exit: checkpoint weight export
let mut host_bf16 = vec![half::bf16::ZERO; param.data.len()];
stream.memcpy_dtoh(&param.data, &mut host_bf16).map_err(|e| { // gpu-exit: checkpoint weight export
MLError::ModelError(format!("weight_to_vec DtoH: {e}"))
})?;
Ok(host)
Ok(host_bf16.iter().map(|x| x.to_f32()).collect())
}
/// Download bias tensor to host. Checkpoint use only. Do not call in hot path.
///
/// Downloads bf16 from GPU, converts to f32 for the caller.
#[cold]
pub fn bias_to_vec(
&self,
@@ -140,11 +144,11 @@ impl GpuLinear {
let param = store.get(&self.bias_name).ok_or_else(|| {
MLError::ModelError(format!("bias_to_vec: '{}' not found", self.bias_name))
})?;
let mut host = vec![0.0_f32; param.data.len()];
stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| { // gpu-exit: checkpoint weight export
let mut host_bf16 = vec![half::bf16::ZERO; param.data.len()];
stream.memcpy_dtoh(&param.data, &mut host_bf16).map_err(|e| { // gpu-exit: checkpoint weight export
MLError::ModelError(format!("bias_to_vec DtoH: {e}"))
})?;
Ok(host)
Ok(host_bf16.iter().map(|x| x.to_f32()).collect())
}
/// Forward pass: `Y = X @ W^T + b`.
@@ -296,47 +300,15 @@ impl GpuLinear {
// ── Helper CUDA kernels ──────────────────────────────────────────────────
/// Compile bias-add and reduce-sum kernels.
fn get_bias_kernels(stream: &Arc<CudaStream>) -> Result<(CudaFunction, CudaFunction), MLError> {
let src = r#"
extern "C" __global__
void add_bias_kernel(float* __restrict__ y,
const float* __restrict__ bias,
int rows, int cols) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < rows * cols) {
int col = idx % cols;
y[idx] += bias[col];
}
}
extern "C" __global__
void reduce_sum_axis0_kernel(const float* __restrict__ x,
float* __restrict__ out,
int rows, int cols) {
int col = blockIdx.x * blockDim.x + threadIdx.x;
if (col < cols) {
float sum = 0.0f;
for (int r = 0; r < rows; r++) {
sum += x[r * cols + col];
}
out[col] = sum;
}
}
"#;
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(src, &context)
.map_err(|e| MLError::ModelError(format!("linear helper kernel compilation: {e}")))?;
let module = context.load_module(ptx).map_err(|e| {
MLError::ModelError(format!("linear helper module load: {e}"))
})?;
let add_fn = module.load_function("add_bias_kernel").map_err(|e| {
MLError::ModelError(format!("add_bias_kernel load: {e}"))
})?;
let reduce_fn = module.load_function("reduce_sum_axis0_kernel").map_err(|e| {
MLError::ModelError(format!("reduce_sum_axis0_kernel load: {e}"))
})?;
Ok((add_fn, reduce_fn))
/// Get bias-add and reduce-sum kernel functions.
///
/// STUB: nvrtc runtime compilation removed. These kernels will be provided
/// by precompiled cubins from build.rs in a future step. For now, this
/// returns an error indicating the kernels are not yet available.
fn get_bias_kernels(_stream: &Arc<CudaStream>) -> Result<(CudaFunction, CudaFunction), MLError> {
Err(MLError::ModelError(
"linear helper kernels not available: nvrtc removed, cubins not yet wired".to_owned(),
))
}
/// Add bias to a 2D tensor: Y[b, j] += bias[j].
@@ -462,9 +434,10 @@ mod tests {
for i in 0..4 {
identity[i * 4 + i] = 1.0;
}
stream.memcpy_htod(&identity, &mut store.get_mut("fc.weight").unwrap().data).unwrap();
let zeros = vec![0.0_f32; 4];
stream.memcpy_htod(&zeros, &mut store.get_mut("fc.bias").unwrap().data).unwrap();
let identity_bf16: Vec<half::bf16> = identity.iter().map(|&x| half::bf16::from_f32(x)).collect();
stream.memcpy_htod(&identity_bf16, &mut store.get_mut("fc.weight").unwrap().data).unwrap();
let zeros_bf16 = vec![half::bf16::ZERO; 4];
stream.memcpy_htod(&zeros_bf16, &mut store.get_mut("fc.bias").unwrap().data).unwrap();
let input = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0];
let x = GpuTensor::from_host(&input, vec![2, 4], &stream).unwrap();

24
crates/ml-core/src/cuda_autograd/loss.rs
View File

@@ -6,7 +6,6 @@
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use crate::MLError;
use super::gpu_tensor::GpuTensor;
@@ -34,22 +33,13 @@ impl std::fmt::Debug for LossResult {
impl LossKernels {
/// Compile all loss kernels.
pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(LOSS_CUDA_SRC, &context)
.map_err(|e| MLError::ModelError(format!("loss kernel compilation: {e}")))?;
let module = context.load_module(ptx).map_err(|e| {
MLError::ModelError(format!("loss module load: {e}"))
})?;
Ok(Self {
mse_kernel: module.load_function("mse_loss_with_grad").map_err(|e| {
MLError::ModelError(format!("mse_loss_with_grad load: {e}"))
})?,
huber_kernel: module.load_function("huber_loss_with_grad").map_err(|e| {
MLError::ModelError(format!("huber_loss_with_grad load: {e}"))
})?,
})
///
/// STUB: nvrtc runtime compilation removed. These kernels will be provided
/// by precompiled cubins from build.rs in a future step.
pub fn new(_stream: &Arc<CudaStream>) -> Result<Self, MLError> {
Err(MLError::ModelError(
"loss kernels not available: nvrtc removed, cubins not yet wired".to_owned(),
))
}
/// MSE loss: `L = mean((pred - target)^2)`, `dL/dpred = 2*(pred - target)/n`.

50
crates/ml-core/src/cuda_autograd/optimizer.rs
View File

@@ -11,7 +11,6 @@ use std::collections::BTreeMap;
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use crate::MLError;
use super::var_store::GpuVarStore;
@@ -97,34 +96,15 @@ impl std::fmt::Debug for GpuAdamW {
impl GpuAdamW {
/// Create a new AdamW optimizer.
///
/// Does NOT pre-allocate moment buffers — they are lazily created on the
/// first `step()` call for each parameter.
/// STUB: nvrtc runtime compilation removed. These kernels will be provided
/// by precompiled cubins from build.rs in a future step.
pub fn new(
config: AdamWConfig,
stream: Arc<CudaStream>,
_config: AdamWConfig,
_stream: Arc<CudaStream>,
) -> Result<Self, MLError> {
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(ADAMW_CUDA_SRC, &context)
.map_err(|e| MLError::ModelError(format!("AdamW kernel compilation: {e}")))?;
let module = context.load_module(ptx).map_err(|e| {
MLError::ModelError(format!("AdamW module load: {e}"))
})?;
let kernel = module.load_function("adamw_update").map_err(|e| {
MLError::ModelError(format!("adamw_update load: {e}"))
})?;
let grad_norm_kernel = module.load_function("grad_l2_norm").map_err(|e| {
MLError::ModelError(format!("grad_l2_norm load: {e}"))
})?;
Ok(Self {
config,
states: BTreeMap::new(),
step_count: 0,
kernel,
grad_norm_kernel,
stream,
})
Err(MLError::ModelError(
"AdamW optimizer not available: nvrtc removed, cubins not yet wired".to_owned(),
))
}
/// Current step count.
@@ -266,13 +246,13 @@ impl GpuAdamW {
}
}
// Read back the squared norm (single scalar, 4 bytes)
let mut norm_sq_host = [0.0_f32];
// Read back the squared norm (single scalar)
let mut norm_sq_host = [half::bf16::ZERO];
self.stream.memcpy_dtoh(&norm_sq_buf, &mut norm_sq_host).map_err(|e| { // gpu-exit: 1 scalar grad_norm for logging
MLError::ModelError(format!("AdamW norm readback: {e}"))
})?;
let norm = norm_sq_host[0].sqrt();
let norm = norm_sq_host[0].to_f32().sqrt();
// Compute scale factor: 1.0 if within budget, else max_norm / norm
let scale = if norm > self.config.max_grad_norm {
@@ -368,9 +348,9 @@ mod tests {
store.linear("fc", 4, 4).unwrap();
// Set weights to 1.0
let ones = vec![1.0_f32; 16];
let ones: Vec<half::bf16> = vec![half::bf16::from_f32(1.0); 16];
stream.memcpy_htod(&ones, &mut store.get_mut("fc.weight").unwrap().data).unwrap();
let bias_ones = vec![1.0_f32; 4];
let bias_ones: Vec<half::bf16> = vec![half::bf16::from_f32(1.0); 4];
stream.memcpy_htod(&bias_ones, &mut store.get_mut("fc.bias").unwrap().data).unwrap();
let mut optimizer = GpuAdamW::new(
@@ -394,9 +374,9 @@ mod tests {
// Weights should have decreased from 1.0
let w_host = {
let param = store.get("fc.weight").unwrap();
let mut h = vec![0.0_f32; param.data.len()];
stream.memcpy_dtoh(&param.data, &mut h).unwrap(); // test-only readback
h
let mut h_bf16 = vec![half::bf16::ZERO; param.data.len()];
stream.memcpy_dtoh(&param.data, &mut h_bf16).unwrap(); // test-only readback
h_bf16.iter().map(|x| x.to_f32()).collect::<Vec<f32>>()
};
for &w in &w_host {

68
crates/ml-core/src/cuda_autograd/reductions.rs
View File

@@ -17,7 +17,6 @@
use std::sync::Arc;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::nvrtc::Ptx;
use crate::MLError;
@@ -49,28 +48,13 @@ pub struct Stats {
impl ReductionKernels {
/// Compile all reduction kernels.
pub fn new(stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let context = stream.context();
let ptx: Ptx = crate::cuda_compile::compile_ptx_for_device(REDUCTION_CUDA_SRC, &context)
.map_err(|e| MLError::ModelError(format!("reduction kernel compilation: {e}")))?;
let module = context.load_module(ptx).map_err(|e| {
MLError::ModelError(format!("reduction module load: {e}"))
})?;
let load = |name: &str| -> Result<CudaFunction, MLError> {
module.load_function(name).map_err(|e| {
MLError::ModelError(format!("reduction kernel '{name}' load: {e}"))
})
};
Ok(Self {
fused_stats_fn: load("fused_stats_reduce")?,
argmax_flat_fn: load("argmax_flat")?,
argmax_rows_fn: load("argmax_rows")?,
sum_reduce_fn: load("sum_reduce")?,
col_sum_fn: load("col_sum_reduce")?,
stream: Arc::clone(stream),
})
///
/// STUB: nvrtc runtime compilation removed. These kernels will be provided
/// by precompiled cubins from build.rs in a future step.
pub fn new(_stream: &Arc<CudaStream>) -> Result<Self, MLError> {
Err(MLError::ModelError(
"reduction kernels not available: nvrtc removed, cubins not yet wired".to_owned(),
))
}
/// Compute min, max, mean, variance in a single kernel launch + single DtoH readback.
@@ -89,7 +73,13 @@ impl ReductionKernels {
})?;
// Initialize min to +INF, max to -INF, sum/sum_sq/count to 0
let init_vals = [f32::INFINITY, f32::NEG_INFINITY, 0.0_f32, 0.0_f32, 0.0_f32];
let init_vals = [
half::bf16::from_f32(f32::INFINITY),
half::bf16::from_f32(f32::NEG_INFINITY),
half::bf16::ZERO,
half::bf16::ZERO,
half::bf16::ZERO,
];
self.stream.memcpy_htod(&init_vals, &mut result).map_err(|e| {
MLError::ModelError(format!("stats init upload: {e}"))
})?;
@@ -121,11 +111,12 @@ impl ReductionKernels {
.map_err(|e| MLError::ModelError(format!("fused_stats_reduce: {e}")))?;
}
// Single DtoH readback: 5 floats = 20 bytes
let mut host = [0.0_f32; 5];
self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
// Single DtoH readback: 5 bf16 values
let mut host_bf16 = [half::bf16::ZERO; 5];
self.stream.memcpy_dtoh(&result, &mut host_bf16).map_err(|e| { // gpu-exit: reduction scalar readback
MLError::ModelError(format!("stats DtoH readback: {e}"))
})?;
let host: Vec<f32> = host_bf16.iter().map(|x| x.to_f32()).collect();
let min_val = host.first().copied().unwrap_or(0.0);
let max_val = host.get(1).copied().unwrap_or(0.0);
@@ -160,7 +151,10 @@ impl ReductionKernels {
let mut result = self.stream.alloc_zeros::<half::bf16>(2).map_err(|e| {
MLError::ModelError(format!("argmax result alloc: {e}"))
})?;
let init_vals = [f32::NEG_INFINITY, 0.0_f32]; // val, idx
let init_vals = [
half::bf16::from_f32(f32::NEG_INFINITY),
half::bf16::ZERO,
]; // val, idx
self.stream.memcpy_htod(&init_vals, &mut result).map_err(|e| {
MLError::ModelError(format!("argmax init upload: {e}"))
})?;
@@ -191,13 +185,13 @@ impl ReductionKernels {
.map_err(|e| MLError::ModelError(format!("argmax_flat: {e}")))?;
}
let mut host = [0.0_f32; 2];
self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
let mut host_bf16 = [half::bf16::ZERO; 2];
self.stream.memcpy_dtoh(&result, &mut host_bf16).map_err(|e| { // gpu-exit: reduction scalar readback
MLError::ModelError(format!("argmax DtoH readback: {e}"))
})?;
// Index is stored as __int_as_float in the kernel
Ok(f32::to_bits(host.get(1).copied().unwrap_or(0.0)))
Ok(f32::to_bits(host_bf16.get(1).map(|x| x.to_f32()).unwrap_or(0.0)))
}
/// Argmax per row for 2D data [rows, cols] -- returns `Vec<u32>` of length `rows`.
@@ -287,12 +281,12 @@ impl ReductionKernels {
.map_err(|e| MLError::ModelError(format!("sum_reduce: {e}")))?;
}
let mut host = [0.0_f32];
self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
let mut host_bf16 = [half::bf16::ZERO];
self.stream.memcpy_dtoh(&result, &mut host_bf16).map_err(|e| { // gpu-exit: reduction scalar readback
MLError::ModelError(format!("sum DtoH readback: {e}"))
})?;
Ok(host.first().copied().unwrap_or(0.0))
Ok(host_bf16.first().map(|x| x.to_f32()).unwrap_or(0.0))
}
/// Per-column sum for 2D data [rows, cols] -- returns `Vec<f32>` of length `cols`.
@@ -336,12 +330,12 @@ impl ReductionKernels {
.map_err(|e| MLError::ModelError(format!("col_sum_reduce: {e}")))?;
}
let mut host = vec![0.0_f32; cols];
self.stream.memcpy_dtoh(&result, &mut host).map_err(|e| { // gpu-exit: reduction scalar readback
let mut host_bf16 = vec![half::bf16::ZERO; cols];
self.stream.memcpy_dtoh(&result, &mut host_bf16).map_err(|e| { // gpu-exit: reduction scalar readback
MLError::ModelError(format!("col_sums DtoH readback: {e}"))
})?;
Ok(host)
Ok(host_bf16.iter().map(|x| x.to_f32()).collect())
}
}

30
crates/ml-core/src/cuda_autograd/stream_ops.rs
View File

@@ -121,6 +121,8 @@ impl StreamTensor {
}
/// Create a tensor from a host `Vec<f32>`.
///
/// Converts f32 to bf16 before upload.
pub fn from_vec(values: Vec<f32>, shape: &[usize], stream: &Arc<CudaStream>) -> Result<Self, MLError> {
let n: usize = shape.iter().product();
if values.len() != n {
@@ -129,8 +131,9 @@ impl StreamTensor {
actual: values.len(),
});
}
let bf16_values: Vec<half::bf16> = values.iter().map(|&x| half::bf16::from_f32(x)).collect();
let data = stream
.clone_htod(&values)
.clone_htod(&bf16_values)
.map_err(|e| MLError::DeviceError(format!("StreamTensor htod: {e}")))?;
Ok(Self {
data,
@@ -158,14 +161,16 @@ impl StreamTensor {
/// Copy tensor data back to host.
///
/// Downloads bf16, converts to f32 for the caller.
/// **WARNING**: This is a synchronous GPU-to-CPU download. It is NOT
/// suitable for use in hot training / inference paths. Use it only for
/// checkpointing, logging, or tests.
#[cold]
pub fn to_vec(&self) -> Result<Vec<f32>, MLError> {
self.stream
let bf16_vec: Vec<half::bf16> = self.stream
.clone_dtoh(&self.data) // gpu-exit: API boundary readback
.map_err(|e| MLError::DeviceError(format!("StreamTensor dtoh: {e}")))
.map_err(|e| MLError::DeviceError(format!("StreamTensor dtoh: {e}")))?;
Ok(bf16_vec.iter().map(|x| x.to_f32()).collect())
}
/// Reshape (no data copy -- just changes shape metadata).
@@ -226,13 +231,15 @@ impl StreamLinear {
let n = out_features * in_features;
let mut rng = rand::thread_rng();
let weight_data: Vec<f32> = (0..n).map(|_| rng.gen_range(-limit..limit)).collect();
let weight_data: Vec<half::bf16> = (0..n)
.map(|_| half::bf16::from_f32(rng.gen_range(-limit..limit)))
.collect();
let weight = stream
.clone_htod(&weight_data)
.map_err(|e| MLError::DeviceError(format!("StreamLinear weight htod: {e}")))?;
let bias_data = vec![0.0_f32; out_features];
let bias_data = vec![half::bf16::ZERO; out_features];
let bias = stream
.clone_htod(&bias_data)
.map_err(|e| MLError::DeviceError(format!("StreamLinear bias htod: {e}")))?;
@@ -342,23 +349,28 @@ impl StreamLinear {
}
/// Collect weight data to host for checkpoint serialization.
///
/// Downloads bf16, converts to f32 for the caller.
#[cold]
pub fn weight_to_vec(&self) -> Result<Vec<f32>, MLError> {
self.stream
let bf16_vec: Vec<half::bf16> = self.stream
.clone_dtoh(&self.weight) // gpu-exit: checkpoint weight export
.map_err(|e| MLError::DeviceError(format!("weight dtoh: {e}")))
.map_err(|e| MLError::DeviceError(format!("weight dtoh: {e}")))?;
Ok(bf16_vec.iter().map(|x| x.to_f32()).collect())
}
/// Collect bias data to host for checkpoint serialization.
///
/// Downloads bf16, converts to f32 for the caller.
#[cold]
pub fn bias_to_vec(&self) -> Result<Option<Vec<f32>>, MLError> {
match &self.bias {
Some(b) => {
let v = self
let bf16_vec: Vec<half::bf16> = self
.stream
.clone_dtoh(b) // gpu-exit: checkpoint weight export
.map_err(|e| MLError::DeviceError(format!("bias dtoh: {e}")))?;
Ok(Some(v))
Ok(Some(bf16_vec.iter().map(|x| x.to_f32()).collect()))
}
None => Ok(None),
}

27
crates/ml-core/src/cuda_autograd/var_store.rs
View File

@@ -164,11 +164,12 @@ impl GpuVarStore {
pub fn all_vars(&self) -> Result<Vec<GpuTensor>, MLError> {
let mut out = Vec::with_capacity(self.params.len());
for (name, param) in &self.params {
let mut host = vec![0.0_f32; param.data.len()];
self.stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| { // gpu-exit: checkpoint export
let mut host_bf16 = vec![half::bf16::ZERO; param.data.len()];
self.stream.memcpy_dtoh(&param.data, &mut host_bf16).map_err(|e| { // gpu-exit: checkpoint export
MLError::ModelError(format!("all_vars DtoH '{name}': {e}"))
})?;
let t = GpuTensor::from_host(&host, param.shape.clone(), &self.stream)?;
let host_f32: Vec<f32> = host_bf16.iter().map(|x| x.to_f32()).collect();
let t = GpuTensor::from_host(&host_f32, param.shape.clone(), &self.stream)?;
out.push(t);
}
Ok(out)
@@ -287,11 +288,12 @@ impl GpuVarStore {
for (name, param) in &self.params {
let offset = host_flat.len();
let count = param.data.len();
let mut host_param = vec![0.0_f32; count];
self.stream.memcpy_dtoh(&param.data, &mut host_param).map_err(|e| { // gpu-exit: checkpoint export
let mut host_bf16 = vec![half::bf16::ZERO; count];
self.stream.memcpy_dtoh(&param.data, &mut host_bf16).map_err(|e| { // gpu-exit: checkpoint export
MLError::ModelError(format!("flatten DtoH '{name}': {e}"))
})?;
host_flat.extend_from_slice(&host_param);
let host_f32: Vec<f32> = host_bf16.iter().map(|x| x.to_f32()).collect();
host_flat.extend_from_slice(&host_f32);
offsets.insert(name.clone(), (offset, count));
}
@@ -311,7 +313,8 @@ impl GpuVarStore {
for (name, &(offset, count)) in offsets {
if let Some(param) = self.params.get_mut(name) {
let slice = &flat_host[offset..offset + count];
self.stream.memcpy_htod(slice, &mut param.data).map_err(|e| {
let bf16_slice: Vec<half::bf16> = slice.iter().map(|&x| half::bf16::from_f32(x)).collect();
self.stream.memcpy_htod(&bf16_slice, &mut param.data).map_err(|e| {
MLError::ModelError(format!("unflatten HtoD '{name}': {e}"))
})?;
}
@@ -331,11 +334,12 @@ impl GpuVarStore {
) -> Result<BTreeMap<String, (Vec<usize>, Vec<f32>)>, MLError> {
let mut out = BTreeMap::new();
for (name, param) in &self.params {
let mut host = vec![0.0_f32; param.data.len()];
self.stream.memcpy_dtoh(&param.data, &mut host).map_err(|e| { // gpu-exit: checkpoint export
let mut host_bf16 = vec![half::bf16::ZERO; param.data.len()];
self.stream.memcpy_dtoh(&param.data, &mut host_bf16).map_err(|e| { // gpu-exit: checkpoint export
MLError::ModelError(format!("export DtoH '{name}': {e}"))
})?;
out.insert(name.clone(), (param.shape.clone(), host));
let host_f32: Vec<f32> = host_bf16.iter().map(|x| x.to_f32()).collect();
out.insert(name.clone(), (param.shape.clone(), host_f32));
}
Ok(out)
}
@@ -364,7 +368,8 @@ impl GpuVarStore {
param.shape, shape
)));
}
self.stream.memcpy_htod(values, &mut param.data).map_err(|e| {
let bf16_values: Vec<half::bf16> = values.iter().map(|&x| half::bf16::from_f32(x)).collect();
self.stream.memcpy_htod(&bf16_values, &mut param.data).map_err(|e| {
MLError::ModelError(format!("import HtoD '{name}': {e}"))
})?;
}

213
crates/ml-core/src/cuda_compile.rs
View File

@@ -1,213 +0,0 @@
//! CUDA kernel precompilation utilities.
//!
//! Compiles `.cu` source to native cubin (SASS) via `nvcc` with full `-O3`
//! optimization and a deterministic SHA-256 disk cache. The cubin is
//! architecture-specific (`sm_XX`) and eliminates driver JIT entirely.
//!
//! Moved from `ml::cuda_pipeline` so that `ml-dqn` (which depends on `ml-core`
//! but not `ml`) can also precompile kernels without a circular dependency.
/// nvcc optimization flags applied to every kernel compilation.
///
/// Changing this array auto-invalidates the cubin cache (flags are hashed into
/// the cache key). No manual version bump needed.
///
/// Note: `--extra-device-vectorization` is deliberately omitted -- it causes
/// catastrophic register spill on the 3000-line experience kernel (sm_90),
/// blowing the per-thread stack and triggering CUDA_ERROR_ILLEGAL_ADDRESS.
const NVCC_OPT_FLAGS: &[&str] = &[
"-O3",
"--use_fast_math",
"--ftz=true",
"--fmad=true",
];
/// Compile CUDA source to a native cubin for the device behind `context`.
///
/// Uses a deterministic SHA-256 disk cache keyed on `(arch, source, flags)`.
/// First call for a given kernel + architecture compiles via `nvcc` (offline);
/// subsequent calls load the cached cubin in <10 ms.
///
/// The experience-collector kernel (3000+ lines of `.cu`) typically
/// takes ~15s to compile via nvcc on H100. With caching, subsequent runs load
/// in <10ms (file existence check only).
///
/// # Feature gate
/// Only available with the `cuda` feature.
///
/// # Compilation strategy
///
/// Uses nvcc (offline compiler) with full `-O3` optimization and `-cubin` output.
/// Produces native SASS for the exact GPU architecture (sm_XX), not virtual PTX.
/// Requires nvcc in `$CUDA_HOME/bin/` or `$PATH`.
pub fn compile_ptx_for_device(
src: &str,
context: &cudarc::driver::CudaContext,
) -> Result<cudarc::nvrtc::Ptx, String> {
use cudarc::driver::sys::CUdevice_attribute;
let major = context
.attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR)
.unwrap_or(0);
let minor = context
.attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR)
.unwrap_or(0);
// Build arch string for nvcc: "sm_XX" (real arch for native cubin output).
// cubin is GPU-specific but eliminates driver JIT entirely.
let arch_str: &'static str = match (major, minor) {
(9, 0) => "sm_90",
(9, _) => "sm_90", // sm_90a and beyond
(8, 9) => "sm_89",
(8, 6) => "sm_86",
(8, 0) => "sm_80",
(7, 5) => "sm_75",
(7, 0) => "sm_70",
_ => "sm_80", // safe default
};
// Try loading cached cubin first
if let Some(cached) = load_cached_cubin(arch_str, src) {
return Ok(cached);
}
// Cache miss -- compile via nvcc to native cubin
let cubin = compile_cubin_with_nvcc(src, arch_str)?;
Ok(cubin)
}
/// Compile CUDA source to native cubin using nvcc with maximum optimization.
///
/// See [`NVCC_OPT_FLAGS`] for the flags applied. `-cubin` and `-arch=sm_XX`
/// are added automatically. Output is a native SASS binary -- no driver JIT.
/// The cubin is written directly to the cache directory and loaded via
/// `Ptx::from_file()` which calls `cuModuleLoad()`.
fn compile_cubin_with_nvcc(
src: &str,
arch: &str,
) -> Result<cudarc::nvrtc::Ptx, String> {
use std::io::Write;
// Find nvcc: prefer $CUDA_HOME/bin/nvcc, fall back to PATH
let nvcc = std::env::var("CUDA_HOME")
.map(|home| std::path::PathBuf::from(home).join("bin/nvcc"))
.unwrap_or_else(|_| std::path::PathBuf::from("nvcc"));
// Write source to temp file (nvcc requires file input)
let tmp_dir = std::env::temp_dir().join("foxhunt_nvcc");
std::fs::create_dir_all(&tmp_dir)
.map_err(|e| format!("Failed to create nvcc temp dir: {e}"))?;
// Ensure cache dir exists -- nvcc writes cubin directly there
let cache_dir = cubin_cache_dir();
std::fs::create_dir_all(&cache_dir)
.map_err(|e| format!("Failed to create cubin cache dir: {e}"))?;
let key = cubin_cache_key(arch, src);
let hash_prefix = key.get(..16).unwrap_or(&key);
let src_path = tmp_dir.join(format!("{hash_prefix}.cu"));
let cubin_path = cache_dir.join(format!("{key}.cubin"));
{
let mut f = std::fs::File::create(&src_path)
.map_err(|e| format!("Failed to write CUDA source to {}: {e}", src_path.display()))?;
f.write_all(src.as_bytes())
.map_err(|e| format!("Failed to write CUDA source: {e}"))?;
}
// Invoke nvcc with -cubin (native SASS) instead of -ptx (virtual ISA).
// This eliminates driver JIT entirely -- cuModuleLoad loads SASS directly.
let arch_flag = format!("-arch={arch}");
let out_path_str = cubin_path.to_str().unwrap_or("out.cubin");
let in_path = src_path.to_str().unwrap_or("in.cu");
let mut args: Vec<&str> = vec!["-cubin"];
args.extend_from_slice(NVCC_OPT_FLAGS);
args.extend_from_slice(&[&arch_flag, "-o", out_path_str, in_path]);
let output = std::process::Command::new(&nvcc)
.args(&args)
.output()
.map_err(|e| format!(
"Failed to execute nvcc at {}: {e}. \
Ensure CUDA toolkit is installed and nvcc is in $CUDA_HOME/bin/ or $PATH.",
nvcc.display()
))?;
if !output.status.success() {
let stderr = String::from_utf8_lossy(&output.stderr);
let stdout = String::from_utf8_lossy(&output.stdout);
return Err(format!(
"nvcc compilation failed (exit={}):\n{stderr}\n{stdout}",
output.status.code().unwrap_or(-1)
));
}
// Verify the cubin was written
let cubin_size = std::fs::metadata(&cubin_path)
.map_err(|e| format!("Failed to read compiled cubin from {}: {e}", cubin_path.display()))?
.len();
// Clean up source temp file (best-effort), cubin stays in cache
drop(std::fs::remove_file(&src_path));
let flags_str = NVCC_OPT_FLAGS.join(" ");
tracing::info!(
arch,
src_len = src.len(),
cubin_len = cubin_size,
flags = %flags_str,
"nvcc cubin compilation succeeded (no driver JIT)"
);
// Load via Ptx::from_file -- cuModuleLoad() handles cubin natively
Ok(cudarc::nvrtc::Ptx::from_file(&cubin_path))
}
/// Resolve the cubin cache directory.
///
/// Prefers `$CARGO_TARGET_DIR/.cubin_cache/` (persisted on CI PVC between runs).
/// Falls back to `/tmp/.cubin_cache/` when `CARGO_TARGET_DIR` is unset.
fn cubin_cache_dir() -> std::path::PathBuf {
let base = std::env::var("CARGO_TARGET_DIR")
.map(std::path::PathBuf::from)
.unwrap_or_else(|_| std::path::PathBuf::from("/tmp"));
base.join(".cubin_cache")
}
/// Compute deterministic SHA-256 cache key from (arch, source, flags).
///
/// Uses SHA-256 (not `DefaultHasher`) because `SipHash` uses random keys
/// per process -- the cache would never hit across separate cargo test runs.
fn cubin_cache_key(arch: &str, src: &str) -> String {
use sha2::{Sha256, Digest};
let mut h = Sha256::new();
h.update(arch.as_bytes());
h.update(b"\x00"); // separator
// Include compile flags so any flag change auto-invalidates cached cubin.
for flag in NVCC_OPT_FLAGS {
h.update(flag.as_bytes());
h.update(b"\x00");
}
h.update(src.as_bytes());
format!("{:x}", h.finalize())
}
/// Try to load cached cubin from disk.
fn load_cached_cubin(
arch: &str,
src: &str,
) -> Option<cudarc::nvrtc::Ptx> {
let key = cubin_cache_key(arch, src);
let cache_path = cubin_cache_dir().join(format!("{key}.cubin"));
if cache_path.exists() {
let cubin_size = std::fs::metadata(&cache_path).map(|m| m.len()).unwrap_or(0);
tracing::info!(
"cubin cache HIT: {} ({} bytes)",
cache_path.display(),
cubin_size,
);
Some(cudarc::nvrtc::Ptx::from_file(&cache_path))
} else {
tracing::info!("cubin cache MISS: {}", cache_path.display());
None
}
}

3
crates/ml-core/src/lib.rs
View File

@@ -134,8 +134,7 @@ pub mod safety;
pub mod memory_optimization;
pub mod batch_size_resolver;
// ========== CUDA PRECOMPILATION (nvcc cubin cache) ==========
pub mod cuda_compile;
// cuda_compile module removed — nvrtc replaced by build.rs cubins
// ========== CUDA-NATIVE AUTOGRAD (replaces Candle backward/VarMap/GradStore) ==========
#[cfg(feature = "cuda")]