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

feat(per-horizon-cfc): add bucket_transition_kernels.cu (5 device kernels)

Browse Source 
Per spec §2.3. All-on-device transition: tau_sort (bitonic-merge),
bucket_assign (static quintile boundaries [0,25,50,75,100,128]),
bucket_iqr (warp-reduction Q1/Q3 per bucket), tau_reorder, heads_compact.
No atomicAdd. No DtoH.

GPU oracle tests pass on RTX 3050 sm_86.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-21 15:46:56 +02:00
parent f13d6c6dcf
commit 1f700f9bfa
3 changed files with 378 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/ml-alpha/build.rs
View File

@@ -24,6 +24,7 @@ const KERNELS: &[&str] = &[
"output_smoothness", // CRT.train: per-horizon adjacent-position prob-jitter penalty "output_smoothness", // CRT.train: per-horizon adjacent-position prob-jitter penalty
"smoothness_lambda_controller", // CRT.train: ISV-driven λ controller anchored on h30 jitter "smoothness_lambda_controller", // CRT.train: ISV-driven λ controller anchored on h30 jitter
"gpu_log_ring", // GPU diagnostic log ring — tick kernel + log_record helper "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. // Cache bust v15 (2026-05-21): gpu_log_helpers.cuh extracted + smoothness_lambda_controller emits records.

193
crates/ml-alpha/cuda/bucket_transition_kernels.cu Normal file
View File

@@ -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();
}
}

184
crates/ml-alpha/tests/bucket_transition_kernels.rs Normal file
View File

@@ -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<CudaStream>, fn_name: &str) -> Result<cudarc::driver::CudaFunction> {
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<T: cudarc::driver::DeviceRepr + cudarc::driver::ValidAsZeroBits + Copy>(stream: &std::sync::Arc<CudaStream>, host: &[T]) -> Result<CudaSlice<T>> {
let mut d = stream.alloc_zeros::<T>(host.len())?;
stream.memcpy_htod(host, &mut d)?;
Ok(d)
}
fn download<T: cudarc::driver::DeviceRepr + Default + Copy>(stream: &std::sync::Arc<CudaStream>, d: &CudaSlice<T>) -> Result<Vec<T>> {
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<f32> = (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::<u32>(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<u32> = (0..HIDDEN_DIM as u32).collect();
let sorted_indices_d = upload(&stream, &sorted_indices_host)?;
let mut bucket_id_d = stream.alloc_zeros::<u8>(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<f32> = (1..=HIDDEN_DIM).map(|x| x as f32).collect();
let sorted_indices_host: Vec<u32> = (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::<f32>(N_HORIZONS)?;
let mut q3_d = stream.alloc_zeros::<f32>(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<f32> = (0..HIDDEN_DIM).map(|c| c as f32).collect();
let bucket_id_host: Vec<u8> = (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
let bucket_offset_host: Vec<u32> = 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::<f32>(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<f32> = (0..N_HORIZONS).flat_map(|h| {
(0..HIDDEN_DIM).map(move |c| (h * 1000 + c) as f32)
}).collect();
let bucket_id_host: Vec<u8> = (0..HIDDEN_DIM).map(|c| ((c / 25).min(4)) as u8).collect();
let bucket_offset_host: Vec<u32> = 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::<f32>(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(())
}