diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 88e62c7cd..5f288dfc7 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -81,6 +81,7 @@ fn main() { "backward_kernels.cu", "iqn_cvar_kernel.cu", "mamba2_temporal_kernel.cu", + "graph_utility_kernels.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a9b7f79d6..94911606a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -71,6 +71,7 @@ static Q_STATS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/q_stats_ static ENSEMBLE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ensemble_kernels.cubin")); static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_grad_kernel.cubin")); static MAMBA2_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mamba2_temporal_kernel.cubin")); +pub(crate) static GRAPH_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/graph_utility_kernels.cubin")); /// Mamba2 temporal scan configuration. const MAMBA2_HISTORY_K: usize = 8; // Rolling history length diff --git a/crates/ml/src/cuda_pipeline/graph_utility_kernels.cu b/crates/ml/src/cuda_pipeline/graph_utility_kernels.cu new file mode 100644 index 000000000..4d6d398b5 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/graph_utility_kernels.cu @@ -0,0 +1,146 @@ +/** + * Graph utility kernels — GPU-native gather + step counter increment. + * + * These kernels replace host-side operations on the training hot path, + * enabling a true single-graph architecture with zero CPU involvement + * per step (single cuGraphExecLaunch). + * + * Kernels: + * gather_f32_rows_padded — gather rows from replay buffer + zero-pad + * to 128-byte-aligned dst_stride (replaces + * 2-step gather + pad_states host pipeline) + * gather_f32_scalar — gather scalar f32 values (rewards, dones, + * IS-weights) from replay buffer indices + * gather_i32_scalar — gather scalar i32 values (actions) from + * replay buffer indices + * increment_step_counters — single-thread kernel that increments all + * per-step Adam/scheduler counters on GPU + * and computes cosine-annealed tau + * + * No #include of common_device_functions.cuh is needed here — the build + * function in build.rs prepends it via string concatenation before + * passing the combined source to nvcc. + */ + +/* ══════════════════════════════════════════════════════════════════════ + * gather_f32_rows_padded + * + * Gathers rows from a flat row-major src buffer into dst, zero-padding + * columns [src_dim, dst_stride) so the output is 128-byte aligned per row. + * + * Launch config: grid=(ceil(batch_size * dst_stride / 256)), block=(256) + * + * Parameters: + * dst — output buffer [batch_size * dst_stride] + * src — source buffer [capacity * src_dim] + * indices — i64 indices into src rows [batch_size] + * src_dim — number of valid f32 elements per src row (unpadded) + * dst_stride — padded row width in f32 elements (dst_dim >= src_dim, + * aligned to 128 bytes = 32 f32 elements) + * batch_size — number of rows to gather + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void gather_f32_rows_padded( + float* __restrict__ dst, + const float* __restrict__ src, + const long long* __restrict__ indices, + int src_dim, int dst_stride, int batch_size) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int total = batch_size * dst_stride; + if (tid >= total) return; + int row = tid / dst_stride; + int col = tid % dst_stride; + long long src_row = indices[row]; + dst[tid] = (col < src_dim) ? src[src_row * src_dim + col] : 0.0f; +} + +/* ══════════════════════════════════════════════════════════════════════ + * gather_f32_scalar + * + * Gathers scalar f32 values (rewards, dones, IS-weights) from the replay + * buffer into a flat output array using the sampled indices. + * + * Launch config: grid=(ceil(batch_size / 256)), block=(256) + * + * Parameters: + * dst — output buffer [batch_size] + * src — source buffer [capacity] (replay buffer scalar field) + * indices — i64 indices [batch_size] + * batch_size — number of elements to gather + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void gather_f32_scalar( + float* __restrict__ dst, + const float* __restrict__ src, + const long long* __restrict__ indices, + int batch_size) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= batch_size) return; + dst[tid] = src[indices[tid]]; +} + +/* ══════════════════════════════════════════════════════════════════════ + * gather_i32_scalar + * + * Gathers scalar i32 values (packed actions) from the replay buffer into + * a flat output array using the sampled indices. + * + * Launch config: grid=(ceil(batch_size / 256)), block=(256) + * + * Parameters: + * dst — output buffer [batch_size] + * src — source buffer [capacity] (replay buffer action field) + * indices — i64 indices [batch_size] + * batch_size — number of elements to gather + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void gather_i32_scalar( + int* __restrict__ dst, + const int* __restrict__ src, + const long long* __restrict__ indices, + int batch_size) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= batch_size) return; + dst[tid] = src[indices[tid]]; +} + +/* ══════════════════════════════════════════════════════════════════════ + * increment_step_counters + * + * Single-thread kernel that atomically increments all per-step counters + * and recomputes the cosine-annealed IQN tau value — all on GPU, zero + * CPU synchronisation required. + * + * All counter pointers must be device-accessible (device alloc or pinned + * host-mapped memory). tau_out is written on every call. + * + * Launch config: grid=(1), block=(1) + * + * Parameters: + * adam_t, sel_t, denoise_t, iql_t, iql_low_t, iqn_t, attn_t, rng_step + * — per-scheduler step counters (int*, device-mapped) + * tau_out — cosine-annealed tau written here (float*, device-mapped) + * tau_init — initial tau value (start of annealing) + * tau_final — final tau value (end of annealing) + * anneal_steps — total number of steps over which to anneal + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void increment_step_counters( + int* adam_t, int* sel_t, int* denoise_t, int* iql_t, + int* iql_low_t, int* iqn_t, int* attn_t, int* rng_step, + float* tau_out, float tau_init, float tau_final, int anneal_steps) +{ + /* adam_t is the master step — its post-increment value drives the + * cosine schedule so that step 1 produces the first non-initial tau. */ + int step = atomicAdd(adam_t, 1) + 1; + atomicAdd(sel_t, 1); + atomicAdd(denoise_t, 1); + atomicAdd(iql_t, 1); + atomicAdd(iql_low_t, 1); + atomicAdd(iqn_t, 1); + atomicAdd(attn_t, 1); + atomicAdd(rng_step, 1); + + float progress = fminf((float)step / (float)anneal_steps, 1.0f); + float cosine = 0.5f * (1.0f + cosf(progress * 3.14159265f)); + *tau_out = tau_final + (tau_init - tau_final) * cosine; +}