diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index ba649762f..3fde4d31a 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -24,6 +24,7 @@ const KERNELS: &[&str] = &[ "output_smoothness", // CRT.train: per-horizon adjacent-position prob-jitter penalty "smoothness_lambda_controller", // CRT.train: ISV-driven λ controller anchored on h30 jitter "gpu_log_ring", // GPU diagnostic log ring — tick kernel + log_record helper + "bucket_transition_kernels", // Per-horizon CfC Phase 1→2 transition: tau_sort, bucket_assign, bucket_iqr, tau_reorder, heads_compact ]; // Cache bust v15 (2026-05-21): gpu_log_helpers.cuh extracted + smoothness_lambda_controller emits records. diff --git a/crates/ml-alpha/cuda/bucket_transition_kernels.cu b/crates/ml-alpha/cuda/bucket_transition_kernels.cu new file mode 100644 index 000000000..6f5a5a1e3 --- /dev/null +++ b/crates/ml-alpha/cuda/bucket_transition_kernels.cu @@ -0,0 +1,193 @@ +// bucket_transition_kernels.cu — Phase 1→2 transition kernels. +// +// All-on-device per feedback_no_htod_htoh_only_mapped_pinned. No atomicAdd +// per feedback_no_atomicadd; reductions use warp-shuffle. Static quintile +// boundaries for HIDDEN_DIM=128: [0, 25, 50, 75, 100, 128]. + +#define HIDDEN_DIM 128 +#define N_HORIZONS 5 +#define BUCKET_DIM_LAST 28 // last bucket absorbs HIDDEN_DIM - 4*25 = 28 +#define MAX_BUCKET_DIM 28 + +// Static bucket boundaries (could also be __constant__ memory). +__device__ const unsigned int BUCKET_OFFSETS[N_HORIZONS + 1] = {0, 25, 50, 75, 100, 128}; + +// ───────────────────────────────────────────────────────────────────── +// tau_sort_kernel: bitonic-merge sort of cfc.tau values. +// +// Launch: 1 block × 32 threads, shared_mem_bytes >= HIDDEN_DIM * 8 bytes +// (4 for tau key, 4 for original index value). +// +// Sorts ascending. Output: sorted_indices_d[p] = original channel index +// whose tau is the p-th smallest. +// +// Algorithm: bitonic merge sort over 128 elements, single block, 4 passes +// of 32 threads cooperating. Each thread handles 4 elements per pass. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void tau_sort_kernel( + const float* __restrict__ tau, // [HIDDEN_DIM] + unsigned int* __restrict__ sorted_indices // [HIDDEN_DIM] +) { + extern __shared__ float smem[]; + float* keys = smem; // [HIDDEN_DIM] + unsigned int* vals = (unsigned int*)(smem + HIDDEN_DIM); // [HIDDEN_DIM] + + int tid = threadIdx.x; + // Each thread loads 4 elements (128 / 32 = 4) + for (int i = 0; i < 4; ++i) { + int idx = tid + i * 32; + keys[idx] = tau[idx]; + vals[idx] = (unsigned int)idx; + } + __syncthreads(); + + // Bitonic sort over 128 elements. + for (int k = 2; k <= HIDDEN_DIM; k <<= 1) { + for (int j = k >> 1; j > 0; j >>= 1) { + for (int i = 0; i < 4; ++i) { + int idx = tid + i * 32; + int partner = idx ^ j; + if (partner > idx) { + bool ascending = ((idx & k) == 0); + bool swap = ascending ? (keys[idx] > keys[partner]) : (keys[idx] < keys[partner]); + if (swap) { + float tk = keys[idx]; keys[idx] = keys[partner]; keys[partner] = tk; + unsigned int tv = vals[idx]; vals[idx] = vals[partner]; vals[partner] = tv; + } + } + } + __syncthreads(); + } + } + + // Write sorted_indices. + for (int i = 0; i < 4; ++i) { + int idx = tid + i * 32; + sorted_indices[idx] = vals[idx]; + } +} + +// ───────────────────────────────────────────────────────────────────── +// bucket_assign_kernel: from sorted_indices, write bucket_id per channel. +// +// Launch: 1 block × HIDDEN_DIM threads. +// Each thread handles one channel. Looks up its rank (position in sorted) +// and assigns bucket id via static quintile boundaries. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void bucket_assign_kernel( + const unsigned int* __restrict__ sorted_indices, // [HIDDEN_DIM] + unsigned char* __restrict__ bucket_id_per_channel // [HIDDEN_DIM] +) { + int tid = threadIdx.x; + if (tid >= HIDDEN_DIM) return; + // tid is a sorted-rank position; sorted_indices[tid] is the original channel + // at this rank. Assign bucket = quintile of rank. + unsigned char bucket = 0; + if (tid >= 100) bucket = 4; + else if (tid >= 75) bucket = 3; + else if (tid >= 50) bucket = 2; + else if (tid >= 25) bucket = 1; + bucket_id_per_channel[sorted_indices[tid]] = bucket; +} + +// ───────────────────────────────────────────────────────────────────── +// bucket_iqr_kernel: per-bucket Q1 and Q3 of tau values. +// +// Launch: N_HORIZONS blocks × 32 threads. Each block computes Q1, Q3 for +// its bucket via warp-shuffle reduction. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void bucket_iqr_kernel( + const float* __restrict__ tau, // [HIDDEN_DIM] + const unsigned int* __restrict__ sorted_indices, // [HIDDEN_DIM] + float* __restrict__ q1_per_bucket, // [N_HORIZONS] + float* __restrict__ q3_per_bucket // [N_HORIZONS] +) { + int bucket = blockIdx.x; + int bucket_start = BUCKET_OFFSETS[bucket]; + int bucket_end = BUCKET_OFFSETS[bucket + 1]; + int bucket_size = bucket_end - bucket_start; + int q1_idx = bucket_start + bucket_size / 4; + int q3_idx = bucket_start + (3 * bucket_size) / 4; + if (threadIdx.x == 0) { + // tau values are already in sorted-rank order via sorted_indices. + q1_per_bucket[bucket] = tau[sorted_indices[q1_idx]]; + q3_per_bucket[bucket] = tau[sorted_indices[q3_idx]]; + } +} + +// ───────────────────────────────────────────────────────────────────── +// tau_reorder_kernel: reorder tau values into bucket-grouped layout. +// +// Output tau_all_d[c] is the tau value of the original channel that +// landed at compact position c in bucket-grouped order. +// +// For simplicity: tau_all_d shares the same indexing as the original tau +// (channels in their natural order), but each channel's value is read +// via the bucket_id to position mapping that's just identity here. +// +// Launch: 1 block × HIDDEN_DIM threads. Per-channel write. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void tau_reorder_kernel( + const float* __restrict__ tau, // [HIDDEN_DIM] + const unsigned char* __restrict__ bucket_id_per_channel, // [HIDDEN_DIM] + const unsigned int* __restrict__ bucket_channel_offset, // [N_HORIZONS + 1] + float* __restrict__ tau_all_d // [HIDDEN_DIM] +) { + extern __shared__ unsigned int per_bucket_cursor[]; // [N_HORIZONS + 1] + int tid = threadIdx.x; + + // Initialize per-bucket cursors to bucket starts. + if (tid <= N_HORIZONS) { + per_bucket_cursor[tid] = bucket_channel_offset[tid]; + } + __syncthreads(); + + // Sequential writes per channel (race avoided by serializing the cursor increment). + // For HIDDEN_DIM=128 this is fast and simple. atomicAdd not used (per + // feedback_no_atomicadd); we use a single-block ordering via __syncthreads. + for (int c = 0; c < HIDDEN_DIM; ++c) { + if (tid == 0) { + unsigned char k = bucket_id_per_channel[c]; + unsigned int pos = per_bucket_cursor[k]; + tau_all_d[pos] = tau[c]; + per_bucket_cursor[k] = pos + 1; + } + __syncthreads(); + } +} + +// ───────────────────────────────────────────────────────────────────── +// heads_compact_kernel: reorder heads_w_skip into compact ragged layout. +// +// Original heads_w_skip is [N_HORIZONS × HIDDEN_DIM]. Output is compact +// [HIDDEN_DIM] where each horizon h owns BUCKET_DIM[h] contiguous floats +// at position bucket_channel_offset[h]. +// +// For each compact-position c in bucket k, read original +// heads_w_skip[k * HIDDEN_DIM + channel_that_landed_at_c]. +// +// Launch: 1 block × HIDDEN_DIM threads. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void heads_compact_kernel( + const float* __restrict__ w_skip_orig, // [N_HORIZONS × HIDDEN_DIM] + const unsigned char* __restrict__ bucket_id_per_channel, // [HIDDEN_DIM] + const unsigned int* __restrict__ bucket_channel_offset, // [N_HORIZONS + 1] + float* __restrict__ w_skip_compact // [HIDDEN_DIM] +) { + extern __shared__ unsigned int per_bucket_cursor[]; // [N_HORIZONS + 1] + int tid = threadIdx.x; + if (tid <= N_HORIZONS) { + per_bucket_cursor[tid] = bucket_channel_offset[tid]; + } + __syncthreads(); + + for (int c = 0; c < HIDDEN_DIM; ++c) { + if (tid == 0) { + unsigned char k = bucket_id_per_channel[c]; + unsigned int pos = per_bucket_cursor[k]; + w_skip_compact[pos] = w_skip_orig[(unsigned int)k * HIDDEN_DIM + c]; + per_bucket_cursor[k] = pos + 1; + } + __syncthreads(); + } +} diff --git a/crates/ml-alpha/tests/bucket_transition_kernels.rs b/crates/ml-alpha/tests/bucket_transition_kernels.rs new file mode 100644 index 000000000..f4e76bda2 --- /dev/null +++ b/crates/ml-alpha/tests/bucket_transition_kernels.rs @@ -0,0 +1,184 @@ +//! GPU oracle tests for bucket_transition_kernels.cu. +//! Run with: SQLX_OFFLINE=true cargo test -p ml-alpha --test bucket_transition_kernels -- --ignored --nocapture + +use anyhow::{Context, Result}; +use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use ml_core::device::MlDevice; + +const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bucket_transition_kernels.cubin")); +const HIDDEN_DIM: usize = 128; +const N_HORIZONS: usize = 5; + +fn load_kernel(stream: &std::sync::Arc, fn_name: &str) -> Result { + let ctx = stream.context(); + let module = ctx.load_cubin(KERNEL_CUBIN.to_vec()).context("load bucket_transition cubin")?; + module.load_function(fn_name).with_context(|| format!("load {fn_name}")) +} + +fn upload(stream: &std::sync::Arc, host: &[T]) -> Result> { + let mut d = stream.alloc_zeros::(host.len())?; + stream.memcpy_htod(host, &mut d)?; + Ok(d) +} + +fn download(stream: &std::sync::Arc, d: &CudaSlice) -> Result> { + let mut h = vec![T::default(); d.len()]; + stream.memcpy_dtoh(d, &mut h)?; + Ok(h) +} + +#[test] +#[ignore = "requires CUDA"] +fn tau_sort_kernel_sorts_tau_ascending() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "tau_sort_kernel")?; + + // Reverse-sorted input: [128, 127, ..., 1] of tau + let tau_host: Vec = (1..=HIDDEN_DIM).rev().map(|x| x as f32).collect(); + let tau_d = upload(&stream, &tau_host)?; + let mut sorted_indices_d = stream.alloc_zeros::(HIDDEN_DIM)?; + + let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (32,1,1), shared_mem_bytes: HIDDEN_DIM as u32 * 8 }; + let mut launch = stream.launch_builder(&func); + launch.arg(&tau_d).arg(&mut sorted_indices_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let sorted_indices = download(&stream, &sorted_indices_d)?; + // After sort ascending, sorted_indices[0] should point to the smallest tau (channel 127), [127] to largest (channel 0) + assert_eq!(sorted_indices[0], (HIDDEN_DIM - 1) as u32, "smallest tau should be at channel HIDDEN_DIM-1"); + assert_eq!(sorted_indices[HIDDEN_DIM - 1], 0, "largest tau should be at channel 0"); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA"] +fn bucket_assign_kernel_assigns_quintiles() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "bucket_assign_kernel")?; + + // sorted_indices[p] = p (identity permutation; channel p has rank p) + let sorted_indices_host: Vec = (0..HIDDEN_DIM as u32).collect(); + let sorted_indices_d = upload(&stream, &sorted_indices_host)?; + let mut bucket_id_d = stream.alloc_zeros::(HIDDEN_DIM)?; + + let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (HIDDEN_DIM as u32,1,1), shared_mem_bytes: 0 }; + let mut launch = stream.launch_builder(&func); + launch.arg(&sorted_indices_d).arg(&mut bucket_id_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let bucket_id = download(&stream, &bucket_id_d)?; + // bucket boundaries: [0,25), [25,50), [50,75), [75,100), [100,128) + assert_eq!(bucket_id[0], 0); + assert_eq!(bucket_id[24], 0); + assert_eq!(bucket_id[25], 1); + assert_eq!(bucket_id[49], 1); + assert_eq!(bucket_id[50], 2); + assert_eq!(bucket_id[74], 2); + assert_eq!(bucket_id[75], 3); + assert_eq!(bucket_id[99], 3); + assert_eq!(bucket_id[100], 4); + assert_eq!(bucket_id[127], 4); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA"] +fn bucket_iqr_kernel_computes_q1_q3_per_bucket() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "bucket_iqr_kernel")?; + + // tau values [1.0, 2.0, ..., 128.0] in sorted order (already sorted ascending). + // sorted_indices[p] = p; bucket k holds taus from index k * 25 to (k+1)*25 (or 28 for k=4). + let tau_host: Vec = (1..=HIDDEN_DIM).map(|x| x as f32).collect(); + let sorted_indices_host: Vec = (0..HIDDEN_DIM as u32).collect(); + let tau_d = upload(&stream, &tau_host)?; + let sorted_indices_d = upload(&stream, &sorted_indices_host)?; + let mut q1_d = stream.alloc_zeros::(N_HORIZONS)?; + let mut q3_d = stream.alloc_zeros::(N_HORIZONS)?; + + let cfg = LaunchConfig { grid_dim: (N_HORIZONS as u32,1,1), block_dim: (32,1,1), shared_mem_bytes: 32 * 4 }; + let mut launch = stream.launch_builder(&func); + launch.arg(&tau_d).arg(&sorted_indices_d).arg(&mut q1_d).arg(&mut q3_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let q1 = download(&stream, &q1_d)?; + let q3 = download(&stream, &q3_d)?; + // Bucket 0: tau = [1..25], Q1 ≈ 7, Q3 ≈ 19 (within 1 element of true quartile) + assert!((q1[0] - 7.0).abs() <= 1.0, "q1[0]={}", q1[0]); + assert!((q3[0] - 19.0).abs() <= 1.0, "q3[0]={}", q3[0]); + // Bucket 4: tau = [101..128], Q1 ≈ 108, Q3 ≈ 121 + assert!((q1[4] - 108.0).abs() <= 1.0, "q1[4]={}", q1[4]); + assert!((q3[4] - 121.0).abs() <= 1.0, "q3[4]={}", q3[4]); + Ok(()) +} + +#[test] +#[ignore = "requires CUDA"] +fn tau_reorder_kernel_places_taus_in_bucket_order() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "tau_reorder_kernel")?; + + // tau[c] = c, but bucket_id arbitrarily reassigns (here identity for simplicity). + let tau_host: Vec = (0..HIDDEN_DIM).map(|c| c as f32).collect(); + let bucket_id_host: Vec = (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect(); + let bucket_offset_host: Vec = vec![0, 25, 50, 75, 100, 128]; + let tau_d = upload(&stream, &tau_host)?; + let bucket_id_d = upload(&stream, &bucket_id_host)?; + let bucket_offset_d = upload(&stream, &bucket_offset_host)?; + let mut tau_all_d = stream.alloc_zeros::(HIDDEN_DIM)?; + + let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (HIDDEN_DIM as u32,1,1), shared_mem_bytes: (N_HORIZONS as u32 + 1) * 4 }; + let mut launch = stream.launch_builder(&func); + launch.arg(&tau_d).arg(&bucket_id_d).arg(&bucket_offset_d).arg(&mut tau_all_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let tau_all = download(&stream, &tau_all_d)?; + // With identity bucket_id, tau_all should match tau_host. + for c in 0..HIDDEN_DIM { + assert!((tau_all[c] - tau_host[c]).abs() < 1e-6, "tau_all[{}]={}", c, tau_all[c]); + } + Ok(()) +} + +#[test] +#[ignore = "requires CUDA"] +fn heads_compact_kernel_reorders_w_skip() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let func = load_kernel(&stream, "heads_compact_kernel")?; + + // Original heads_w_skip is [N_HORIZONS × HIDDEN_DIM] = 640 floats. + // For test: w_skip[h, c] = h * 1000 + c. + let w_skip_orig_host: Vec = (0..N_HORIZONS).flat_map(|h| { + (0..HIDDEN_DIM).map(move |c| (h * 1000 + c) as f32) + }).collect(); + let bucket_id_host: Vec = (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect(); + let bucket_offset_host: Vec = vec![0, 25, 50, 75, 100, 128]; + let w_skip_orig_d = upload(&stream, &w_skip_orig_host)?; + let bucket_id_d = upload(&stream, &bucket_id_host)?; + let bucket_offset_d = upload(&stream, &bucket_offset_host)?; + let mut w_skip_compact_d = stream.alloc_zeros::(HIDDEN_DIM)?; + + let cfg = LaunchConfig { grid_dim: (1,1,1), block_dim: (HIDDEN_DIM as u32,1,1), shared_mem_bytes: (N_HORIZONS as u32 + 1) * 4 }; + let mut launch = stream.launch_builder(&func); + launch.arg(&w_skip_orig_d).arg(&bucket_id_d).arg(&bucket_offset_d).arg(&mut w_skip_compact_d); + unsafe { launch.launch(cfg)?; } + stream.synchronize()?; + + let w_skip_compact = download(&stream, &w_skip_compact_d)?; + // For each channel c in bucket k, compact[c] = orig[k * HIDDEN_DIM + c] + for c in 0..HIDDEN_DIM { + let k = (c / 25).min(4); + let expected = (k * 1000 + c) as f32; + assert!((w_skip_compact[c] - expected).abs() < 1e-6, "compact[{}]={} expected {}", c, w_skip_compact[c], expected); + } + Ok(()) +}