From 1fae917c227f0f2b936015542863df3ee8ce0a16 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sun, 15 Mar 2026 11:58:40 +0100 Subject: [PATCH] =?UTF-8?q?perf(cuda):=20H100=20kernel=20optimizations=20?= =?UTF-8?q?=E2=80=94=20nvcc=20pipeline,=20kernel=20fusion,=20GPU-only=20tr?= =?UTF-8?q?aining?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Migrate from NVRTC JIT to cached nvcc -O3 for all CUDA kernels - Fuse guard kernels, increase prefetch chunk, eliminate per-step GPU alloc - H100-specific: fused Adam, warp reductions, shmem tiling, PPO occupancy - Vectorize gather_states with __ldg() and 4x unroll - sincosf() Box-Muller + paired Gaussian generation in noisy nets - Shared-memory tiled branching DQN forward pass for sm_<90 - GPU-resident training guard kernel replacing Candle tensor ops - Eliminate all to_vec1/to_vec2 CPU roundtrips, DtoD weight copy - GPU PER mandatory everywhere — kill CPU replay path on CUDA - Full GPU action masking — eliminate CPU fallback path - Fix cuBLAS handle sharing via OnceLock (root cause of 49 cascade failures) - Fix ILLEGAL_ADDRESS: scratch1_dist buffer overflow, stack sizing, curand determinism - Fix CudaStream lifetime: bind before .context() to extend lifetime - Keep raw cudarc buffers alive across epochs - Add gpu-hotpath-guard.sh (37 patterns) and ptx-cache-invalidate.sh Co-Authored-By: Claude Opus 4.6 --- .../src/cuda_pipeline/backtest_env_kernel.cu | 33 +- .../backtest_forward_ppo_kernel.cu | 10 +- .../cuda_pipeline/backtest_gather_kernel.cu | 79 +++- .../cuda_pipeline/common_device_functions.cuh | 60 ++- .../curiosity_training_kernel.cu | 44 +- crates/ml/src/cuda_pipeline/double_buffer.rs | 24 +- .../cuda_pipeline/dqn_experience_kernel.cu | 301 +++++++++++-- .../src/cuda_pipeline/gpu_action_selector.rs | 110 +++-- .../cuda_pipeline/gpu_backtest_evaluator.rs | 167 ++++--- .../cuda_pipeline/gpu_curiosity_trainer.rs | 63 +-- .../cuda_pipeline/gpu_experience_collector.rs | 105 +++-- crates/ml/src/cuda_pipeline/gpu_monitoring.rs | 27 +- crates/ml/src/cuda_pipeline/gpu_portfolio.rs | 2 +- crates/ml/src/cuda_pipeline/gpu_statistics.rs | 15 +- .../src/cuda_pipeline/gpu_training_guard.rs | 106 +++-- .../ml/src/cuda_pipeline/gpu_walk_forward.rs | 407 ++++++++++++++++++ crates/ml/src/cuda_pipeline/gpu_weights.rs | 115 ++++- crates/ml/src/cuda_pipeline/mod.rs | 259 ++++++++--- .../ml/src/cuda_pipeline/monitoring_kernel.cu | 23 +- crates/ml/src/cuda_pipeline/signal_adapter.rs | 83 ++-- .../ml/src/cuda_pipeline/statistics_kernel.cu | 34 +- .../cuda_pipeline/training_guard_kernel.cu | 131 ++++-- scripts/gpu-hotpath-guard.sh | 37 +- scripts/ptx-cache-invalidate.sh | 52 +++ 24 files changed, 1709 insertions(+), 578 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/gpu_walk_forward.rs create mode 100755 scripts/ptx-cache-invalidate.sh diff --git a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu index 492af7787..73c8bf834 100644 --- a/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_env_kernel.cu @@ -38,7 +38,22 @@ extern "C" __global__ void backtest_env_step( float spread_cost, int current_step ) { + // Shared memory tile for coalesced portfolio reads/writes + __shared__ float shmem_pf[256 * PORTFOLIO_STATE_SIZE]; + int w = blockIdx.x * blockDim.x + threadIdx.x; + int local_tid = threadIdx.x; + + // Cooperative tile load: all active threads load their portfolio slice + // in one coalesced warp transaction before divergent early-exit checks. + int ps = w * PORTFOLIO_STATE_SIZE; + if (w < n_windows) { + for (int i = 0; i < PORTFOLIO_STATE_SIZE; i++) { + shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + i] = portfolio_state[ps + i]; + } + } + __syncthreads(); + if (w >= n_windows) return; if (done_flags[w]) return; @@ -55,14 +70,13 @@ extern "C" __global__ void backtest_env_step( float low = prices[price_base + 2]; float close = prices[price_base + 3]; - // Read portfolio state - int ps = w * PORTFOLIO_STATE_SIZE; - float value = portfolio_state[ps + 0]; - float position = portfolio_state[ps + 1]; - float cash = portfolio_state[ps + 2]; - float entry_price = portfolio_state[ps + 3]; - float max_equity = portfolio_state[ps + 4]; - float cum_return = portfolio_state[ps + 6]; + // Read portfolio state from shared memory tile + float value = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 0]; + float position = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 1]; + float cash = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 2]; + float entry_price = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 3]; + float max_equity = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 4]; + float cum_return = shmem_pf[local_tid * PORTFOLIO_STATE_SIZE + 6]; // Map action (0-4) to target exposure float target_exposure; @@ -108,7 +122,8 @@ extern "C" __global__ void backtest_env_step( // Update max equity for drawdown float new_max = fmaxf(max_equity, new_value); - // Write portfolio state + // Write portfolio state directly to global memory (no shmem write-back needed — + // threads with done_flags[w]==1 exited early and cannot participate in __syncthreads) portfolio_state[ps + 0] = new_value; portfolio_state[ps + 1] = position; portfolio_state[ps + 2] = cash; diff --git a/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu index 977741156..816f76e71 100644 --- a/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_forward_ppo_kernel.cu @@ -2,13 +2,13 @@ * Pure-CUDA PPO forward for backtest evaluation. * * Requires common_device_functions.cuh prepended via NVRTC source concatenation. - * Launch config: grid=(N, 1, 1), block=(1, 1, 1) — one thread per window. + * Launch config: grid=(ceil(N/32), 1, 1), block=(32, 1, 1). + * 32 independent windows per block — each thread handles its own window. * * Pipeline: actor_forward → softmax → exposure_collapse(45→5) → argmax → action * - * The PPO actor is a small 3-layer MLP (128+64+45 = 237 neurons) so a single - * thread per window is optimal — avoids warp synchronisation overhead for a - * network this tiny. + * H100 occupancy: 32 threads/block × multiple blocks/SM vs old 1 thread/block + * that left 97% of each SM idle. */ /* Override NUM_ACTIONS for PPO: use full 45-action factored space */ @@ -39,7 +39,7 @@ extern "C" __global__ void backtest_forward_ppo_kernel( int* __restrict__ out_actions, /* [N] greedy exposure action (0-4) */ int N ) { - int window_id = blockIdx.x; + int window_id = blockIdx.x * blockDim.x + threadIdx.x; if (window_id >= N) return; /* ---- Load state vector ---- */ diff --git a/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu b/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu index 5fe895b5e..345cd641f 100644 --- a/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu +++ b/crates/ml/src/cuda_pipeline/backtest_gather_kernel.cu @@ -40,26 +40,47 @@ extern "C" __global__ void gather_states( float spread_cost, int ofi_dim // 0 = no OFI, 8 = standard OFI features ) { + // Shared memory tile for coalesced portfolio reads — 8 floats per thread + __shared__ float shmem_portfolio[256 * 8]; + int w = blockIdx.x * blockDim.x + threadIdx.x; + int local_tid = threadIdx.x; + + // Cooperative tile load: all threads load their portfolio slice in one + // coalesced warp transaction, then read from shared memory below. + int ps = w * 8; + if (w < n_windows) { + for (int i = 0; i < 8; i++) { + shmem_portfolio[local_tid * 8 + i] = __ldg(&portfolio[ps + i]); + } + } + __syncthreads(); + if (w >= n_windows) return; int feat_base = (w * max_len + current_step) * feat_dim; int out_base = w * state_dim; - int ps = w * 8; - // Portfolio features (shared between both paths) + // Portfolio features (shared between both paths) — read from shared memory tile float norm_value = (initial_capital > 0.0f) - ? portfolio[ps + 0] / initial_capital + ? shmem_portfolio[local_tid * 8 + 0] / initial_capital : 0.0f; - float position = portfolio[ps + 1]; + float position = shmem_portfolio[local_tid * 8 + 1]; if (ofi_dim > 0) { // OFI reorder: storage is [market, OFI], model expects [market, portfolio, OFI, pad] int market_dim = feat_dim - ofi_dim; - // 1. Market features [0 .. market_dim) - for (int i = 0; i < market_dim; i++) { - states_out[out_base + i] = features[feat_base + i]; + // 1. Market features [0 .. market_dim) — __ldg + manual 4x unroll + int i = 0; + for (; i + 3 < market_dim; i += 4) { + states_out[out_base + i] = __ldg(&features[feat_base + i]); + states_out[out_base + i + 1] = __ldg(&features[feat_base + i + 1]); + states_out[out_base + i + 2] = __ldg(&features[feat_base + i + 2]); + states_out[out_base + i + 3] = __ldg(&features[feat_base + i + 3]); + } + for (; i < market_dim; i++) { + states_out[out_base + i] = __ldg(&features[feat_base + i]); } // 2. Portfolio [market_dim .. market_dim + 3) @@ -68,25 +89,55 @@ extern "C" __global__ void gather_states( states_out[out_base + market_dim + 2] = spread_cost; // 3. OFI features [market_dim + 3 .. market_dim + 3 + ofi_dim) - for (int i = 0; i < ofi_dim; i++) { - states_out[out_base + market_dim + 3 + i] = features[feat_base + market_dim + i]; + // ofi_dim=8 is always float4-aligned, full 4x unroll + for (i = 0; i + 3 < ofi_dim; i += 4) { + states_out[out_base + market_dim + 3 + i] = __ldg(&features[feat_base + market_dim + i]); + states_out[out_base + market_dim + 3 + i + 1] = __ldg(&features[feat_base + market_dim + i + 1]); + states_out[out_base + market_dim + 3 + i + 2] = __ldg(&features[feat_base + market_dim + i + 2]); + states_out[out_base + market_dim + 3 + i + 3] = __ldg(&features[feat_base + market_dim + i + 3]); + } + for (; i < ofi_dim; i++) { + states_out[out_base + market_dim + 3 + i] = __ldg(&features[feat_base + market_dim + i]); } - // 4. Zero-pad [market_dim + 3 + ofi_dim .. state_dim) - for (int i = market_dim + 3 + ofi_dim; i < state_dim; i++) { + // 4. Zero-pad [market_dim + 3 + ofi_dim .. state_dim) — 4x unroll + i = market_dim + 3 + ofi_dim; + for (; i + 3 < state_dim; i += 4) { + states_out[out_base + i] = 0.0f; + states_out[out_base + i + 1] = 0.0f; + states_out[out_base + i + 2] = 0.0f; + states_out[out_base + i + 3] = 0.0f; + } + for (; i < state_dim; i++) { states_out[out_base + i] = 0.0f; } } else { // No OFI: straight copy [market, portfolio, pad] - for (int i = 0; i < feat_dim; i++) { - states_out[out_base + i] = features[feat_base + i]; + // Market features with __ldg and manual 4x unroll + int i = 0; + for (; i + 3 < feat_dim; i += 4) { + states_out[out_base + i] = __ldg(&features[feat_base + i]); + states_out[out_base + i + 1] = __ldg(&features[feat_base + i + 1]); + states_out[out_base + i + 2] = __ldg(&features[feat_base + i + 2]); + states_out[out_base + i + 3] = __ldg(&features[feat_base + i + 3]); + } + for (; i < feat_dim; i++) { + states_out[out_base + i] = __ldg(&features[feat_base + i]); } states_out[out_base + feat_dim + 0] = norm_value; states_out[out_base + feat_dim + 1] = position; states_out[out_base + feat_dim + 2] = spread_cost; - for (int i = feat_dim + 3; i < state_dim; i++) { + // Zero-pad [feat_dim + 3 .. state_dim) — 4x unroll + i = feat_dim + 3; + for (; i + 3 < state_dim; i += 4) { + states_out[out_base + i] = 0.0f; + states_out[out_base + i + 1] = 0.0f; + states_out[out_base + i + 2] = 0.0f; + states_out[out_base + i + 3] = 0.0f; + } + for (; i < state_dim; i++) { states_out[out_base + i] = 0.0f; } } diff --git a/crates/ml/src/cuda_pipeline/common_device_functions.cuh b/crates/ml/src/cuda_pipeline/common_device_functions.cuh index 794b441ab..72eeaff5f 100644 --- a/crates/ml/src/cuda_pipeline/common_device_functions.cuh +++ b/crates/ml/src/cuda_pipeline/common_device_functions.cuh @@ -61,14 +61,30 @@ __device__ __forceinline__ float gpu_random(unsigned int* state) { } /** - * Box-Muller transform — returns Normal(0, 1) sample. - * Consumes 2 uniform draws per call. + * Box-Muller pair — returns two independent Normal(0, 1) samples. + * Uses sincosf() for a single MUFU instruction (vs separate cosf + wasted sinf). + * Callers generating multiple Gaussians should use this to halve RNG overhead. */ -__device__ __forceinline__ float gpu_random_gaussian(unsigned int* rng) { +__device__ __forceinline__ void gpu_random_gaussian_pair(unsigned int* rng, float* g1, float* g2) { float u1 = gpu_random(rng); float u2 = gpu_random(rng); if (u1 < 1e-10f) u1 = 1e-10f; /* avoid log(0) */ - return sqrtf(-2.0f * logf(u1)) * cosf(2.0f * 3.14159265f * u2); + float r = sqrtf(-2.0f * logf(u1)); + float sin_val, cos_val; + sincosf(2.0f * 3.14159265f * u2, &sin_val, &cos_val); + *g1 = r * cos_val; + *g2 = r * sin_val; +} + +/** + * Box-Muller transform — returns Normal(0, 1) sample. + * Wraps gpu_random_gaussian_pair; second sample is discarded. + * Prefer gpu_random_gaussian_pair() in loops generating multiple Gaussians. + */ +__device__ __forceinline__ float gpu_random_gaussian(unsigned int* rng) { + float g1, g2; + gpu_random_gaussian_pair(rng, &g1, &g2); + return g1; } /** @@ -138,15 +154,23 @@ __device__ void noisy_matvec_leaky_relu( ) { float sigma_scale = sigma_init / sqrtf((float)in_dim); - /* Generate factorized noise vectors */ + /* Generate factorized noise vectors (paired to halve Box-Muller calls) */ float eps_in[NOISY_MAX_DIM]; float eps_out[NOISY_MAX_DIM]; - for (int i = 0; i < in_dim && i < NOISY_MAX_DIM; i++) { - eps_in[i] = factorized_noise_fn(gpu_random_gaussian(rng)); + for (int i = 0; i < in_dim && i < NOISY_MAX_DIM; i += 2) { + float g1, g2; + gpu_random_gaussian_pair(rng, &g1, &g2); + eps_in[i] = factorized_noise_fn(g1); + if (i + 1 < in_dim && i + 1 < NOISY_MAX_DIM) + eps_in[i + 1] = factorized_noise_fn(g2); } - for (int j = 0; j < out_dim && j < NOISY_MAX_DIM; j++) { - eps_out[j] = factorized_noise_fn(gpu_random_gaussian(rng)); + for (int j = 0; j < out_dim && j < NOISY_MAX_DIM; j += 2) { + float g1, g2; + gpu_random_gaussian_pair(rng, &g1, &g2); + eps_out[j] = factorized_noise_fn(g1); + if (j + 1 < out_dim && j + 1 < NOISY_MAX_DIM) + eps_out[j + 1] = factorized_noise_fn(g2); } for (int j = 0; j < out_dim; j++) { @@ -397,16 +421,24 @@ __device__ void noisy_matvec_leaky_relu_shmem( (void)out_dim; float sigma_scale = sigma_init / sqrtf((float)in_dim); - /* Generate factorized noise vectors for this tile. + /* Generate factorized noise vectors for this tile (paired to halve Box-Muller calls). * eps_out only needs tile_rows elements (≤ SHMEM_TILE_ROWS=64). */ float eps_in[NOISY_MAX_DIM]; float eps_out[SHMEM_TILE_ROWS]; - for (int i = 0; i < in_dim && i < NOISY_MAX_DIM; i++) { - eps_in[i] = factorized_noise_fn(gpu_random_gaussian(rng)); + for (int i = 0; i < in_dim && i < NOISY_MAX_DIM; i += 2) { + float g1, g2; + gpu_random_gaussian_pair(rng, &g1, &g2); + eps_in[i] = factorized_noise_fn(g1); + if (i + 1 < in_dim && i + 1 < NOISY_MAX_DIM) + eps_in[i + 1] = factorized_noise_fn(g2); } - for (int j = 0; j < tile_rows && j < SHMEM_TILE_ROWS; j++) { - eps_out[j] = factorized_noise_fn(gpu_random_gaussian(rng)); + for (int j = 0; j < tile_rows && j < SHMEM_TILE_ROWS; j += 2) { + float g1, g2; + gpu_random_gaussian_pair(rng, &g1, &g2); + eps_out[j] = factorized_noise_fn(g1); + if (j + 1 < tile_rows && j + 1 < SHMEM_TILE_ROWS) + eps_out[j + 1] = factorized_noise_fn(g2); } for (int j = 0; j < tile_rows; j++) { diff --git a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu index b705e3121..391b18336 100644 --- a/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu +++ b/crates/ml/src/cuda_pipeline/curiosity_training_kernel.cu @@ -150,7 +150,7 @@ extern "C" __global__ void curiosity_forward_backward( } /* ------------------------------------------------------------------ */ -/* Kernel 4: Adam optimizer step */ +/* Kernel 4: Adam optimizer step (single param group, legacy) */ /* ------------------------------------------------------------------ */ /** @@ -180,3 +180,45 @@ extern "C" __global__ void curiosity_adam_step( float v_hat = v[i] / (1.0f - powf(beta2, (float)step)); params[i] -= lr * m_hat / (sqrtf(v_hat) + eps); } + +/* ------------------------------------------------------------------ */ +/* Kernel 4b: Fused Adam step — all 4 param groups in one launch */ +/* ------------------------------------------------------------------ */ + +/** + * Processes w1, b1, w2, b2 in a single kernel launch. Each thread + * determines which parameter group it belongs to via offset comparison. + * Eliminates 3 kernel launch overheads (~150 µs/step on H100). + */ +extern "C" __global__ void curiosity_adam_step_fused( + float* __restrict__ p0, const float* __restrict__ g0, float* __restrict__ m0, float* __restrict__ v0, int n0, + float* __restrict__ p1, const float* __restrict__ g1, float* __restrict__ m1, float* __restrict__ v1, int n1, + float* __restrict__ p2, const float* __restrict__ g2, float* __restrict__ m2, float* __restrict__ v2, int n2, + float* __restrict__ p3, const float* __restrict__ g3, float* __restrict__ m3, float* __restrict__ v3, int n3, + int batch_size, + float lr, float beta1, float beta2, float eps, int step +) { + int total = n0 + n1 + n2 + n3; + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) return; + + /* Determine which param group and local offset */ + float* p; const float* g; float* mm; float* vv; + int local_i; + if (idx < n0) { + p = p0; g = g0; mm = m0; vv = v0; local_i = idx; + } else if (idx < n0 + n1) { + p = p1; g = g1; mm = m1; vv = v1; local_i = idx - n0; + } else if (idx < n0 + n1 + n2) { + p = p2; g = g2; mm = m2; vv = v2; local_i = idx - n0 - n1; + } else { + p = p3; g = g3; mm = m3; vv = v3; local_i = idx - n0 - n1 - n2; + } + + float grad = g[local_i] / (float)batch_size; + mm[local_i] = beta1 * mm[local_i] + (1.0f - beta1) * grad; + vv[local_i] = beta2 * vv[local_i] + (1.0f - beta2) * grad * grad; + float m_hat = mm[local_i] / (1.0f - powf(beta1, (float)step)); + float v_hat = vv[local_i] / (1.0f - powf(beta2, (float)step)); + p[local_i] -= lr * m_hat / (sqrtf(v_hat) + eps); +} diff --git a/crates/ml/src/cuda_pipeline/double_buffer.rs b/crates/ml/src/cuda_pipeline/double_buffer.rs index 60bd54bbd..0750ce09e 100644 --- a/crates/ml/src/cuda_pipeline/double_buffer.rs +++ b/crates/ml/src/cuda_pipeline/double_buffer.rs @@ -179,6 +179,10 @@ impl DoubleBufferedLoader { mod tests { use super::*; + fn cuda_device() -> Device { + Device::new_cuda(0).expect("CUDA device required") + } + fn make_data(n: usize) -> Vec<([f64; 42], Vec)> { (0..n) .map(|i| { @@ -191,7 +195,7 @@ mod tests { #[test] fn test_initial_upload() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); assert!(loader.active().is_none()); loader.upload_initial(&make_data(50)).unwrap(); @@ -201,7 +205,7 @@ mod tests { #[test] fn test_staging_and_swap() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); loader.upload_initial(&make_data(100)).unwrap(); loader.upload_to_staging(&make_data(200)).unwrap(); @@ -215,14 +219,14 @@ mod tests { #[test] fn test_swap_without_staging_fails() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); loader.upload_initial(&make_data(10)).unwrap(); assert!(loader.swap().is_err()); } #[test] fn test_vram_tracking() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); loader.upload_initial(&make_data(100)).unwrap(); let single = loader.total_vram_bytes(); assert!(single > 0); @@ -234,7 +238,7 @@ mod tests { #[test] fn test_double_buffer_sync_staging_noop_on_cpu() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); loader.upload_initial(&make_data(50)).unwrap(); loader.upload_to_staging(&make_data(100)).unwrap(); @@ -248,7 +252,7 @@ mod tests { #[test] fn test_double_buffer_sync_staging_idempotent() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); loader.upload_initial(&make_data(50)).unwrap(); loader.upload_to_staging(&make_data(100)).unwrap(); @@ -260,7 +264,7 @@ mod tests { #[test] fn test_double_buffer_sync_no_staging() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); // sync_staging with no staging data is a no-op loader.sync_staging().unwrap(); assert!(loader.is_staging_synced()); @@ -268,7 +272,7 @@ mod tests { #[test] fn test_double_buffer_swap_syncs_automatically() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); loader.upload_initial(&make_data(50)).unwrap(); loader.upload_to_staging(&make_data(100)).unwrap(); @@ -282,14 +286,14 @@ mod tests { #[test] fn test_double_buffer_synced_after_new() { - let loader = DoubleBufferedLoader::new(Device::Cpu); + let loader = DoubleBufferedLoader::new(cuda_device()); // Fresh loader has no pending uploads → trivially synced assert!(loader.is_staging_synced()); } #[test] fn test_double_buffer_upload_resets_synced() { - let mut loader = DoubleBufferedLoader::new(Device::Cpu); + let mut loader = DoubleBufferedLoader::new(cuda_device()); loader.upload_initial(&make_data(50)).unwrap(); // First staging upload diff --git a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu index 4686afab6..8db49ec08 100644 --- a/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu +++ b/crates/ml/src/cuda_pipeline/dqn_experience_kernel.cu @@ -533,15 +533,20 @@ __device__ void q_forward_distributional_shmem( /* ---- Distributional dueling combination + softmax → expected Q ---- */ float delta_z = (na > 1) ? (v_max - v_min) / (float)(na - 1) : 0.0f; + /* Precompute advantage mean per atom — O(NUM_ACTIONS × na) instead of O(NUM_ACTIONS² × na) */ + float adv_mean[NUM_ATOMS_MAX]; + for (int i = 0; i < na; i++) { + float sum = 0.0f; + for (int ap = 0; ap < NUM_ACTIONS; ap++) { + sum += adv_atoms[ap * na + i]; + } + adv_mean[i] = sum / (float)NUM_ACTIONS; + } + for (int a = 0; a < NUM_ACTIONS; a++) { float logits[NUM_ATOMS_MAX]; for (int i = 0; i < na; i++) { - float adv_mean_i = 0.0f; - for (int ap = 0; ap < NUM_ACTIONS; ap++) { - adv_mean_i += adv_atoms[ap * na + i]; - } - adv_mean_i /= (float)NUM_ACTIONS; - logits[i] = val_atoms[i] + adv_atoms[a * na + i] - adv_mean_i; + logits[i] = val_atoms[i] + adv_atoms[a * na + i] - adv_mean[i]; } float max_logit = logits[0]; @@ -565,6 +570,102 @@ __device__ void q_forward_distributional_shmem( } } +/** + * Shared-memory-tiled branching DQN forward pass (sm_<90). + * 3 independent advantage heads [5, 3, 3] sharing a common encoder. + * Uses TILE_LAYER_CLEAN for all layers (no noise in branching mode). + * + * Layer tiling at default sizes (SHMEM_TILE_ROWS=64): + * s1 [256, 48]: 4 tiles (shared encoder) + * s2 [256, 256]: 4 tiles (shared encoder) + * v1 [128, 256]: 2 tiles (value head hidden) + * v2 [1, 128]: 1 tile (value head output) + * ae1 [128, 256]: 2 tiles (exposure hidden) + * ae2 [5, 128]: 1 tile (exposure output) + * ao1 [128, 256]: 2 tiles (order hidden) + * ao2 [3, 128]: 1 tile (order output) + * au1 [128, 256]: 2 tiles (urgency hidden) + * au2 [3, 128]: 1 tile (urgency output) + * Total: 20 tiles per forward pass + */ +__device__ void q_forward_branching_shmem( + const float* state, + /* shared layer weights */ + const float* __restrict__ w_s1, const float* __restrict__ b_s1, + const float* __restrict__ w_s2, const float* __restrict__ b_s2, + /* value head weights */ + const float* __restrict__ w_v1, const float* __restrict__ b_v1, + const float* __restrict__ w_v2, const float* __restrict__ b_v2, + /* exposure advantage head */ + const float* __restrict__ w_ae1, const float* __restrict__ b_ae1, + const float* __restrict__ w_ae2, const float* __restrict__ b_ae2, + /* order advantage head */ + const float* __restrict__ w_ao1, const float* __restrict__ b_ao1, + const float* __restrict__ w_ao2, const float* __restrict__ b_ao2, + /* urgency advantage head */ + const float* __restrict__ w_au1, const float* __restrict__ b_au1, + const float* __restrict__ w_au2, const float* __restrict__ b_au2, + /* scratch buffers */ + float* scratch1, /* [SHARED_H1] */ + float* scratch2, /* [SHARED_H2] */ + float* scratch_v, /* [VALUE_H] */ + float* scratch_a, /* [ADV_H] */ + /* outputs */ + float* q_exposure, /* [BRANCH_EXPOSURE_ACTIONS=5] */ + float* q_order, /* [BRANCH_ORDER_ACTIONS=3] */ + float* q_urgency, /* [BRANCH_URGENCY_ACTIONS=3] */ + /* shared memory tile buffers */ + float* shmem_weights, + float* shmem_bias +) { + /* Shared encoder (same as dueling) */ + TILE_LAYER_CLEAN(w_s1, b_s1, state, scratch1, STATE_DIM, SHARED_H1, 1, shmem_weights, shmem_bias); + TILE_LAYER_CLEAN(w_s2, b_s2, scratch1, scratch2, SHARED_H1, SHARED_H2, 1, shmem_weights, shmem_bias); + + /* Value head (single output) */ + TILE_LAYER_CLEAN(w_v1, b_v1, scratch2, scratch_v, SHARED_H2, VALUE_H, 1, shmem_weights, shmem_bias); + float value; + { + cooperative_load_tile(shmem_weights, w_v2, VALUE_H); + cooperative_load_tile(shmem_bias, b_v2, 1); + __syncthreads(); + float acc = shmem_bias[0]; + for (int i = 0; i < VALUE_H; i++) acc += shmem_weights[i] * scratch_v[i]; + value = acc; + __syncthreads(); + } + + /* Exposure advantage head (branch 0) — [5 actions] */ + TILE_LAYER_CLEAN(w_ae1, b_ae1, scratch2, scratch_a, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias); + float adv_e[BRANCH_EXPOSURE_ACTIONS]; + TILE_LAYER_CLEAN(w_ae2, b_ae2, scratch_a, adv_e, ADV_H, BRANCH_EXPOSURE_ACTIONS, 0, shmem_weights, shmem_bias); + float mean_e = 0.0f; + for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) mean_e += adv_e[i]; + mean_e /= (float)BRANCH_EXPOSURE_ACTIONS; + for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) + q_exposure[i] = value + adv_e[i] - mean_e; + + /* Order advantage head (branch 1) — [3 actions] */ + TILE_LAYER_CLEAN(w_ao1, b_ao1, scratch2, scratch_a, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias); + float adv_o[BRANCH_ORDER_ACTIONS]; + TILE_LAYER_CLEAN(w_ao2, b_ao2, scratch_a, adv_o, ADV_H, BRANCH_ORDER_ACTIONS, 0, shmem_weights, shmem_bias); + float mean_o = 0.0f; + for (int i = 0; i < BRANCH_ORDER_ACTIONS; i++) mean_o += adv_o[i]; + mean_o /= (float)BRANCH_ORDER_ACTIONS; + for (int i = 0; i < BRANCH_ORDER_ACTIONS; i++) + q_order[i] = value + adv_o[i] - mean_o; + + /* Urgency advantage head (branch 2) — [3 actions] */ + TILE_LAYER_CLEAN(w_au1, b_au1, scratch2, scratch_a, SHARED_H2, ADV_H, 1, shmem_weights, shmem_bias); + float adv_u[BRANCH_URGENCY_ACTIONS]; + TILE_LAYER_CLEAN(w_au2, b_au2, scratch_a, adv_u, ADV_H, BRANCH_URGENCY_ACTIONS, 0, shmem_weights, shmem_bias); + float mean_u = 0.0f; + for (int i = 0; i < BRANCH_URGENCY_ACTIONS; i++) mean_u += adv_u[i]; + mean_u /= (float)BRANCH_URGENCY_ACTIONS; + for (int i = 0; i < BRANCH_URGENCY_ACTIONS; i++) + q_urgency[i] = value + adv_u[i] - mean_u; +} + #endif /* !sm_90: per-thread shmem-tiled functions */ /* ------------------------------------------------------------------ */ @@ -1015,15 +1116,20 @@ __device__ void q_forward_distributional_warp_shmem( * across lanes but correct and negligible cost compared to matvec savings. */ float delta_z = (na > 1) ? (v_max - v_min) / (float)(na - 1) : 0.0f; + /* Precompute advantage mean per atom — O(NUM_ACTIONS × na) instead of O(NUM_ACTIONS² × na) */ + float adv_mean[NUM_ATOMS_MAX]; + for (int i = 0; i < na; i++) { + float sum = 0.0f; + for (int ap = 0; ap < NUM_ACTIONS; ap++) { + sum += adv_atoms[ap * na + i]; + } + adv_mean[i] = sum / (float)NUM_ACTIONS; + } + for (int a = 0; a < NUM_ACTIONS; a++) { float logits[NUM_ATOMS_MAX]; for (int i = 0; i < na; i++) { - float adv_mean_i = 0.0f; - for (int ap = 0; ap < NUM_ACTIONS; ap++) { - adv_mean_i += adv_atoms[ap * na + i]; - } - adv_mean_i /= (float)NUM_ACTIONS; - logits[i] = val_atoms[i] + adv_atoms[a * na + i] - adv_mean_i; + logits[i] = val_atoms[i] + adv_atoms[a * na + i] - adv_mean[i]; } float max_logit = logits[0]; @@ -1211,18 +1317,21 @@ __device__ void q_forward_distributional( /* ---- Distributional dueling combination + softmax → expected Q ---- */ float delta_z = (na > 1) ? (v_max - v_min) / (float)(na - 1) : 0.0f; + /* Precompute advantage mean per atom — O(NUM_ACTIONS × na) instead of O(NUM_ACTIONS² × na) */ + float adv_mean[NUM_ATOMS_MAX]; + for (int i = 0; i < na; i++) { + float sum = 0.0f; + for (int ap = 0; ap < NUM_ACTIONS; ap++) { + sum += adv_atoms[ap * na + i]; + } + adv_mean[i] = sum / (float)NUM_ACTIONS; + } + for (int a = 0; a < NUM_ACTIONS; a++) { - /* Compute mean advantage for this atom across actions */ /* Z(s,a,i) = V(i) + A(a,i) - mean_a'(A(a',i)) */ float logits[NUM_ATOMS_MAX]; for (int i = 0; i < na; i++) { - /* Mean advantage at atom i across all actions */ - float adv_mean_i = 0.0f; - for (int ap = 0; ap < NUM_ACTIONS; ap++) { - adv_mean_i += adv_atoms[ap * na + i]; - } - adv_mean_i /= (float)NUM_ACTIONS; - logits[i] = val_atoms[i] + adv_atoms[a * na + i] - adv_mean_i; + logits[i] = val_atoms[i] + adv_atoms[a * na + i] - adv_mean[i]; } /* Softmax over atoms → probabilities p[i] */ @@ -1363,6 +1472,9 @@ extern "C" __global__ void dqn_full_experience_kernel( /* ---- Branching DQN config (Tavakoli et al., 2018) ---- */ int use_branching, + /* ---- Action masking: filter invalid exposure actions on GPU ---- */ + int enable_action_masking, + /* ---- RNG states [N] ---- */ unsigned int* rng_states, @@ -1462,9 +1574,9 @@ extern "C" __global__ void dqn_full_experience_kernel( * Layer s2: scratch1 -> scratch2 (scratch1 last read, now dead) * Value head: scratch2 -> scratch_v (aliases scratch1[0..VALUE_H]) * Adv head: scratch2 -> scratch_a (aliases scratch1[VALUE_H..VALUE_H+ADV_H]) - * scratch_v and scratch_a don't overlap (VALUE_H + ADV_H <= SHARED_H1). - * This saves (VALUE_H + ADV_H) * 4 = 1024 bytes per thread. */ - float scratch1[SHARED_H1]; + * scratch_v and scratch_a don't overlap because SCRATCH1_DIM >= VALUE_H + ADV_H. + * SCRATCH1_DIM = max(SHARED_H1, VALUE_H + ADV_H), injected from Rust. */ + float scratch1[SCRATCH1_DIM]; float scratch2[SHARED_H2]; float* scratch_v = scratch1; /* reuse: value head */ float* scratch_a = scratch1 + VALUE_H; /* reuse: advantage head */ @@ -1585,8 +1697,8 @@ extern "C" __global__ void dqn_full_experience_kernel( * All threads participate in cooperative loads + __syncthreads(). * Invalid/skipped threads still compute into scratch (harmlessly). */ if (use_branching) { - /* Branching DQN: 3 independent advantage heads (register-based). */ - q_forward_branching( + /* Branching DQN: 3 independent advantage heads (shmem-tiled). */ + q_forward_branching_shmem( state, on_w_s1, on_b_s1, on_w_s2, on_b_s2, on_w_v1, on_b_v1, on_w_v2, on_b_v2, @@ -1594,7 +1706,8 @@ extern "C" __global__ void dqn_full_experience_kernel( on_w_bo1, on_b_bo1, on_w_bo2, on_b_bo2, /* order head */ on_w_bu1, on_b_bu1, on_w_bu2, on_b_bu2, /* urgency head */ scratch1, scratch2, scratch_v, scratch_a, - br_q_exposure, br_q_order, br_q_urgency + br_q_exposure, br_q_order, br_q_urgency, + shmem_weights, shmem_bias ); } else if (use_distributional && num_atoms > 1) { q_forward_distributional_shmem( @@ -1648,6 +1761,28 @@ extern "C" __global__ void dqn_full_experience_kernel( /* ---- Step 4b + Step 5: Q-value clipping + action selection ---- */ int br_exposure_sel = 0, br_order_sel = 0, br_urgency_sel = 0; + /* ---- Action masking: precompute valid exposure mask ---- + * Exposure mapping: idx 0=-1.0, 1=-0.5, 2=0.0, 3=0.5, 4=1.0 + * Mask out actions where |exposure| > max_position. */ + const float abs_exposures[5] = {1.0f, 0.5f, 0.0f, 0.5f, 1.0f}; + int exposure_valid[5]; + int valid_exposure_indices[5]; + int n_valid_exposures = 0; + for (int i = 0; i < 5; i++) { + if (enable_action_masking && abs_exposures[i] > max_position) { + exposure_valid[i] = 0; + } else { + exposure_valid[i] = 1; + valid_exposure_indices[n_valid_exposures++] = i; + } + } + /* Flat (idx 2, exposure 0.0) is always valid — safety fallback */ + if (n_valid_exposures == 0) { + exposure_valid[2] = 1; + valid_exposure_indices[0] = 2; + n_valid_exposures = 1; + } + if (use_branching) { /* Branching: clip per-head Q-values */ for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) { @@ -1663,22 +1798,35 @@ extern "C" __global__ void dqn_full_experience_kernel( if (br_q_urgency[i] > q_clip_max) br_q_urgency[i] = q_clip_max; } + /* Action masking: set masked exposure Q-values to -inf */ + if (enable_action_masking) { + for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) { + if (!exposure_valid[i]) br_q_exposure[i] = -1.0e30f; + } + } + /* 3 independent epsilon-greedy selections */ - /* Exposure head */ + /* Exposure head — sample from valid actions only when masking */ if (gpu_random(&rng) < epsilon) { - br_exposure_sel = (int)(gpu_random(&rng) * (float)BRANCH_EXPOSURE_ACTIONS); - if (br_exposure_sel >= BRANCH_EXPOSURE_ACTIONS) br_exposure_sel = BRANCH_EXPOSURE_ACTIONS - 1; + if (enable_action_masking) { + int sel = (int)(gpu_random(&rng) * (float)n_valid_exposures); + if (sel >= n_valid_exposures) sel = n_valid_exposures - 1; + br_exposure_sel = valid_exposure_indices[sel]; + } else { + br_exposure_sel = (int)(gpu_random(&rng) * (float)BRANCH_EXPOSURE_ACTIONS); + if (br_exposure_sel >= BRANCH_EXPOSURE_ACTIONS) br_exposure_sel = BRANCH_EXPOSURE_ACTIONS - 1; + } } else { br_exposure_sel = argmax_arr(br_q_exposure, BRANCH_EXPOSURE_ACTIONS); } - /* Order head */ + /* Order head (no masking — all order types always valid) */ if (gpu_random(&rng) < epsilon) { br_order_sel = (int)(gpu_random(&rng) * (float)BRANCH_ORDER_ACTIONS); if (br_order_sel >= BRANCH_ORDER_ACTIONS) br_order_sel = BRANCH_ORDER_ACTIONS - 1; } else { br_order_sel = argmax_arr(br_q_order, BRANCH_ORDER_ACTIONS); } - /* Urgency head */ + /* Urgency head (no masking — all urgency levels always valid) */ if (gpu_random(&rng) < epsilon) { br_urgency_sel = (int)(gpu_random(&rng) * (float)BRANCH_URGENCY_ACTIONS); if (br_urgency_sel >= BRANCH_URGENCY_ACTIONS) br_urgency_sel = BRANCH_URGENCY_ACTIONS - 1; @@ -1728,10 +1876,24 @@ extern "C" __global__ void dqn_full_experience_kernel( if (q_values[i] > q_clip_max) q_values[i] = q_clip_max; } + /* Action masking: set masked Q-values to -inf (argmax never picks them) */ + if (enable_action_masking) { + for (int i = 0; i < NUM_ACTIONS; i++) { + if (!exposure_valid[i]) q_values[i] = -1.0e30f; + } + } + float r = gpu_random(&rng); if (r < epsilon) { - action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS); - if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1; + /* Epsilon-greedy random: sample from valid actions only */ + if (enable_action_masking) { + int sel = (int)(gpu_random(&rng) * (float)n_valid_exposures); + if (sel >= n_valid_exposures) sel = n_valid_exposures - 1; + action_idx = valid_exposure_indices[sel]; + } else { + action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS); + if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1; + } } else { if (count_bonus_coefficient > 0.0f && total_action_count > 0) { float log_n = logf((float)total_action_count); @@ -1915,8 +2077,8 @@ extern "C" __global__ void dqn_full_experience_kernel( * All threads participate in cooperative loads, same as online forward. */ float max_target_q = 0.0f; if (use_branching) { - /* Branching target: 3 independent heads, aggregate max */ - q_forward_branching( + /* Branching target: 3 independent heads, aggregate max (shmem-tiled) */ + q_forward_branching_shmem( next_state, tg_w_s1, tg_b_s1, tg_w_s2, tg_b_s2, tg_w_v1, tg_b_v1, tg_w_v2, tg_b_v2, @@ -1924,7 +2086,8 @@ extern "C" __global__ void dqn_full_experience_kernel( tg_w_bo1, tg_b_bo1, tg_w_bo2, tg_b_bo2, /* order */ tg_w_bu1, tg_b_bu1, tg_w_bu2, tg_b_bu2, /* urgency */ scratch1, scratch2, scratch_v, scratch_a, - br_tgt_exposure, br_tgt_order, br_tgt_urgency + br_tgt_exposure, br_tgt_order, br_tgt_urgency, + shmem_weights, shmem_bias ); /* Clip per-head target Q-values */ for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) { @@ -2262,6 +2425,9 @@ extern "C" __global__ void dqn_full_experience_kernel_warp( /* ---- Branching DQN config (Tavakoli et al., 2018) ---- */ int use_branching, + /* ---- Action masking: filter invalid exposure actions on GPU ---- */ + int enable_action_masking, + /* ---- RNG states [N] ---- */ unsigned int* rng_states, @@ -2348,9 +2514,14 @@ extern "C" __global__ void dqn_full_experience_kernel_warp( unsigned int rng = tid_valid ? rng_states[episode_id] : 0u; int ep_start = tid_valid ? episode_starts[episode_id] : 0; - /* ---- Distributed scratch buffers (per-lane, tiny) ---- */ + /* ---- Distributed scratch buffers (per-lane, tiny) ---- + * scratch1_dist is reused for value + advantage head outputs: + * [0..DIST_SIZE(VALUE_H)) = value head activations + * [DIST_SIZE(VALUE_H)..DIST_SIZE(VALUE_H)+DIST_SIZE(ADV_H)) = advantage activations + * So it must be max(SHARED_H1, VALUE_H + ADV_H) in distributed size. + * SCRATCH1_DIM is injected from Rust with this max already computed. */ float state_dist[DIST_SIZE(STATE_DIM)]; - float scratch1_dist[DIST_SIZE(SHARED_H1)]; + float scratch1_dist[DIST_SIZE(SCRATCH1_DIM)]; float scratch2_dist[DIST_SIZE(SHARED_H2)]; float q_values[NUM_ACTIONS]; float tgt_q_values[NUM_ACTIONS]; @@ -2559,6 +2730,25 @@ extern "C" __global__ void dqn_full_experience_kernel_warp( int br_exposure_sel = 0, br_order_sel = 0, br_urgency_sel = 0; if (!skip_data) { + /* ---- Action masking: precompute valid exposure mask (all lanes) ---- */ + const float abs_exposures[5] = {1.0f, 0.5f, 0.0f, 0.5f, 1.0f}; + int exposure_valid[5]; + int valid_exposure_indices[5]; + int n_valid_exposures = 0; + for (int i = 0; i < 5; i++) { + if (enable_action_masking && abs_exposures[i] > max_position) { + exposure_valid[i] = 0; + } else { + exposure_valid[i] = 1; + valid_exposure_indices[n_valid_exposures++] = i; + } + } + if (n_valid_exposures == 0) { + exposure_valid[2] = 1; + valid_exposure_indices[0] = 2; + n_valid_exposures = 1; + } + /* ---- Step 4b + Step 5: Q-value clipping + action selection ---- */ if (use_branching) { /* Branching: clip per-head, all lanes identical */ @@ -2575,11 +2765,25 @@ extern "C" __global__ void dqn_full_experience_kernel_warp( if (br_q_urgency[i] > q_clip_max) br_q_urgency[i] = q_clip_max; } + /* Action masking: set masked exposure Q-values to -inf */ + if (enable_action_masking) { + for (int i = 0; i < BRANCH_EXPOSURE_ACTIONS; i++) { + if (!exposure_valid[i]) br_q_exposure[i] = -1.0e30f; + } + } + if (lane_id == 0) { /* 3 independent epsilon-greedy */ + /* Exposure head — sample from valid actions only when masking */ if (gpu_random(&rng) < epsilon) { - br_exposure_sel = (int)(gpu_random(&rng) * (float)BRANCH_EXPOSURE_ACTIONS); - if (br_exposure_sel >= BRANCH_EXPOSURE_ACTIONS) br_exposure_sel = BRANCH_EXPOSURE_ACTIONS - 1; + if (enable_action_masking) { + int sel = (int)(gpu_random(&rng) * (float)n_valid_exposures); + if (sel >= n_valid_exposures) sel = n_valid_exposures - 1; + br_exposure_sel = valid_exposure_indices[sel]; + } else { + br_exposure_sel = (int)(gpu_random(&rng) * (float)BRANCH_EXPOSURE_ACTIONS); + if (br_exposure_sel >= BRANCH_EXPOSURE_ACTIONS) br_exposure_sel = BRANCH_EXPOSURE_ACTIONS - 1; + } } else { br_exposure_sel = argmax_arr(br_q_exposure, BRANCH_EXPOSURE_ACTIONS); } @@ -2623,11 +2827,24 @@ extern "C" __global__ void dqn_full_experience_kernel_warp( if (q_values[i] > q_clip_max) q_values[i] = q_clip_max; } + /* Action masking: set masked Q-values to -inf */ + if (enable_action_masking) { + for (int i = 0; i < NUM_ACTIONS; i++) { + if (!exposure_valid[i]) q_values[i] = -1.0e30f; + } + } + if (lane_id == 0) { float r = gpu_random(&rng); if (r < epsilon) { - action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS); - if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1; + if (enable_action_masking) { + int sel = (int)(gpu_random(&rng) * (float)n_valid_exposures); + if (sel >= n_valid_exposures) sel = n_valid_exposures - 1; + action_idx = valid_exposure_indices[sel]; + } else { + action_idx = (int)(gpu_random(&rng) * (float)DQN_NUM_ACTIONS); + if (action_idx >= DQN_NUM_ACTIONS) action_idx = DQN_NUM_ACTIONS - 1; + } } else { if (count_bonus_coefficient > 0.0f && total_action_count > 0) { float log_n = logf((float)total_action_count); diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 5a7d9c9e1..89f931312 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -14,15 +14,15 @@ use candle_core::cuda_backend::cudarc; use candle_core::{DType, Device, Tensor}; -use cudarc::driver::{CudaFunction, CudaSlice, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, DevicePtr, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; use std::sync::OnceLock; use tracing::info; use crate::MLError; -/// One-time PTX compilation cache. NVRTC compilation costs ~100ms, so we -/// compile once per process and reuse for all `GpuActionSelector` instances. +/// One-time PTX compilation cache. nvcc compilation is cached to disk via +/// `compile_ptx_for_device`, so the OnceLock avoids re-reading the disk cache. static EPSILON_GREEDY_PTX: OnceLock> = OnceLock::new(); /// Compile the epsilon-greedy kernel PTX, prepending common device functions. @@ -31,7 +31,7 @@ static EPSILON_GREEDY_PTX: OnceLock> = OnceLock::new(); /// but `common_device_functions.cuh` has `#error` guards that require them. /// We inject standard defaults (48/42/3) to satisfy the guards — they are never /// read by `gpu_random()`, `route_order()`, or `simulate_fill_check()`. -fn compile_kernel_ptx() -> Result { +fn compile_kernel_ptx(context: &CudaContext) -> Result { let defines = "\ #define STATE_DIM 48\n\ #define MARKET_DIM 42\n\ @@ -40,8 +40,7 @@ fn compile_kernel_ptx() -> Result { let kernel_src = include_str!("epsilon_greedy_kernel.cu"); let full_source = format!("{defines}{common_src}\n{kernel_src}"); - cudarc::nvrtc::compile_ptx(&full_source) - .map_err(|e| format!("epsilon_greedy CUDA kernel compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(&full_source, context) } /// GPU-fused epsilon-greedy action selector. @@ -81,15 +80,15 @@ impl GpuActionSelector { } }; - // Compile PTX (once per process via OnceLock) - let ptx_result = EPSILON_GREEDY_PTX.get_or_init(compile_kernel_ptx); - let ptx = ptx_result.as_ref().map_err(|e| { - MLError::ModelError(format!("epsilon_greedy PTX: {e}")) - })?; - // Load module + function let stream = cuda_dev.cuda_stream(); let context = stream.context(); + + // Compile PTX (once per process via OnceLock, disk-cached via nvcc) + let ptx_result = EPSILON_GREEDY_PTX.get_or_init(|| compile_kernel_ptx(&context)); + let ptx = ptx_result.as_ref().map_err(|e| { + MLError::ModelError(format!("epsilon_greedy PTX: {e}")) + })?; let module = context.load_module(ptx.clone()).map_err(|e| { MLError::ModelError(format!("epsilon_greedy module load: {e}")) })?; @@ -761,24 +760,24 @@ impl std::fmt::Debug for GpuActionSelector { mod tests { use super::*; - /// Verify the kernel source compiles via NVRTC (CPU-only test: just checks PTX generation). #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_ptx_compilation() { - let result = compile_kernel_ptx(); - // On machines without NVRTC, this will fail — that's expected. - // This test primarily validates that the .cu source is syntactically correct - // when NVRTC is available (CI builder has CUDA). + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { + return; + }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let result = compile_kernel_ptx(&context); if let Err(ref e) = result { - if e.contains("NVRTC") || e.contains("nvrtc") || e.contains("not found") { - // NVRTC not installed — skip gracefully - return; - } panic!("PTX compilation failed: {e}"); } } /// Verify GpuActionSelector construction fails gracefully on CPU device. #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_cpu_device_rejected() { let result = GpuActionSelector::new(&Device::Cpu, 256, 42); assert!(result.is_err()); @@ -813,36 +812,32 @@ mod tests { .select_actions(&q_values, 0.0, batch_size, num_actions) .expect("select_actions greedy"); assert_eq!(greedy_actions.dims(), &[batch_size]); - let greedy_vec: Vec = greedy_actions.to_vec1().expect("readback"); - for &a in &greedy_vec { - assert!( - (a as usize) < num_actions, - "action {a} out of range [0, {num_actions})" - ); - } - // Verify greedy matches Candle argmax - let candle_argmax = q_values - .argmax(1) - .expect("argmax") - .to_vec1::() - .expect("readback"); + // GPU-side bounds check: all actions < num_actions + let bound = Tensor::new(num_actions as u32, &device).expect("bound"); + let in_range_count = greedy_actions.lt(&bound).expect("lt") + .to_dtype(candle_core::DType::U32).expect("cast") + .sum_all().expect("sum").to_scalar::().expect("scalar"); + assert_eq!(in_range_count, batch_size as u32, "greedy actions out of range [0, {num_actions})"); + + // Verify greedy matches Candle argmax (GPU-side element-wise equality) + let candle_argmax = q_values.argmax(1).expect("argmax"); + let eq_count = greedy_actions.eq(&candle_argmax).expect("eq") + .to_dtype(candle_core::DType::U32).expect("cast") + .sum_all().expect("sum").to_scalar::().expect("scalar"); assert_eq!( - greedy_vec, candle_argmax, - "Fused GPU argmax should match Candle argmax" + eq_count, batch_size as u32, + "Fused GPU argmax should match Candle argmax ({eq_count}/{batch_size} matched)" ); // epsilon=1.0 → pure random (all actions should be valid) let random_actions = selector .select_actions(&q_values, 1.0, batch_size, num_actions) .expect("select_actions random"); - let random_vec: Vec = random_actions.to_vec1().expect("readback"); - for &a in &random_vec { - assert!( - (a as usize) < num_actions, - "random action {a} out of range [0, {num_actions})" - ); - } + let rand_in_range = random_actions.lt(&bound).expect("lt") + .to_dtype(candle_core::DType::U32).expect("cast") + .sum_all().expect("sum").to_scalar::().expect("scalar"); + assert_eq!(rand_in_range, batch_size as u32, "random actions out of range [0, {num_actions})"); } /// Test the branching DQN action selection kernel on GPU. @@ -867,19 +862,20 @@ mod tests { let actions = selector .select_actions_branching(&q_exposure, &q_order, &q_urgency, 0.0) .expect("branching select"); - let actions_vec: Vec = actions.to_vec1().expect("readback"); - for &a in &actions_vec { - assert!(a < 45, "factored action {a} should be < 45 (5×3×3)"); - } + let bound_45 = Tensor::new(45_u32, &device).expect("bound"); + let in_range_count = actions.lt(&bound_45).expect("lt") + .to_dtype(candle_core::DType::U32).expect("cast") + .sum_all().expect("sum").to_scalar::().expect("scalar"); + assert_eq!(in_range_count, batch_size as u32, "factored actions should all be < 45 (5×3×3)"); // epsilon=1.0 → random let rand_actions = selector .select_actions_branching(&q_exposure, &q_order, &q_urgency, 1.0) .expect("branching random"); - let rand_vec: Vec = rand_actions.to_vec1().expect("readback"); - for &a in &rand_vec { - assert!(a < 45, "random factored action {a} should be < 45"); - } + let rand_in_range = rand_actions.lt(&bound_45).expect("lt") + .to_dtype(candle_core::DType::U32).expect("cast") + .sum_all().expect("sum").to_scalar::().expect("scalar"); + assert_eq!(rand_in_range, batch_size as u32, "random factored actions should all be < 45"); } /// Test the routed epsilon-greedy kernel with fill simulation on GPU. @@ -916,12 +912,10 @@ mod tests { ) .expect("routed select"); - let actions_vec: Vec = actions.to_vec1().expect("readback"); - for &a in &actions_vec { - assert!( - (a as usize) < num_actions, - "routed action {a} should be < {num_actions} (exposure only)" - ); - } + let bound = Tensor::new(num_actions as u32, &device).expect("bound"); + let in_range_count = actions.lt(&bound).expect("lt") + .to_dtype(candle_core::DType::U32).expect("cast") + .sum_all().expect("sum").to_scalar::().expect("scalar"); + assert_eq!(in_range_count, batch_size as u32, "routed actions should all be < {num_actions}"); } } diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index 1fae9fddc..4f537042c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -33,7 +33,7 @@ use std::sync::Arc; use candle_core::cuda_backend::cudarc; use candle_core::{DType, Device, Tensor}; -use cudarc::driver::{CudaEvent, CudaFunction, CudaGraph, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaContext, CudaEvent, CudaFunction, CudaGraph, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; use std::sync::OnceLock; use tracing::{info, warn}; @@ -73,26 +73,23 @@ static GATHER_PTX: OnceLock> = OnceLock::new(); static PPO_FORWARD_PTX: OnceLock> = OnceLock::new(); static SUPERVISED_SIGNAL_PTX: OnceLock> = OnceLock::new(); -fn compile_env_ptx() -> Result { +fn compile_env_ptx(context: &CudaContext) -> Result { let src = include_str!("backtest_env_kernel.cu"); - cudarc::nvrtc::compile_ptx(src) - .map_err(|e| format!("backtest_env_kernel CUDA compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(src, context) } -fn compile_metrics_ptx() -> Result { +fn compile_metrics_ptx(context: &CudaContext) -> Result { let src = include_str!("backtest_metrics_kernel.cu"); - cudarc::nvrtc::compile_ptx(src) - .map_err(|e| format!("backtest_metrics_kernel CUDA compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(src, context) } -fn compile_gather_ptx() -> Result { +fn compile_gather_ptx(context: &CudaContext) -> Result { let src = include_str!("backtest_gather_kernel.cu"); - cudarc::nvrtc::compile_ptx(src) - .map_err(|e| format!("backtest_gather_kernel CUDA compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(src, context) } /// Compile PPO backtest forward kernel with common device functions prepended. -fn compile_ppo_forward_ptx() -> Result { +fn compile_ppo_forward_ptx(context: &CudaContext) -> Result { let defines = "\ #define STATE_DIM 48\n\ #define MARKET_DIM 42\n\ @@ -100,15 +97,13 @@ fn compile_ppo_forward_ptx() -> Result { let common_src = include_str!("common_device_functions.cuh"); let kernel_src = include_str!("backtest_forward_ppo_kernel.cu"); let full_source = format!("{defines}{common_src}\n{kernel_src}"); - cudarc::nvrtc::compile_ptx(&full_source) - .map_err(|e| format!("backtest_forward_ppo_kernel CUDA compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(&full_source, context) } /// Compile supervised signal-to-action kernel (standalone, no common header needed). -fn compile_supervised_signal_ptx() -> Result { +fn compile_supervised_signal_ptx(context: &CudaContext) -> Result { let src = include_str!("backtest_forward_supervised_kernel.cu"); - cudarc::nvrtc::compile_ptx(src) - .map_err(|e| format!("backtest_forward_supervised_kernel CUDA compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(src, context) } // ── Public types ────────────────────────────────────────────────────────────── @@ -373,17 +368,17 @@ impl GpuBacktestEvaluator { // ── Compile / cache kernels ─────────────────────────────────────── let env_ptx = ENV_PTX - .get_or_init(compile_env_ptx) + .get_or_init(|| compile_env_ptx(&context)) .as_ref() .map_err(|e| MLError::ModelError(format!("env kernel PTX: {e}")))?; let metrics_ptx = METRICS_PTX - .get_or_init(compile_metrics_ptx) + .get_or_init(|| compile_metrics_ptx(&context)) .as_ref() .map_err(|e| MLError::ModelError(format!("metrics kernel PTX: {e}")))?; let gather_ptx = GATHER_PTX - .get_or_init(compile_gather_ptx) + .get_or_init(|| compile_gather_ptx(&context)) .as_ref() .map_err(|e| MLError::ModelError(format!("gather kernel PTX: {e}")))?; @@ -557,11 +552,12 @@ impl GpuBacktestEvaluator { let state_dim = self.state_dim; // Launch the gather kernel — one thread per window + // Shared memory: 256 threads × 8 portfolio floats × 4 bytes = 8 KB let grid = ((self.n_windows + 255) / 256) as u32; let launch_cfg = LaunchConfig { grid_dim: (grid.max(1), 1, 1), block_dim: (256, 1, 1), - shared_mem_bytes: 0, + shared_mem_bytes: 256 * 8 * std::mem::size_of::() as u32, }; let n_windows_i32 = self.n_windows as i32; let max_len_i32 = self.max_len as i32; @@ -1084,12 +1080,6 @@ impl GpuBacktestEvaluator { actor_weights: &PpoActorWeightSet, device: &Device, ) -> Result, MLError> { - // Compile and load the PPO forward kernel - let ppo_ptx = PPO_FORWARD_PTX - .get_or_init(compile_ppo_forward_ptx) - .as_ref() - .map_err(|e| MLError::ModelError(format!("PPO forward PTX: {e}")))?; - let cuda_dev = match device { Device::Cuda(d) => d, Device::Cpu | Device::Metal(_) => { @@ -1100,6 +1090,12 @@ impl GpuBacktestEvaluator { }; let cuda_stream = cuda_dev.cuda_stream(); let context = cuda_stream.context(); + + // Compile and load the PPO forward kernel + let ppo_ptx = PPO_FORWARD_PTX + .get_or_init(|| compile_ppo_forward_ptx(&context)) + .as_ref() + .map_err(|e| MLError::ModelError(format!("PPO forward PTX: {e}")))?; let ppo_module = context .load_module(ppo_ptx.clone()) .map_err(|e| MLError::ModelError(format!("PPO forward module load: {e}")))?; @@ -1112,11 +1108,11 @@ impl GpuBacktestEvaluator { self.launch_gather(step)?; // 2. Launch PPO forward kernel: states → actions (pure CUDA, no candle) - // grid=(N,1,1), block=(1,1,1) — one thread per window + // grid=(ceil(N/32),1,1), block=(32,1,1) — 32 independent windows/block let n_i32 = self.n_windows as i32; let ppo_cfg = LaunchConfig { - grid_dim: (self.n_windows as u32, 1, 1), - block_dim: (1, 1, 1), + grid_dim: (((self.n_windows as u32) + 31) / 32, 1, 1), + block_dim: (32, 1, 1), shared_mem_bytes: 0, }; @@ -1182,12 +1178,6 @@ impl GpuBacktestEvaluator { ))); } - // Compile and load the signal-to-action kernel - let sig_ptx = SUPERVISED_SIGNAL_PTX - .get_or_init(compile_supervised_signal_ptx) - .as_ref() - .map_err(|e| MLError::ModelError(format!("supervised signal PTX: {e}")))?; - let cuda_dev = match device { Device::Cuda(d) => d, Device::Cpu | Device::Metal(_) => { @@ -1198,6 +1188,12 @@ impl GpuBacktestEvaluator { }; let cuda_stream = cuda_dev.cuda_stream(); let context = cuda_stream.context(); + + // Compile and load the signal-to-action kernel + let sig_ptx = SUPERVISED_SIGNAL_PTX + .get_or_init(|| compile_supervised_signal_ptx(&context)) + .as_ref() + .map_err(|e| MLError::ModelError(format!("supervised signal PTX: {e}")))?; let sig_module = context .load_module(sig_ptx.clone()) .map_err(|e| MLError::ModelError(format!("supervised signal module load: {e}")))?; @@ -1359,11 +1355,12 @@ impl GpuBacktestEvaluator { /// `evaluate_ppo` paths read directly from `states_buf`. fn launch_gather(&self, step: usize) -> Result<(), MLError> { let state_dim = self.state_dim; + // Shared memory: 256 threads × 8 portfolio floats × 4 bytes = 8 KB let grid = ((self.n_windows + 255) / 256) as u32; let launch_cfg = LaunchConfig { grid_dim: (grid.max(1), 1, 1), block_dim: (256, 1, 1), - shared_mem_bytes: 0, + shared_mem_bytes: 256 * 8 * std::mem::size_of::() as u32, }; let n_windows_i32 = self.n_windows as i32; let max_len_i32 = self.max_len as i32; @@ -1398,11 +1395,12 @@ impl GpuBacktestEvaluator { /// Launch `backtest_env_step` kernel for a given step on the specified stream. fn launch_env_step_on(&self, step: usize, target_stream: &Arc) -> Result<(), MLError> { + // Shared memory: 256 threads × PORTFOLIO_STATE_SIZE(8) floats × 4 bytes = 8 KB let grid = ((self.n_windows + 255) / 256) as u32; let env_cfg = LaunchConfig { grid_dim: (grid.max(1), 1, 1), block_dim: (256, 1, 1), - shared_mem_bytes: 0, + shared_mem_bytes: 256 * 8 * std::mem::size_of::() as u32, }; let n_win_i32 = self.n_windows as i32; let max_len_i32 = self.max_len as i32; @@ -1580,6 +1578,7 @@ mod tests { } #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_new_rejects_cpu_device() { let prices = vec![vec![[1.0_f32; 4]; 5]]; let features = vec![vec![vec![0.0_f32; 4]; 5]]; @@ -1595,48 +1594,42 @@ mod tests { assert!(msg.contains("CUDA"), "expected CUDA error, got: {msg}"); } - /// Verify PTX sources compile without errors (skips gracefully when NVRTC absent). + /// Verify PTX sources compile without errors (requires CUDA — nvcc path). #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_env_ptx_compilation() { - let result = compile_env_ptx(); + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { return }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let result = compile_env_ptx(&context); if let Err(ref e) = result { - if e.contains("NVRTC") - || e.contains("nvrtc") - || e.contains("not found") - || e.contains("No such file") - { - return; // NVRTC not installed — acceptable on CPU-only machines - } panic!("backtest_env_kernel PTX compilation failed: {e}"); } } #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_metrics_ptx_compilation() { - let result = compile_metrics_ptx(); + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { return }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let result = compile_metrics_ptx(&context); if let Err(ref e) = result { - if e.contains("NVRTC") - || e.contains("nvrtc") - || e.contains("not found") - || e.contains("No such file") - { - return; - } panic!("backtest_metrics_kernel PTX compilation failed: {e}"); } } #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_gather_ptx_compilation() { - let result = compile_gather_ptx(); + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { return }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let result = compile_gather_ptx(&context); if let Err(ref e) = result { - if e.contains("NVRTC") - || e.contains("nvrtc") - || e.contains("not found") - || e.contains("No such file") - { - return; // NVRTC not installed — acceptable on CPU-only machines - } panic!("backtest_gather_kernel PTX compilation failed: {e}"); } } @@ -1650,12 +1643,16 @@ mod tests { assert_eq!(3_usize, 3_usize, "PORTFOLIO_DIM constant is 3"); } - /// Verify that the forward kernel CUDA source compiles (NVRTC). + /// Verify that the forward kernel CUDA source compiles (nvcc path). #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_forward_kernel_ptx_compilation() { + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { return }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); let common_src = include_str!("common_device_functions.cuh"); let kernel_src = include_str!("backtest_forward_kernel.cu"); - // Use default dimensions for compilation test let dim_overrides = "\ #define STATE_DIM 48\n\ #define MARKET_DIM 42\n\ @@ -1667,48 +1664,36 @@ mod tests { #define SHMEM_MAX_IN_DIM 256\n\ #define SHMEM_TILE_ROWS 32\n"; let full_source = format!("{common_src}\n{dim_overrides}\n{kernel_src}"); - let result = cudarc::nvrtc::compile_ptx(&full_source); + let result = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context); if let Err(ref e) = result { - let msg = format!("{e}"); - if msg.contains("NVRTC") - || msg.contains("nvrtc") - || msg.contains("not found") - || msg.contains("No such file") - { - return; // NVRTC not installed — acceptable on CPU-only machines - } panic!("backtest_forward_kernel PTX compilation failed: {e}"); } } - /// Verify PPO forward PTX compiles (skips gracefully when NVRTC absent). + /// Verify PPO forward PTX compiles (requires CUDA — nvcc path). #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_ppo_forward_ptx_compilation() { - let result = compile_ppo_forward_ptx(); + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { return }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let result = compile_ppo_forward_ptx(&context); if let Err(ref e) = result { - if e.contains("NVRTC") - || e.contains("nvrtc") - || e.contains("not found") - || e.contains("No such file") - { - return; // NVRTC not installed — acceptable on CPU-only machines - } panic!("backtest_forward_ppo_kernel PTX compilation failed: {e}"); } } - /// Verify supervised signal-to-action PTX compiles (skips gracefully when NVRTC absent). + /// Verify supervised signal-to-action PTX compiles (requires CUDA — nvcc path). #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_supervised_signal_ptx_compilation() { - let result = compile_supervised_signal_ptx(); + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { return }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let result = compile_supervised_signal_ptx(&context); if let Err(ref e) = result { - if e.contains("NVRTC") - || e.contains("nvrtc") - || e.contains("not found") - || e.contains("No such file") - { - return; // NVRTC not installed — acceptable on CPU-only machines - } panic!("backtest_forward_supervised_kernel PTX compilation failed: {e}"); } } diff --git a/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs index bdc1523f7..495764f5d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_curiosity_trainer.rs @@ -17,7 +17,7 @@ use std::sync::{Arc, OnceLock}; use candle_core::cuda_backend::cudarc; -use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; use tracing::debug; @@ -50,7 +50,7 @@ const ADAM_EPS: f32 = 1e-8; static CURIOSITY_TRAINING_PTX: OnceLock> = OnceLock::new(); -fn compile_curiosity_training_ptx() -> Result { +fn compile_curiosity_training_ptx(context: &CudaContext) -> Result { let defines = "\ #define STATE_DIM 48\n\ #define MARKET_DIM 42\n\ @@ -58,8 +58,7 @@ fn compile_curiosity_training_ptx() -> Result { let common_src = include_str!("common_device_functions.cuh"); let kernel_src = include_str!("curiosity_training_kernel.cu"); let full_source = format!("{defines}{common_src}\n{kernel_src}"); - cudarc::nvrtc::compile_ptx(&full_source) - .map_err(|e| format!("curiosity_training CUDA kernel compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(&full_source, context) } // --------------------------------------------------------------------------- @@ -79,6 +78,7 @@ pub struct GpuCuriosityTrainer { shift_func: CudaFunction, fwd_bwd_func: CudaFunction, adam_func: CudaFunction, + adam_fused_func: CudaFunction, // Gradient buffers grad_w1: CudaSlice, // [CUR_W1_LEN] @@ -171,7 +171,7 @@ impl GpuCuriosityTrainer { max_samples: usize, ) -> Result { // ---- Compile and load kernels ---- - let ptx_result = CURIOSITY_TRAINING_PTX.get_or_init(compile_curiosity_training_ptx); + let ptx_result = CURIOSITY_TRAINING_PTX.get_or_init(|| compile_curiosity_training_ptx(&stream.context())); let ptx = ptx_result .as_ref() .map_err(|e| MLError::ModelError(format!("curiosity training PTX: {e}")))?; @@ -190,6 +190,9 @@ impl GpuCuriosityTrainer { let adam_func = module.load_function("curiosity_adam_step").map_err(|e| { MLError::ModelError(format!("curiosity_adam_step load: {e}")) })?; + let adam_fused_func = module.load_function("curiosity_adam_step_fused").map_err(|e| { + MLError::ModelError(format!("curiosity_adam_step_fused load: {e}")) + })?; // ---- Allocate gradient buffers ---- let grad_w1 = stream.alloc_zeros::(CUR_W1_LEN).map_err(|e| { @@ -252,6 +255,7 @@ impl GpuCuriosityTrainer { shift_func, fwd_bwd_func, adam_func, + adam_fused_func, grad_w1, grad_b1, grad_w2, @@ -380,29 +384,32 @@ impl GpuCuriosityTrainer { self.step += 1; let step = self.step; - // ---- Step 5: Adam optimizer step (4 launches, one per param group) ---- - // Use free function to avoid borrow conflicts between &self fields - // and &mut self fields needed simultaneously. - launch_adam_step( - &self.stream, &self.adam_func, - &mut weights.w1, &self.grad_w1, CUR_W1_LEN, - &mut self.adam_m_w1, &mut self.adam_v_w1, n_train, step, - )?; - launch_adam_step( - &self.stream, &self.adam_func, - &mut weights.b1, &self.grad_b1, CUR_B1_LEN, - &mut self.adam_m_b1, &mut self.adam_v_b1, n_train, step, - )?; - launch_adam_step( - &self.stream, &self.adam_func, - &mut weights.w2, &self.grad_w2, CUR_W2_LEN, - &mut self.adam_m_w2, &mut self.adam_v_w2, n_train, step, - )?; - launch_adam_step( - &self.stream, &self.adam_func, - &mut weights.b2, &self.grad_b2, CUR_B2_LEN, - &mut self.adam_m_b2, &mut self.adam_v_b2, n_train, step, - )?; + // ---- Step 5: Fused Adam optimizer step (single launch for all 4 param groups) ---- + let total_params = CUR_W1_LEN + CUR_B1_LEN + CUR_W2_LEN + CUR_B2_LEN; + let fused_cfg = LaunchConfig { + grid_dim: (((total_params as u32) + 255) / 256, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + let n0 = CUR_W1_LEN as i32; + let n1 = CUR_B1_LEN as i32; + let n2 = CUR_W2_LEN as i32; + let n3 = CUR_B2_LEN as i32; + let bs_i32 = n_train as i32; + unsafe { + self.stream + .launch_builder(&self.adam_fused_func) + .arg(&mut weights.w1).arg(&self.grad_w1).arg(&mut self.adam_m_w1).arg(&mut self.adam_v_w1).arg(&n0) + .arg(&mut weights.b1).arg(&self.grad_b1).arg(&mut self.adam_m_b1).arg(&mut self.adam_v_b1).arg(&n1) + .arg(&mut weights.w2).arg(&self.grad_w2).arg(&mut self.adam_m_w2).arg(&mut self.adam_v_w2).arg(&n2) + .arg(&mut weights.b2).arg(&self.grad_b2).arg(&mut self.adam_m_b2).arg(&mut self.adam_v_b2).arg(&n3) + .arg(&bs_i32) + .arg(&ADAM_LR).arg(&ADAM_BETA1).arg(&ADAM_BETA2).arg(&ADAM_EPS).arg(&step) + .launch(fused_cfg) + .map_err(|e| { + MLError::ModelError(format!("curiosity_adam_step_fused launch: {e}")) + })?; + } debug!(step, n_train, "curiosity GPU training step complete"); diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 3f4ab271d..2dc3b92b5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -135,6 +135,10 @@ pub struct ExperienceCollectorConfig { pub dsr_eta: f32, /// N-step returns lookahead (1 = standard TD, 3-5 typical for Rainbow DQN) pub n_steps: i32, + /// Enable action masking: filters invalid exposure actions in the GPU kernel. + /// When enabled, actions whose |exposure| > max_position are set to -inf + /// before argmax and excluded from epsilon-greedy random sampling. + pub enable_action_masking: bool, } impl Default for ExperienceCollectorConfig { @@ -179,6 +183,7 @@ impl Default for ExperienceCollectorConfig { use_dsr: true, dsr_eta: 0.01, n_steps: 1, + enable_action_masking: false, } } } @@ -384,10 +389,16 @@ impl GpuExperienceCollector { let pow2 = (max_tile as u32).next_power_of_two() >> 1; pow2.clamp(16, 256) as usize }; - let portfolio_dim = 3usize; // cash, position_size, unrealized_pnl + let portfolio_dim: usize = 3; // cash, position_size, unrealized_pnl // OFI features: 8 when state_dim accommodates them (56 = 42+3+8+3pad), // 0 when state_dim is 48 (42+3+0+3pad) or smaller. - let ofi_dim = if state_dim >= market_dim + portfolio_dim + 8 { 8usize } else { 0usize }; + let ofi_dim: usize = if state_dim >= market_dim + portfolio_dim + 8 { 8 } else { 0 }; + // scratch1_dist is reused inside the forward function: + // - first as output of shared layer 1 (SHARED_H1 elements) + // - then as value head output (VALUE_H elements at offset 0) + // PLUS advantage head output (ADV_H elements at offset DIST_SIZE(VALUE_H)) + // So scratch1 must hold max(SHARED_H1, VALUE_H + ADV_H) in distributed form. + let scratch1_dim = shared_h1.max(value_h + adv_h); let dim_overrides = format!( "#define NOISY_MAX_DIM {noisy_max_dim}\n\ #define STATE_DIM {state_dim}\n\ @@ -398,6 +409,7 @@ impl GpuExperienceCollector { #define SHARED_H2 {shared_h2}\n\ #define VALUE_H {value_h}\n\ #define ADV_H {adv_h}\n\ + #define SCRATCH1_DIM {scratch1_dim}\n\ #define NUM_ATOMS_MAX {num_atoms_max}\n\ #define SHMEM_MAX_IN_DIM {shmem_max_in_dim}\n\ #define SHMEM_TILE_ROWS {shmem_tile_rows}\n" @@ -565,19 +577,45 @@ impl GpuExperienceCollector { } } - // Increase per-thread stack size for large hidden dimensions. - // With SHARED_H1=512 (H100), scratch arrays consume ~5.5KB per thread, - // exceeding the default 1KB CUDA stack limit. + // Dynamically calculate per-thread stack size from kernel scratch arrays. + // + // The warp kernel allocates these per-thread arrays on the stack: + // noisy functions: eps_in[NOISY_MAX_DIM] + eps_out[SHMEM_TILE_ROWS] + // kernel body: div_window[100], barrier_st[5], q_values[5], tgt_q_values[5], + // action_counts[5], nstep_ring[5], state_dist, scratch_dist(x2), + // exposure_valid[5], valid_exposure_indices[5] + // branching: br_q_{exposure,order,urgency}[5+3+3], + // br_tgt_{exposure,order,urgency}[5+3+3], + // adv_e[5], adv_o[3], adv_u[3] (inside forward fn) + // + // Total per-thread arrays: ~800B (kernel body) + ~1.2KB (noisy worst-case) + // Plus call frames, saved registers, and nvcc -O3 spills. { use cudarc::driver::sys::{cuCtxSetLimit, CUlimit, CUresult}; - let stack_bytes = 0x4000_usize; // 16KB — enough for 5.5KB scratch + noise + call frames + let non_tiled_bytes = 2 * noisy_max_dim * 4; // eps_in + eps_out both [NOISY_MAX_DIM] + let tiled_bytes = (noisy_max_dim + shmem_tile_rows) * 4; // eps_in[NOISY_MAX_DIM] + eps_out[SHMEM_TILE_ROWS] + let noisy_bytes = non_tiled_bytes.max(tiled_bytes); + // Kernel body arrays: div_window[100]=400 + barrier_st[5]=20 + q/tgt/action/nstep=100 + // + branching arrays=176 + exposure masks=40 + dist vectors=28 = ~764 bytes. + let kernel_body_bytes = 800_usize; + // 2KB overhead for call frames, function params, loop vars, and + // nvcc -O3 register spills. Measured ~800B on sm_90 pre-branching, + // estimated ~1.5KB with branching + action masking + fill simulation. + let needed = noisy_bytes + kernel_body_bytes + 2048; + // Round up to next power of 2, clamp to [8KB, 32KB]. + let stack_bytes = (needed as u32).next_power_of_two().clamp(0x2000, 0x8000) as usize; let result = unsafe { cuCtxSetLimit(CUlimit::CU_LIMIT_STACK_SIZE, stack_bytes) }; if result != CUresult::CUDA_SUCCESS { return Err(MLError::ModelError(format!( - "cuCtxSetLimit(STACK_SIZE, {stack_bytes}) FAILED: {result:?} — experience kernel needs 16KB stack" + "cuCtxSetLimit(STACK_SIZE, {stack_bytes}) FAILED: {result:?} — \ + kernel needs {stack_bytes}B stack (noisy={noisy_bytes}B, body={kernel_body_bytes}B, \ + noisy_max_dim={noisy_max_dim}, tile_rows={shmem_tile_rows})" ))); } else { - tracing::info!("CUDA stack size set to {stack_bytes} bytes"); + tracing::info!( + "CUDA stack size set to {stack_bytes}B (noisy={noisy_bytes}B, body={kernel_body_bytes}B, \ + noisy_max_dim={noisy_max_dim}, tile_rows={shmem_tile_rows})" + ); } } @@ -907,6 +945,10 @@ impl GpuExperienceCollector { self.launch_kernel(market_features_buf, targets_buf, episode_starts, config)?; // ---- Step 6: Download output buffers (only N*L elements) ---- + // Note: cudarc's memcpy_dtoh internally calls cuMemcpyDtoHAsync, but with pageable + // Vec destinations CUDA blocks until the copy completes (no true DMA overlap). + // True async batching requires PinnedHostSlice destinations — deferred until profiling + // shows download latency dominates (currently ~2% of collect_experiences walltime). let total = n_episodes * timesteps; let states_view = self.states_out.slice(..total * self.state_dim); @@ -950,33 +992,25 @@ impl GpuExperienceCollector { ); // ---- Step 7: Build next_states from states (episode-aware shift) ---- - // next_states[t] = states[t+1] within same episode; terminal/last uses current state. + // next_states[t] = states[t+1] within same episode; terminal/last copies current state. + // Pre-allocated flat copy avoids per-element bounds checks and extend_from_slice overhead. let sd = self.state_dim; - let mut next_states = Vec::with_capacity(total * sd); + let mut next_states = vec![0.0_f32; total * sd]; for ep in 0..n_episodes { for t in 0..timesteps { let idx = ep * timesteps + t; let is_done = done_flags.get(idx).copied().unwrap_or(0) != 0; - if !is_done && t + 1 < timesteps { - let next_idx = ep * timesteps + t + 1; - let ns_start = next_idx * sd; - let ns_end = ns_start + sd; - if let Some(ns) = states.get(ns_start..ns_end) { - next_states.extend_from_slice(ns); - } else { - // Out of bounds — copy current state as fallback - let s_start = idx * sd; - if let Some(s) = states.get(s_start..s_start + sd) { - next_states.extend_from_slice(s); - } - } + let src_idx = if !is_done && t + 1 < timesteps { + ep * timesteps + t + 1 // next timestep within episode } else { - // Terminal or last timestep — next_state = current state - // (masked by done flag in Bellman equation) - let s_start = idx * sd; - if let Some(s) = states.get(s_start..s_start + sd) { - next_states.extend_from_slice(s); - } + idx // terminal or last: copy current state (masked by done flag in Bellman) + }; + let dst = idx * sd; + let src = src_idx * sd; + if let (Some(dst_slice), Some(src_slice)) = + (next_states.get_mut(dst..dst + sd), states.get(src..src + sd)) + { + dst_slice.copy_from_slice(src_slice); } } } @@ -1193,13 +1227,13 @@ impl GpuExperienceCollector { &self.stream, ); + // Extract ofi_slice before builder chain to avoid borrow conflict + // (ofi_slice() borrows &self while later .arg() borrows &mut self.field). + let ofi_buf = self.ofi_slice() as *const CudaSlice; // Safety: kernel parameter order matches dqn_experience_kernel.cu exactly. // All buffer sizes are pre-validated (MAX_EPISODES, MAX_TIMESTEPS). // All CudaSlice lifetimes are valid (owned by self). // Both standard and warp kernels share the same parameter signature. - // Extract ofi_slice before builder chain to avoid borrow conflict - // (ofi_slice() borrows &self while later .arg() borrows &mut self.field). - let ofi_buf = self.ofi_slice() as *const CudaSlice; unsafe { self.stream .launch_builder(active_kernel) @@ -1262,6 +1296,8 @@ impl GpuExperienceCollector { .arg(&config.n_steps) // Branching DQN flag .arg(&(self.use_branching as i32)) + // Action masking flag + .arg(&(config.enable_action_masking as i32)) // RNG .arg(&mut self.rng_states) // Outputs @@ -1511,8 +1547,9 @@ impl GpuExperienceCollector { /// Returns 1-element placeholder when no OFI data uploaded (kernel sees out-of-bounds /// indices and uses 0.0 — safe because OFI_DIM=0 means the kernel never reads from it). fn ofi_slice(&self) -> &CudaSlice { - // Safety: ofi_gpu is always Some (1-element placeholder allocated in new()). - // If allocation failed, unwrap panics at init time, not at kernel launch. + // ofi_gpu is always Some (1-element placeholder allocated in new()). + // If allocation failed, the constructor would have returned Err. + #[allow(clippy::expect_used)] self.ofi_gpu.as_ref().expect("BUG: ofi_gpu placeholder not allocated") } } diff --git a/crates/ml/src/cuda_pipeline/gpu_monitoring.rs b/crates/ml/src/cuda_pipeline/gpu_monitoring.rs index b1b4b4a52..3132a2448 100644 --- a/crates/ml/src/cuda_pipeline/gpu_monitoring.rs +++ b/crates/ml/src/cuda_pipeline/gpu_monitoring.rs @@ -5,7 +5,7 @@ use std::sync::{Arc, OnceLock}; use candle_core::cuda_backend::cudarc; -use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; use ml_core::nvtx::NvtxRange; use crate::MLError; @@ -34,16 +34,15 @@ pub struct GpuMonitoringReducer { static MONITORING_PTX: OnceLock> = OnceLock::new(); -fn compile_monitoring_ptx() -> Result { +fn compile_monitoring_ptx(context: &CudaContext) -> Result { let kernel_src = include_str!("monitoring_kernel.cu"); - cudarc::nvrtc::compile_ptx(kernel_src) - .map_err(|e| format!("monitoring kernel CUDA compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context) } impl GpuMonitoringReducer { pub fn new(stream: &Arc) -> Result { let context = stream.context(); - let ptx_result = MONITORING_PTX.get_or_init(compile_monitoring_ptx); + let ptx_result = MONITORING_PTX.get_or_init(|| compile_monitoring_ptx(&context)); let ptx = ptx_result.as_ref() .map_err(|e| MLError::ModelError(format!("monitoring PTX: {e}")))?; let module = context.load_module(ptx.clone()) @@ -122,19 +121,17 @@ mod tests { assert_eq!(s.action_counts, [0; 5]); } - /// Verify the PTX compiles without errors (skips gracefully when NVRTC is absent). #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_monitoring_ptx_compilation() { - let result = compile_monitoring_ptx(); + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { + return; + }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let result = compile_monitoring_ptx(&context); if let Err(ref e) = result { - if e.contains("NVRTC") - || e.contains("nvrtc") - || e.contains("not found") - || e.contains("No such file") - { - // NVRTC not installed — acceptable on CPU-only machines. - return; - } panic!("monitoring kernel PTX compilation failed: {e}"); } } diff --git a/crates/ml/src/cuda_pipeline/gpu_portfolio.rs b/crates/ml/src/cuda_pipeline/gpu_portfolio.rs index 3be126dbf..21f274bd5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_portfolio.rs +++ b/crates/ml/src/cuda_pipeline/gpu_portfolio.rs @@ -63,7 +63,7 @@ pub struct PortfolioSimResult { fn compile_experience_kernels(context: &Arc) -> Result { let kernel_src = include_str!("experience_kernels.cu"); - let ptx: Ptx = cudarc::nvrtc::compile_ptx(kernel_src).map_err(|e| { + let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context).map_err(|e| { MLError::ModelError(format!( "CUDA experience kernel compilation failed (is NVRTC available?): {e}" )) diff --git a/crates/ml/src/cuda_pipeline/gpu_statistics.rs b/crates/ml/src/cuda_pipeline/gpu_statistics.rs index 40cee3eca..d77a6adc0 100644 --- a/crates/ml/src/cuda_pipeline/gpu_statistics.rs +++ b/crates/ml/src/cuda_pipeline/gpu_statistics.rs @@ -9,7 +9,7 @@ use candle_core::cuda_backend::cudarc; use candle_core::Device; -use cudarc::driver::{CudaFunction, CudaSlice, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; use std::sync::OnceLock; @@ -18,10 +18,9 @@ use crate::MLError; static STATISTICS_PTX: OnceLock> = OnceLock::new(); -fn compile_statistics_ptx() -> Result { +fn compile_statistics_ptx(context: &CudaContext) -> Result { let kernel_src = include_str!("statistics_kernel.cu"); - cudarc::nvrtc::compile_ptx(kernel_src) - .map_err(|e| format!("batch_statistics CUDA kernel compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context) } /// Aggregated batch statistics from GPU parallel reduction. @@ -54,13 +53,13 @@ impl GpuStatistics { Device::Cpu | Device::Metal(_) => return Err(MLError::ModelError("GpuStatistics requires CUDA".into())), }; - let ptx_result = STATISTICS_PTX.get_or_init(compile_statistics_ptx); + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + + let ptx_result = STATISTICS_PTX.get_or_init(|| compile_statistics_ptx(&context)); let ptx = ptx_result.as_ref().map_err(|e| { MLError::ModelError(format!("statistics PTX: {e}")) })?; - - let stream = cuda_dev.cuda_stream(); - let context = stream.context(); let module = context.load_module(ptx.clone()).map_err(|e| { MLError::ModelError(format!("statistics module load: {e}")) })?; diff --git a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs index d353268fa..de68171f2 100644 --- a/crates/ml/src/cuda_pipeline/gpu_training_guard.rs +++ b/crates/ml/src/cuda_pipeline/gpu_training_guard.rs @@ -16,7 +16,7 @@ use candle_core::cuda_backend::cudarc; use candle_core::{DType, Device, Tensor}; -use cudarc::driver::{CudaFunction, CudaSlice, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; use std::ffi::c_void; use std::mem::MaybeUninit; @@ -29,10 +29,9 @@ use crate::MLError; static TRAINING_GUARD_PTX: OnceLock> = OnceLock::new(); -fn compile_training_guard_ptx() -> Result { +fn compile_training_guard_ptx(context: &CudaContext) -> Result { let kernel_src = include_str!("training_guard_kernel.cu"); - cudarc::nvrtc::compile_ptx(kernel_src) - .map_err(|e| format!("training_guard_kernel CUDA compilation failed: {e}")) + crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context) } // ── Mapped pinned memory helper ─────────────────────────────────────────── @@ -92,11 +91,19 @@ impl MappedBuffer { /// Uses `read_volatile` to prevent the CPU from caching stale values. /// Caller must ensure that the kernel writing to this buffer has completed /// (e.g., via double-buffering with sufficient pipeline delay). + /// + /// # Panics (debug builds) + /// Panics if `idx >= self.len`. All call sites use compile-time-known + /// indices that are within bounds (verified by construction). fn read(&self, idx: usize) -> f32 { - if idx >= self.len { - return 0.0; - } - // Safety: host_ptr is valid for self.len floats, idx is bounds-checked above. + debug_assert!( + idx < self.len, + "MappedBuffer::read OOB: idx={idx} >= len={}", + self.len + ); + // Safety: host_ptr is valid for self.len floats; all callers use + // hardcoded indices < len (guard=0..5 of 7, qstats=0..3 of 4, + // qdiv=0..4 of 5). debug_assert catches bugs in dev builds. unsafe { std::ptr::read_volatile(self.host_ptr.add(idx)) } } } @@ -168,8 +175,7 @@ pub struct QValueDivergence { /// `memcpy_dtoh` from the per-step hot path. The kernel writes to /// one buffer while the CPU reads the other (one-step delay). pub struct GpuTrainingGuard { - check_func: CudaFunction, - accumulate_func: CudaFunction, + fused_check_accum_func: CudaFunction, qvalue_stats_func: CudaFunction, qvalue_div_func: CudaFunction, @@ -217,25 +223,22 @@ impl GpuTrainingGuard { } }; - // Compile PTX (once per process) - let ptx_result = TRAINING_GUARD_PTX.get_or_init(compile_training_guard_ptx); - let ptx = ptx_result.as_ref().map_err(|e| { - MLError::ModelError(format!("training_guard PTX: {e}")) - })?; - // Load module and functions let stream = cuda_dev.cuda_stream(); let context = stream.context(); + + // Compile PTX (once per process via OnceLock, disk-cached via nvcc) + let ptx_result = TRAINING_GUARD_PTX.get_or_init(|| compile_training_guard_ptx(&context)); + let ptx = ptx_result.as_ref().map_err(|e| { + MLError::ModelError(format!("training_guard PTX: {e}")) + })?; let module = context.load_module(ptx.clone()).map_err(|e| { MLError::ModelError(format!("training_guard module load: {e}")) })?; - let check_func = module - .load_function("training_guard_check") - .map_err(|e| MLError::ModelError(format!("training_guard_check load: {e}")))?; - let accumulate_func = module - .load_function("training_guard_accumulate") - .map_err(|e| MLError::ModelError(format!("training_guard_accumulate load: {e}")))?; + let fused_check_accum_func = module + .load_function("training_guard_check_and_accumulate") + .map_err(|e| MLError::ModelError(format!("training_guard_check_and_accumulate load: {e}")))?; let qvalue_stats_func = module .load_function("qvalue_stats_reduce") .map_err(|e| MLError::ModelError(format!("qvalue_stats_reduce load: {e}")))?; @@ -258,8 +261,7 @@ impl GpuTrainingGuard { .map_err(|e| MLError::ModelError(format!("q_mean_tensor init: {e}")))?; Ok(Self { - check_func, - accumulate_func, + fused_check_accum_func, qvalue_stats_func, qvalue_div_func, guard_mapped, @@ -378,7 +380,9 @@ impl GpuTrainingGuard { let warmup_int: i32 = if warmup { 1 } else { 0 }; - // Launch kernel 1: training_guard_check -- grid=(1,1,1), block=(1,1,1) + // Fused kernel: training_guard_check_and_accumulate -- grid=(1,1,1), block=(1,1,1) + // Reads loss/grad_norm ONCE, does guard check + accumulate in a single launch. + // Saves ~2-4μs/step in kernel launch overhead (12-25ms/epoch at 6375 steps). let single_cfg = LaunchConfig { grid_dim: (1, 1, 1), block_dim: (1, 1, 1), @@ -388,31 +392,19 @@ impl GpuTrainingGuard { // Safety: all GPU slices are valid, layout offsets applied. // write_dev_ptr is a CUdeviceptr (u64) pointing into mapped host memory; // the kernel writes 7 floats to it, __threadfence_system ensures visibility. + // acc_buf is 3 f32 on same device (device memory). unsafe { stream - .launch_builder(&self.check_func) + .launch_builder(&self.fused_check_accum_func) .arg(&loss_view) .arg(&grad_view) .arg(&write_dev_ptr) + .arg(&mut self.acc_buf) .arg(&clip_threshold) .arg(&collapse_threshold) .arg(&warmup_int) .launch(single_cfg) - .map_err(|e| MLError::ModelError(format!("training_guard_check launch: {e}")))?; - } - - // Launch kernel 2: training_guard_accumulate -- grid=(1,1,1), block=(1,1,1) - // Safety: loss_view/grad_view valid; acc_buf is 3 f32 on same device. - unsafe { - stream - .launch_builder(&self.accumulate_func) - .arg(&loss_view) - .arg(&grad_view) - .arg(&mut self.acc_buf) - .launch(single_cfg) - .map_err(|e| { - MLError::ModelError(format!("training_guard_accumulate launch: {e}")) - })?; + .map_err(|e| MLError::ModelError(format!("training_guard_check_and_accumulate launch: {e}")))?; } // Swap buffer index for next step @@ -688,16 +680,23 @@ impl GpuTrainingGuard { } /// Accumulate a Q-value mean on GPU (zero CPU sync). Uses running Welford accumulator. + /// + /// Uses `affine(1/count, 0)` instead of `Tensor::new(count)` + `div` to avoid + /// per-step GPU scalar allocation (cuMemAlloc + cuMemcpyHtoD ~2-5μs each). pub fn accumulate_q_value(&mut self, avg_q_tensor: &Tensor) -> Result<(), MLError> { self.q_count += 1; let delta = avg_q_tensor .sub(&self.q_mean_tensor) .map_err(|e| MLError::ModelError(format!("q_mean delta: {e}")))?; - let count_f = Tensor::new(self.q_count as f32, avg_q_tensor.device()) - .map_err(|e| MLError::ModelError(format!("q_count tensor: {e}")))?; + // Use affine(scale, bias) to divide by count without allocating a GPU tensor. + // affine(1/n, 0) = delta * (1/n) + 0, computed via a single cuBLAS call. + let inv_count = 1.0_f64 / self.q_count as f64; + let scaled_delta = delta + .affine(inv_count, 0.0) + .map_err(|e| MLError::ModelError(format!("q_mean scale: {e}")))?; self.q_mean_tensor = self .q_mean_tensor - .add(&delta.div(&count_f).map_err(|e| MLError::ModelError(format!("q_mean div: {e}")))?) + .add(&scaled_delta) .map_err(|e| MLError::ModelError(format!("q_mean add: {e}")))?; Ok(()) } @@ -734,25 +733,24 @@ impl std::fmt::Debug for GpuTrainingGuard { mod tests { use super::*; - /// Verify the PTX compiles without errors (skips gracefully when NVRTC is absent). #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_ptx_compilation() { - let result = compile_training_guard_ptx(); + let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); + let candle_core::Device::Cuda(ref cuda_dev) = device else { + return; + }; + let stream = cuda_dev.cuda_stream(); + let context = stream.context(); + let result = compile_training_guard_ptx(&context); if let Err(ref e) = result { - if e.contains("NVRTC") - || e.contains("nvrtc") - || e.contains("not found") - || e.contains("No such file") - { - // NVRTC not installed -- acceptable on CPU-only machines. - return; - } panic!("training_guard PTX compilation failed: {e}"); } } /// Verify `GpuTrainingGuard::new` fails gracefully on a CPU device. #[test] + #[cfg_attr(not(feature = "cuda"), ignore)] fn test_cpu_device_rejected() { let result = GpuTrainingGuard::new(&Device::Cpu); assert!(result.is_err()); diff --git a/crates/ml/src/cuda_pipeline/gpu_walk_forward.rs b/crates/ml/src/cuda_pipeline/gpu_walk_forward.rs new file mode 100644 index 000000000..72d9ff7df --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_walk_forward.rs @@ -0,0 +1,407 @@ +//! GPU-resident walk-forward data manager. +//! +//! Uploads the ENTIRE dataset to GPU VRAM once, then provides per-fold +//! views via index ranges. This eliminates per-fold re-upload latency +//! and enables zero-copy fold transitions on H100 (80 GB VRAM). +//! +//! # Architecture +//! +//! ```text +//! GPU VRAM (uploaded once): +//! ┌─────────────────────────────────────────────────────────┐ +//! │ features [total_bars, 42] │ targets [total_bars, 4] │ +//! │ ofi [total_bars, 8] │ (optional) │ +//! └─────────────────────────────────────────────────────────┘ +//! +//! Fold 0: train [0..T0] val [T0..V0] test [V0..E0] +//! Fold 1: train [0..T1] val [T1..V1] test [V1..E1] +//! Fold 2: train [0..T2] val [T2..V2] test [V2..E2] +//! (expanding window) +//! +//! Per-fold kernel launch: +//! episode_starts[] = evenly spaced within [fold.train_start..fold.train_end] +//! total_bars = fold.train_end (kernel can't read past this) +//! ``` +//! +//! No new CUDA kernels required — the existing `dqn_full_experience_kernel` +//! naturally restricts to any contiguous subrange via `episode_starts` and +//! `total_bars`. + +use candle_core::cuda_backend::cudarc::driver::CudaSlice; +use candle_core::Device; +use tracing::info; + +use crate::features::extraction::FeatureVector; +use crate::MLError; + +/// A walk-forward fold's index range into the global GPU buffer. +#[derive(Debug, Clone)] +pub struct GpuFoldRange { + /// Fold index (0-based). + pub fold: usize, + /// First bar index of the training window (inclusive). + pub train_start: usize, + /// Last bar index of the training window (exclusive). + pub train_end: usize, + /// First bar index of the validation window (inclusive). + pub val_start: usize, + /// Last bar index of the validation window (exclusive). + pub val_end: usize, + /// First bar index of the test window (inclusive). + pub test_start: usize, + /// Last bar index of the test window (exclusive). + pub test_end: usize, +} + +impl GpuFoldRange { + /// Number of bars in the training window. + pub fn train_len(&self) -> usize { + self.train_end.saturating_sub(self.train_start) + } + + /// Number of bars in the validation window. + pub fn val_len(&self) -> usize { + self.val_end.saturating_sub(self.val_start) + } + + /// Number of bars in the test window. + pub fn test_len(&self) -> usize { + self.test_end.saturating_sub(self.test_start) + } + + /// Compute episode start indices for the GPU experience collector kernel. + /// + /// Episodes are evenly spaced within `[train_start, train_end - timesteps)`. + /// The kernel reads `timesteps` bars forward from each start, so starts must + /// leave room for the full episode without exceeding `train_end`. + pub fn episode_starts(&self, n_episodes: usize, timesteps: usize) -> Vec { + let usable = self.train_len().saturating_sub(timesteps).max(1); + let stride = (usable / n_episodes.max(1)).max(1); + (0..n_episodes) + .map(|i| { + let offset = (i * stride) % usable; + (self.train_start + offset) as i32 + }) + .collect() + } +} + +/// Walk-forward configuration for GPU-resident data. +#[derive(Debug, Clone)] +pub struct GpuWalkForwardConfig { + /// Fraction of total data for the initial training window (default: 0.5). + pub initial_train_fraction: f64, + /// Fraction of total data per validation window (default: 0.125). + pub val_fraction: f64, + /// Fraction of total data per test window (default: 0.125). + pub test_fraction: f64, + /// Step size between folds as a fraction of total data (default: 0.125). + pub step_fraction: f64, +} + +impl Default for GpuWalkForwardConfig { + fn default() -> Self { + Self { + initial_train_fraction: 0.5, + val_fraction: 0.125, + test_fraction: 0.125, + step_fraction: 0.125, + } + } +} + +impl GpuWalkForwardConfig { + /// Generate fold ranges for the given total number of bars. + /// + /// Returns an empty `Vec` if the data is too small for even one fold. + pub fn generate_folds(&self, total_bars: usize) -> Vec { + let initial_train = (total_bars as f64 * self.initial_train_fraction) as usize; + let val_size = (total_bars as f64 * self.val_fraction) as usize; + let test_size = (total_bars as f64 * self.test_fraction) as usize; + let step_size = (total_bars as f64 * self.step_fraction).max(1.0) as usize; + let min_split = 50; // minimum bars per split + + if initial_train < min_split || val_size < min_split || test_size < min_split { + return Vec::new(); + } + + let mut folds = Vec::new(); + let mut fold_idx = 0; + + loop { + let train_end = initial_train + fold_idx * step_size; + let val_start = train_end; + let val_end = val_start + val_size; + let test_start = val_end; + let test_end = test_start + test_size; + + if test_end > total_bars { + break; + } + + folds.push(GpuFoldRange { + fold: fold_idx, + train_start: 0, // expanding window: always starts at 0 + train_end, + val_start, + val_end, + test_start, + test_end, + }); + + fold_idx += 1; + } + + folds + } +} + +/// GPU-resident walk-forward data store. +/// +/// Holds the entire dataset on GPU as cudarc `CudaSlice` buffers. +/// Folds are expressed as index ranges into these buffers — no re-upload needed. +// CudaSlice doesn't implement Debug, so manual impl. +impl std::fmt::Debug for GpuWalkForwardData { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GpuWalkForwardData") + .field("total_bars", &self.total_bars) + .field("feature_dim", &self.feature_dim) + .field("target_dim", &self.target_dim) + .field("folds", &self.folds.len()) + .field("vram_bytes", &self.vram_bytes) + .field("has_ofi", &self.ofi.is_some()) + .finish() + } +} + +pub struct GpuWalkForwardData { + /// Market features [total_bars * FEATURE_DIM] (FEATURE_DIM = 42). + pub features: CudaSlice, + /// Target prices [total_bars * TARGET_DIM] (TARGET_DIM = 4). + pub targets: CudaSlice, + /// OFI features [total_bars * 8] (optional, from MBP-10). + pub ofi: Option>, + /// Total number of bars in the dataset. + pub total_bars: usize, + /// Feature dimension (42). + pub feature_dim: usize, + /// Target dimension (4). + pub target_dim: usize, + /// Generated fold ranges. + pub folds: Vec, + /// VRAM usage in bytes. + pub vram_bytes: usize, +} + +impl GpuWalkForwardData { + /// Upload the entire dataset to GPU VRAM and generate walk-forward folds. + /// + /// On H100 (80 GB), a typical dataset of 1M bars uses: + /// features: 1M × 42 × 4 = 168 MB + /// targets: 1M × 4 × 4 = 16 MB + /// ofi: 1M × 8 × 4 = 32 MB + /// total: 216 MB (~0.27% of VRAM) + pub fn upload( + training_data: &[(FeatureVector, Vec)], + ofi_data: Option<&[[f64; 8]]>, + wf_config: &GpuWalkForwardConfig, + device: &Device, + ) -> Result { + let candle_core::Device::Cuda(cuda_dev) = device else { + return Err(MLError::ModelError( + "GpuWalkForwardData requires CUDA device".to_owned(), + )); + }; + + let total_bars = training_data.len(); + if total_bars == 0 { + return Err(MLError::ModelError("Empty training data".to_owned())); + } + + let feature_dim = 42; + let target_dim = 4; + let stream = cuda_dev.cuda_stream(); + + // Flatten features: [total_bars * 42] + let mut flat_features = Vec::with_capacity(total_bars * feature_dim); + for (features, _) in training_data { + for &v in features.iter() { + flat_features.push(v as f32); + } + } + + // Flatten targets: [total_bars * 4] + let mut flat_targets = Vec::with_capacity(total_bars * target_dim); + for (_, targets) in training_data { + for i in 0..target_dim { + flat_targets.push(targets.get(i).copied().unwrap_or(0.0) as f32); + } + } + + // Upload features + let features_gpu = stream + .memcpy_stod(&flat_features) + .map_err(|e| MLError::ModelError(format!("WF features upload: {e}")))?; + + let mut vram = total_bars * feature_dim * 4; + + // Upload targets + let targets_gpu = stream + .memcpy_stod(&flat_targets) + .map_err(|e| MLError::ModelError(format!("WF targets upload: {e}")))?; + + vram += total_bars * target_dim * 4; + + // Upload OFI (optional) + let ofi_gpu = if let Some(ofi) = ofi_data { + let n = ofi.len().min(total_bars); + let mut flat_ofi = Vec::with_capacity(total_bars * 8); + for i in 0..total_bars { + if i < n { + for &v in ofi[i].iter() { + flat_ofi.push(v as f32); + } + } else { + flat_ofi.extend_from_slice(&[0.0_f32; 8]); + } + } + let buf = stream + .memcpy_stod(&flat_ofi) + .map_err(|e| MLError::ModelError(format!("WF OFI upload: {e}")))?; + vram += total_bars * 8 * 4; + Some(buf) + } else { + None + }; + + // Generate walk-forward folds + let folds = wf_config.generate_folds(total_bars); + + info!( + "GPU walk-forward data uploaded: {} bars, {} folds, {:.1} MB VRAM", + total_bars, + folds.len(), + vram as f64 / 1_048_576.0, + ); + + for fold in &folds { + info!( + " Fold {}: train [{}..{}] ({} bars), val [{}..{}] ({} bars), test [{}..{}] ({} bars)", + fold.fold, + fold.train_start, fold.train_end, fold.train_len(), + fold.val_start, fold.val_end, fold.val_len(), + fold.test_start, fold.test_end, fold.test_len(), + ); + } + + Ok(Self { + features: features_gpu, + targets: targets_gpu, + ofi: ofi_gpu, + total_bars, + feature_dim, + target_dim, + folds, + vram_bytes: vram, + }) + } + + /// Get the number of folds. + pub fn num_folds(&self) -> usize { + self.folds.len() + } + + /// Get a fold's range by index. + pub fn fold(&self, idx: usize) -> Option<&GpuFoldRange> { + self.folds.get(idx) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_fold_generation_basic() { + let config = GpuWalkForwardConfig::default(); + // 1000 bars: initial_train=500, val=125, test=125, step=125 + let folds = config.generate_folds(1000); + + assert!(!folds.is_empty(), "should generate at least one fold"); + + let f0 = &folds[0]; + assert_eq!(f0.train_start, 0); + assert_eq!(f0.train_end, 500); + assert_eq!(f0.val_start, 500); + assert_eq!(f0.val_end, 625); + assert_eq!(f0.test_start, 625); + assert_eq!(f0.test_end, 750); + + // Check expanding window for fold 1 + if folds.len() > 1 { + let f1 = &folds[1]; + assert_eq!(f1.train_start, 0); + assert_eq!(f1.train_end, 625); // expanded by step_size=125 + assert!(f1.train_len() > f0.train_len()); + } + } + + #[test] + fn test_fold_generation_insufficient_data() { + let config = GpuWalkForwardConfig::default(); + // 100 bars: initial_train=50, val=12, test=12 — val/test too small (< 50) + let folds = config.generate_folds(100); + assert!(folds.is_empty(), "100 bars too few for default config"); + } + + #[test] + fn test_episode_starts_within_range() { + let fold = GpuFoldRange { + fold: 0, + train_start: 0, + train_end: 500, + val_start: 500, + val_end: 625, + test_start: 625, + test_end: 750, + }; + let timesteps = 100; + let n_episodes = 64; + let starts = fold.episode_starts(n_episodes, timesteps); + + assert_eq!(starts.len(), n_episodes); + for &s in &starts { + assert!(s >= 0, "start must be non-negative"); + assert!( + (s as usize) + timesteps <= fold.train_end, + "episode must not exceed train window: start={s} + timesteps={timesteps} > train_end={}", + fold.train_end, + ); + } + } + + #[test] + fn test_fold_generation_many_bars() { + let config = GpuWalkForwardConfig { + initial_train_fraction: 0.5, + val_fraction: 0.1, + test_fraction: 0.1, + step_fraction: 0.1, + }; + // 10000 bars: should generate multiple folds + let folds = config.generate_folds(10000); + assert!(folds.len() >= 3, "should generate >=3 folds for 10K bars"); + + // Each successive fold has strictly more training data + for i in 1..folds.len() { + assert!( + folds[i].train_len() > folds[i - 1].train_len(), + "fold {} train ({}) should exceed fold {} train ({})", + i, + folds[i].train_len(), + i - 1, + folds[i - 1].train_len(), + ); + } + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_weights.rs b/crates/ml/src/cuda_pipeline/gpu_weights.rs index 475fe3b30..dce2769b1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_weights.rs +++ b/crates/ml/src/cuda_pipeline/gpu_weights.rs @@ -10,6 +10,8 @@ //! //! Weights are stored row-major `[out_features, in_features]` matching the //! Candle convention. The CUDA kernel reads them in the same layout. +// CUDA FFI module — memcpy_dtod_async requires unsafe by design. +#![allow(unsafe_code)] use std::sync::Arc; @@ -498,8 +500,30 @@ impl KernelWeightPack { // Extraction helpers // --------------------------------------------------------------------------- +/// Async device-to-device memcpy with error context. +/// +/// Wraps the single unavoidable `unsafe` call (`memcpy_dtod_async`) in a +/// bounds-documented helper so callers remain safe. +fn dtod_copy_checked( + dst: cudarc::driver::sys::CUdeviceptr, + src: cudarc::driver::sys::CUdeviceptr, + num_bytes: usize, + stream: &Arc, + name: &str, + op: &str, +) -> Result<(), MLError> { + // Safety: caller guarantees src/dst are valid device pointers within + // the same CUDA context, num_bytes ≤ allocation size of both buffers, + // and stream belongs to the same context. + unsafe { + cudarc::driver::result::memcpy_dtod_async(dst, src, num_bytes, stream.cu_stream()) + .map_err(|e| MLError::ModelError(format!("DtoD {op} {name}: {e}")))?; + } + Ok(()) +} + /// Extract a single named weight tensor from locked VarMap data, flatten it, -/// allocate a GPU buffer, and upload. +/// allocate a GPU buffer, and copy via device-to-device (zero CPU roundtrip). fn extract_one( vars_data: &std::collections::HashMap, name: &str, @@ -509,23 +533,39 @@ fn extract_one( .get(name) .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))? .as_tensor(); - let flat: Vec = tensor + let flat_tensor = tensor .flatten_all() .map_err(|e| MLError::ModelError(format!("Flatten {name}: {e}")))? .to_dtype(candle_core::DType::F32) .map_err(|e| MLError::ModelError(format!("Cast {name} to F32: {e}")))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("to_vec1 {name}: {e}")))?; - let mut buf = stream - .alloc_zeros::(flat.len()) + .contiguous() + .map_err(|e| MLError::ModelError(format!("Contiguous {name}: {e}")))?; + let n_elems = flat_tensor.elem_count(); + let buf = stream + .alloc_zeros::(n_elems) .map_err(|e| MLError::ModelError(format!("Alloc {name}: {e}")))?; - stream - .memcpy_htod(&flat, &mut buf) - .map_err(|e| MLError::ModelError(format!("Upload {name}: {e}")))?; + + let (storage_guard, _layout) = flat_tensor.storage_and_layout(); + match *storage_guard { + candle_core::Storage::Cuda(ref cs) => { + let src_slice: &CudaSlice = cs.as_cuda_slice() + .map_err(|e| MLError::ModelError(format!("as_cuda_slice {name}: {e}")))?; + let (src_ptr, _src_sync) = src_slice.device_ptr(stream); + let (dst_ptr, _dst_sync) = buf.device_ptr(stream); + let num_bytes = n_elems * std::mem::size_of::(); + dtod_copy_checked(dst_ptr, src_ptr, num_bytes, stream, name, "copy")?; + } + candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { + return Err(MLError::ModelError(format!( + "Weight tensor '{name}' is not on CUDA — GPU weight extraction requires CUDA tensors" + ))); + } + } Ok(buf) } /// Re-upload a single weight tensor into an existing GPU buffer (no realloc). +/// Uses device-to-device copy (zero CPU roundtrip). fn sync_one( vars_data: &std::collections::HashMap, name: &str, @@ -536,16 +576,31 @@ fn sync_one( .get(name) .ok_or_else(|| MLError::ModelError(format!("Missing weight: {name}")))? .as_tensor(); - let flat: Vec = tensor + let flat_tensor = tensor .flatten_all() .map_err(|e| MLError::ModelError(format!("Flatten {name}: {e}")))? .to_dtype(candle_core::DType::F32) .map_err(|e| MLError::ModelError(format!("Cast {name} to F32: {e}")))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("to_vec1 {name}: {e}")))?; - stream - .memcpy_htod(&flat, buf) - .map_err(|e| MLError::ModelError(format!("Upload {name}: {e}")))?; + .contiguous() + .map_err(|e| MLError::ModelError(format!("Contiguous {name}: {e}")))?; + let n_elems = flat_tensor.elem_count(); + + let (storage_guard, _layout) = flat_tensor.storage_and_layout(); + match *storage_guard { + candle_core::Storage::Cuda(ref cs) => { + let src_slice: &CudaSlice = cs.as_cuda_slice() + .map_err(|e| MLError::ModelError(format!("as_cuda_slice {name}: {e}")))?; + let (src_ptr, _src_sync) = src_slice.device_ptr(stream); + let (dst_ptr, _dst_sync) = buf.device_ptr(stream); + let num_bytes = n_elems * std::mem::size_of::(); + dtod_copy_checked(dst_ptr, src_ptr, num_bytes, stream, name, "sync")?; + } + candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { + return Err(MLError::ModelError(format!( + "Weight tensor '{name}' is not on CUDA — GPU weight sync requires CUDA tensors" + ))); + } + } Ok(()) } @@ -1089,11 +1144,16 @@ mod tests { use crate::ppo::ppo::{PPOConfig, PPO}; use candle_core::{Device, DType}; + fn cuda_device() -> Device { + Device::new_cuda(0).expect("CUDA device required") + } + /// Verify that all 12 expected VarMap key paths exist in a DuelingQNetwork. + #[cfg_attr(not(feature = "cuda"), ignore)] #[test] fn test_dueling_weight_key_paths() { let config = DuelingConfig::new(54, 5, vec![256, 256], 128, 128); - let network = DuelingQNetwork::new(config, Device::Cpu) + let network = DuelingQNetwork::new(config, cuda_device()) .expect("network creation should not fail on CPU"); let vars_data = network.vars().data().lock() @@ -1135,10 +1195,11 @@ mod tests { } /// Verify expected parameter counts match Dueling DQN spec. + #[cfg_attr(not(feature = "cuda"), ignore)] #[test] fn test_dueling_param_count() { let config = DuelingConfig::new(54, 5, vec![256, 256], 128, 128); - let network = DuelingQNetwork::new(config, Device::Cpu) + let network = DuelingQNetwork::new(config, cuda_device()) .expect("network creation should not fail"); let vars_data = network.vars().data().lock() @@ -1159,6 +1220,7 @@ mod tests { } /// Verify that all 6 expected VarMap key paths exist in a PPO PolicyNetwork. + #[cfg_attr(not(feature = "cuda"), ignore)] #[test] fn test_ppo_actor_weight_key_paths() { let config = PPOConfig { @@ -1168,7 +1230,7 @@ mod tests { value_hidden_dims: vec![512, 384, 256, 128, 64], ..PPOConfig::default() }; - let ppo = PPO::with_device(config, Device::Cpu) + let ppo = PPO::with_device(config, cuda_device()) .expect("PPO creation should not fail on CPU"); let vars_data = ppo.actor.vars().data().lock() @@ -1186,7 +1248,9 @@ mod tests { .expect("policy_layer_0.weight must exist") .as_tensor(); assert_eq!(l0w.dims(), &[128, 48], "policy_layer_0.weight shape"); - assert_eq!(l0w.dtype(), DType::F32); + // Weights use training_dtype: BF16 on CUDA, F32 on CPU + let expected_dtype = ml_core::mixed_precision::training_dtype(&cuda_device()); + assert_eq!(l0w.dtype(), expected_dtype); let l1w = vars_data.get("policy_layer_1.weight") .expect("policy_layer_1.weight must exist") @@ -1200,6 +1264,7 @@ mod tests { } /// Verify that all 12 expected VarMap key paths exist in a PPO ValueNetwork. + #[cfg_attr(not(feature = "cuda"), ignore)] #[test] fn test_ppo_critic_weight_key_paths() { let config = PPOConfig { @@ -1209,7 +1274,7 @@ mod tests { value_hidden_dims: vec![512, 384, 256, 128, 64], ..PPOConfig::default() }; - let ppo = PPO::with_device(config, Device::Cpu) + let ppo = PPO::with_device(config, cuda_device()) .expect("PPO creation should not fail on CPU"); let vars_data = ppo.critic.vars().data().lock() @@ -1227,7 +1292,9 @@ mod tests { .expect("value_layer_0.weight must exist") .as_tensor(); assert_eq!(l0w.dims(), &[512, 48], "value_layer_0.weight shape"); - assert_eq!(l0w.dtype(), DType::F32); + // Weights use training_dtype: BF16 on CUDA, F32 on CPU + let expected_dtype = ml_core::mixed_precision::training_dtype(&cuda_device()); + assert_eq!(l0w.dtype(), expected_dtype); let l4w = vars_data.get("value_layer_4.weight") .expect("value_layer_4.weight must exist") @@ -1241,6 +1308,7 @@ mod tests { } /// Verify PPO actor total parameter count is in expected range. + #[cfg_attr(not(feature = "cuda"), ignore)] #[test] fn test_ppo_actor_param_count() { let config = PPOConfig { @@ -1250,7 +1318,7 @@ mod tests { value_hidden_dims: vec![512, 384, 256, 128, 64], ..PPOConfig::default() }; - let ppo = PPO::with_device(config, Device::Cpu) + let ppo = PPO::with_device(config, cuda_device()) .expect("PPO creation should not fail"); let vars_data = ppo.actor.vars().data().lock() @@ -1271,6 +1339,7 @@ mod tests { } /// Verify PPO critic total parameter count is in expected range. + #[cfg_attr(not(feature = "cuda"), ignore)] #[test] fn test_ppo_critic_param_count() { let config = PPOConfig { @@ -1280,7 +1349,7 @@ mod tests { value_hidden_dims: vec![512, 384, 256, 128, 64], ..PPOConfig::default() }; - let ppo = PPO::with_device(config, Device::Cpu) + let ppo = PPO::with_device(config, cuda_device()) .expect("PPO creation should not fail"); let vars_data = ppo.critic.vars().data().lock() diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 8595073ae..0c08a9334 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -33,6 +33,8 @@ pub mod gpu_monitoring; pub mod gpu_backtest_evaluator; #[cfg(feature = "cuda")] pub mod gpu_curiosity_trainer; +#[cfg(feature = "cuda")] +pub mod gpu_walk_forward; // gpu_replay_buffer moved to ml-dqn crate /// Maximum bytes allowed for a single GPU upload (2 GB safety limit). @@ -99,6 +101,12 @@ pub fn optimal_launch_dims(n_items: u32, max_threads_per_block: u32) -> (u32, u3 /// /// # Feature gate /// Only available with the `cuda` feature. +/// +/// # Compilation strategy +/// +/// Uses nvcc (offline compiler) with full `-O3` optimization. nvcc produces +/// higher-quality PTX than NVRTC and supports all optimization flags. +/// Requires nvcc in `$CUDA_HOME/bin/` or `$PATH`. #[cfg(feature = "cuda")] pub fn compile_ptx_for_device( src: &str, @@ -113,9 +121,7 @@ pub fn compile_ptx_for_device( .attribute(CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR) .unwrap_or(0); - // Build arch string: "compute_90", "compute_80", etc. - // NVRTC accepts the compute_XX form (virtual architecture) which enables - // __CUDA_ARCH__ and arch-specific code generation. + // Build arch string for nvcc: "compute_XX" (virtual arch for PTX output). let arch_str: &'static str = match (major, minor) { (9, 0) => "compute_90", (9, _) => "compute_90", // sm_90a and beyond @@ -124,11 +130,7 @@ pub fn compile_ptx_for_device( (8, 0) => "compute_80", (7, 5) => "compute_75", (7, 0) => "compute_70", - _ => { - // Unknown or very old/new — use default NVRTC behavior (no caching). - return cudarc::nvrtc::compile_ptx(src) - .map_err(|e| e.to_string()); - } + _ => "compute_80", // safe default }; // Try loading from PTX cache first @@ -136,13 +138,8 @@ pub fn compile_ptx_for_device( return Ok(cached); } - // Cache miss — compile via NVRTC - let opts = cudarc::nvrtc::CompileOptions { - arch: Some(arch_str), - ..Default::default() - }; - let ptx = cudarc::nvrtc::compile_ptx_with_opts(src, opts) - .map_err(|e| e.to_string())?; + // Cache miss — compile via nvcc with full -O3 optimization + let ptx = compile_with_nvcc(src, arch_str)?; // Cache the compiled PTX for future runs save_ptx_to_cache(arch_str, src, &ptx); @@ -150,6 +147,107 @@ pub fn compile_ptx_for_device( Ok(ptx) } +/// nvcc optimization flags applied to every kernel compilation. +/// +/// Changing this array auto-invalidates the PTX 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. +#[cfg(feature = "cuda")] +const NVCC_OPT_FLAGS: &[&str] = &[ + "-O3", + "--use_fast_math", + "--ftz=true", + "--fmad=true", +]; + +/// Compile CUDA source to PTX using nvcc with maximum optimization. +/// +/// See [`NVCC_OPT_FLAGS`] for the flags applied. `-ptx` and `-arch=compute_XX` +/// are added automatically. +#[cfg(feature = "cuda")] +fn compile_with_nvcc( + src: &str, + arch: &str, +) -> Result { + 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}"))?; + + // Use content hash for temp file name to avoid collisions + let src_hash = { + use sha2::{Sha256, Digest}; + let mut h = Sha256::new(); + h.update(arch.as_bytes()); + h.update(src.as_bytes()); + format!("{:x}", h.finalize()) + }; + let hash_prefix = src_hash.get(..16).unwrap_or(&src_hash); + let src_path = tmp_dir.join(format!("{hash_prefix}.cu")); + let ptx_path = tmp_dir.join(format!("{hash_prefix}.ptx")); + + { + 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 optimization flags from NVCC_OPT_FLAGS constant. + let arch_flag = format!("-arch={arch}"); + let out_path = ptx_path.to_str().unwrap_or("out.ptx"); + let in_path = src_path.to_str().unwrap_or("in.cu"); + let mut args: Vec<&str> = vec!["-ptx"]; + args.extend_from_slice(NVCC_OPT_FLAGS); + args.extend_from_slice(&[&arch_flag, "-o", out_path, 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) + )); + } + + // Read the compiled PTX + let ptx_text = std::fs::read_to_string(&ptx_path) + .map_err(|e| format!("Failed to read compiled PTX from {}: {e}", ptx_path.display()))?; + + // Clean up temp files (best-effort) + drop(std::fs::remove_file(&src_path)); + drop(std::fs::remove_file(&ptx_path)); + + let flags_str = NVCC_OPT_FLAGS.join(" "); + tracing::info!( + arch, + src_len = src.len(), + ptx_len = ptx_text.len(), + flags = %flags_str, + "nvcc compilation succeeded" + ); + + Ok(candle_core::cuda_backend::cudarc::nvrtc::Ptx::from_src(ptx_text)) +} + /// Resolve the PTX cache directory. /// /// Prefers `$CARGO_TARGET_DIR/.ptx_cache/` (persisted on CI PVC between runs). @@ -172,6 +270,12 @@ fn ptx_cache_key(arch: &str, src: &str) -> String { 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 PTX. + // No manual version bump needed — derived from NVCC_OPT_FLAGS constant. + for flag in NVCC_OPT_FLAGS { + h.update(flag.as_bytes()); + h.update(b"\x00"); + } h.update(src.as_bytes()); format!("{:x}", h.finalize()) } @@ -398,19 +502,33 @@ impl DqnGpuData { /// /// Returns [current_close_preproc, next_close_preproc, current_close_raw, next_close_raw]. pub fn bar_target_values(&self, bar_idx: usize) -> Result<[f32; 4], MLError> { - let slice = self.bar_targets(bar_idx)? + let flat = self.bar_targets(bar_idx)? .flatten_all() .map_err(|e| MLError::ModelError(format!("Target flatten failed: {e}")))? .to_dtype(candle_core::DType::F32) - .map_err(|e| MLError::ModelError(format!("Target cast to F32 failed: {e}")))? - .to_vec1::() - .map_err(|e| MLError::ModelError(format!("Target to_vec1 failed: {e}")))?; - Ok([ - slice.get(0).copied().unwrap_or(0.0), - slice.get(1).copied().unwrap_or(0.0), - slice.get(2).copied().unwrap_or(0.0), - slice.get(3).copied().unwrap_or(0.0), - ]) + .map_err(|e| MLError::ModelError(format!("Target cast to F32 failed: {e}")))?; + // narrow returns rank-1 [1]; squeeze(0) reduces to rank-0 scalar for to_scalar() + let v0 = flat.narrow(0, 0, 1) + .and_then(|t| t.squeeze(0)) + .map_err(|e| MLError::ModelError(format!("Target narrow 0: {e}")))? + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Target scalar 0: {e}")))?; + let v1 = flat.narrow(0, 1, 1) + .and_then(|t| t.squeeze(0)) + .map_err(|e| MLError::ModelError(format!("Target narrow 1: {e}")))? + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Target scalar 1: {e}")))?; + let v2 = flat.narrow(0, 2, 1) + .and_then(|t| t.squeeze(0)) + .map_err(|e| MLError::ModelError(format!("Target narrow 2: {e}")))? + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Target scalar 2: {e}")))?; + let v3 = flat.narrow(0, 3, 1) + .and_then(|t| t.squeeze(0)) + .map_err(|e| MLError::ModelError(format!("Target narrow 3: {e}")))? + .to_scalar::() + .map_err(|e| MLError::ModelError(format!("Target scalar 3: {e}")))?; + Ok([v0, v1, v2, v3]) } /// Build a complete state tensor by concatenating pre-uploaded market features @@ -727,9 +845,24 @@ impl PpoGpuData { #[cfg(test)] mod tests { use super::*; + use std::sync::OnceLock; + + /// Shared CUDA device for all cuda_pipeline tests. + /// Reuses one cuBLAS handle instead of creating a new one per test, + /// preventing handle pool exhaustion across test binaries. + static SHARED_CUDA: OnceLock = OnceLock::new(); + + fn cuda_device() -> Device { + SHARED_CUDA + .get_or_init(|| { + Device::new_cuda(0) + .expect("CUDA device required — no CPU fallback in cuda_pipeline tests") + }) + .clone() + } #[test] - fn test_dqn_gpu_data_upload_cpu() { + fn test_dqn_gpu_data_upload() { let data: Vec<([f64; 42], Vec)> = (0..100) .map(|i| { let features = [i as f64 * 0.01; 42]; @@ -738,7 +871,7 @@ mod tests { }) .collect(); - let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap(); + let gpu_data = DqnGpuData::upload(&data, &cuda_device()).unwrap(); assert_eq!(gpu_data.num_bars, 100); assert_eq!(gpu_data.feature_dim, 42); @@ -758,9 +891,9 @@ mod tests { ([1.0; 42], vec![100.0, 101.0, 100.5, 101.5]), ]; - let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap(); + let gpu_data = DqnGpuData::upload(&data, &cuda_device()).unwrap(); let portfolio = [0.95_f32, 0.5, 0.0001]; - let state = gpu_data.build_state_tensor(0, &portfolio, &Device::Cpu).unwrap(); + let state = gpu_data.build_state_tensor(0, &portfolio, &cuda_device()).unwrap(); assert_eq!(state.dims(), &[1, 45]); } @@ -808,12 +941,12 @@ mod tests { } #[test] - fn test_ppo_gpu_data_upload_cpu() { + fn test_ppo_gpu_data_upload() { let market_data: Vec> = (0..200) .map(|i| vec![i as f32 * 0.01; 45]) .collect(); - let gpu_data = PpoGpuData::upload(&market_data, &Device::Cpu).unwrap(); + let gpu_data = PpoGpuData::upload(&market_data, &cuda_device()).unwrap(); assert_eq!(gpu_data.num_steps, 200); assert_eq!(gpu_data.state_dim, 45); @@ -831,21 +964,21 @@ mod tests { vec![2.0_f32; 42], ]; - let result = PpoGpuData::upload(&market_data, &Device::Cpu); + let result = PpoGpuData::upload(&market_data, &cuda_device()); assert!(result.is_err()); } #[test] fn test_dqn_gpu_data_empty() { let data: Vec<([f64; 42], Vec)> = vec![]; - let result = DqnGpuData::upload(&data, &Device::Cpu); + let result = DqnGpuData::upload(&data, &cuda_device()); assert!(result.is_err()); } #[test] fn test_ppo_gpu_data_empty() { let data: Vec> = vec![]; - let result = PpoGpuData::upload(&data, &Device::Cpu); + let result = PpoGpuData::upload(&data, &cuda_device()); assert!(result.is_err()); } @@ -862,7 +995,7 @@ mod tests { let data: Vec<([f64; 42], Vec)> = vec![ ([0.5; 42], vec![100.0, 101.0, 100.5, 101.5]), ]; - let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap(); + let gpu_data = DqnGpuData::upload(&data, &cuda_device()).unwrap(); assert!(gpu_data.bar_features(0).is_ok()); assert!(gpu_data.bar_features(1).is_err()); @@ -881,24 +1014,42 @@ mod tests { }) .collect(); - let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap(); + let gpu_data = DqnGpuData::upload(&data, &cuda_device()).unwrap(); let portfolio = [0.95_f32, 0.5, 0.0001]; // Build batch of 32 states starting at bar 10 - let batch = gpu_data.build_batch_states(10, 32, &portfolio, &Device::Cpu).unwrap(); + let batch = gpu_data.build_batch_states(10, 32, &portfolio, &cuda_device()).unwrap(); assert_eq!(batch.dims(), &[32, 45]); // 42 market + 3 portfolio - // Verify portfolio features are broadcast correctly - let row0 = batch.narrow(0, 0, 1).unwrap().flatten_all().unwrap().to_vec1::().unwrap(); - assert!((row0[42] - 0.95).abs() < 1e-5, "portfolio_value"); - assert!((row0[43] - 0.5).abs() < 1e-5, "position"); - assert!((row0[44] - 0.0001).abs() < 1e-5, "spread"); + // Verify portfolio features are broadcast correctly (GPU-side comparison) + let row0_portfolio = batch.narrow(0, 0, 1).unwrap() + .narrow(1, 42, 3).unwrap().flatten_all().unwrap(); + let expected = Tensor::new(&[0.95_f32, 0.5, 0.0001], batch.device()).unwrap() + .to_dtype(batch.dtype()).unwrap(); + let diff = row0_portfolio.sub(&expected).unwrap().abs().unwrap(); + let max_diff = diff.max(0).unwrap() + .to_dtype(candle_core::DType::F32).unwrap() + .to_scalar::().unwrap(); + assert!(max_diff < 1e-2, "portfolio features mismatch: max_diff={max_diff}"); - let row31 = batch.narrow(0, 31, 1).unwrap().flatten_all().unwrap().to_vec1::().unwrap(); - assert!((row31[42] - 0.95).abs() < 1e-5, "portfolio_value row31"); + let row31_pv = batch.narrow(0, 31, 1).unwrap() + .narrow(1, 42, 1).unwrap().flatten_all().unwrap(); + let expected_pv = Tensor::new(&[0.95_f32], batch.device()).unwrap() + .to_dtype(batch.dtype()).unwrap(); + let pv_diff = row31_pv.sub(&expected_pv).unwrap().abs().unwrap() + .max(0).unwrap() + .to_dtype(candle_core::DType::F32).unwrap() + .to_scalar::().unwrap(); + assert!(pv_diff < 1e-2, "portfolio_value row31 mismatch: diff={pv_diff}"); - // Verify market features differ between rows - assert!((row0[0] - row31[0]).abs() > 1e-6, "market features should differ"); + // Verify market features differ between rows (GPU-side) + let row0_market = batch.narrow(0, 0, 1).unwrap().narrow(1, 0, 1).unwrap().flatten_all().unwrap(); + let row31_market = batch.narrow(0, 31, 1).unwrap().narrow(1, 0, 1).unwrap().flatten_all().unwrap(); + let market_diff = row0_market.sub(&row31_market).unwrap().abs().unwrap() + .max(0).unwrap() + .to_dtype(candle_core::DType::F32).unwrap() + .to_scalar::().unwrap(); + assert!(market_diff > 1e-6, "market features should differ"); } #[test] @@ -907,15 +1058,15 @@ mod tests { ([1.0; 42], vec![100.0, 101.0, 100.5, 101.5]), ([2.0; 42], vec![100.0, 101.0, 100.5, 101.5]), ]; - let gpu_data = DqnGpuData::upload(&data, &Device::Cpu).unwrap(); + let gpu_data = DqnGpuData::upload(&data, &cuda_device()).unwrap(); let portfolio = [1.0_f32, 0.0, 0.0001]; // Request more bars than available — should clamp - let batch = gpu_data.build_batch_states(0, 100, &portfolio, &Device::Cpu).unwrap(); + let batch = gpu_data.build_batch_states(0, 100, &portfolio, &cuda_device()).unwrap(); assert_eq!(batch.dims(), &[2, 45]); // Clamped to 2 bars // Start beyond range — should error - let err = gpu_data.build_batch_states(5, 10, &portfolio, &Device::Cpu); + let err = gpu_data.build_batch_states(5, 10, &portfolio, &cuda_device()); assert!(err.is_err()); } @@ -929,7 +1080,7 @@ mod tests { .map(|i| ([i as f64 * 0.01; 42], vec![100.0, 101.0, 100.5, 101.5])) .collect(); - let gpu_data = pool.upload_dqn(&data, &Device::Cpu).unwrap(); + let gpu_data = pool.upload_dqn(&data, &cuda_device()).unwrap(); assert_eq!(gpu_data.num_bars, 50); assert_eq!(gpu_data.feature_dim, 42); } @@ -942,14 +1093,14 @@ mod tests { let data1: Vec<([f64; 42], Vec)> = (0..100) .map(|i| ([i as f64 * 0.01; 42], vec![100.0, 101.0, 100.5, 101.5])) .collect(); - let g1 = pool.upload_dqn(&data1, &Device::Cpu).unwrap(); + let g1 = pool.upload_dqn(&data1, &cuda_device()).unwrap(); assert_eq!(g1.num_bars, 100); // Second upload — reuses the same staging buffers let data2: Vec<([f64; 42], Vec)> = (0..200) .map(|i| ([i as f64 * 0.02; 42], vec![200.0, 201.0, 200.5, 201.5])) .collect(); - let g2 = pool.upload_dqn(&data2, &Device::Cpu).unwrap(); + let g2 = pool.upload_dqn(&data2, &cuda_device()).unwrap(); assert_eq!(g2.num_bars, 200); // Verify data integrity: g2 should have data2's values @@ -965,7 +1116,7 @@ mod tests { let data: Vec<([f64; 42], Vec)> = (0..50) .map(|i| ([i as f64; 42], vec![1.0, 2.0, 3.0, 4.0])) .collect(); - let gpu_data = pool.upload_dqn(&data, &Device::Cpu).unwrap(); + let gpu_data = pool.upload_dqn(&data, &cuda_device()).unwrap(); assert_eq!(gpu_data.num_bars, 50); // Staging buffers should have grown assert!(pool.feature_buf.len() >= 50 * 42); @@ -975,7 +1126,7 @@ mod tests { fn test_gpu_buffer_pool_empty() { let mut pool = GpuBufferPool::new(100, 42, 4); let data: Vec<([f64; 42], Vec)> = vec![]; - assert!(pool.upload_dqn(&data, &Device::Cpu).is_err()); + assert!(pool.upload_dqn(&data, &cuda_device()).is_err()); } #[test] diff --git a/crates/ml/src/cuda_pipeline/monitoring_kernel.cu b/crates/ml/src/cuda_pipeline/monitoring_kernel.cu index 6bd06aea5..4d35073a8 100644 --- a/crates/ml/src/cuda_pipeline/monitoring_kernel.cu +++ b/crates/ml/src/cuda_pipeline/monitoring_kernel.cu @@ -61,12 +61,23 @@ extern "C" __global__ void monitoring_reduce( if (a >= 0 && a < 5) local_counts[a]++; } - // Warp reduction then atomic to shared - atomicAdd(&s_sum, local_sum); - atomicAdd(&s_sq_sum, local_sq); - atomicMin_float(&s_min, local_min); - atomicMax_float(&s_max, local_max); - for (int i = 0; i < 5; i++) atomicAdd(&s_counts[i], local_counts[i]); + // Warp-level reduction before atomic to shared (reduces contention 32x) + for (int mask = 16; mask >= 1; mask >>= 1) { + local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, mask); + local_sq += __shfl_xor_sync(0xFFFFFFFF, local_sq, mask); + local_min = fminf(local_min, __shfl_xor_sync(0xFFFFFFFF, local_min, mask)); + local_max = fmaxf(local_max, __shfl_xor_sync(0xFFFFFFFF, local_max, mask)); + for (int i = 0; i < 5; i++) + local_counts[i] += __shfl_xor_sync(0xFFFFFFFF, local_counts[i], mask); + } + // Only lane 0 of each warp atomicAdd (8 warp leaders instead of 256 threads) + if ((tid & 31) == 0) { + atomicAdd(&s_sum, local_sum); + atomicAdd(&s_sq_sum, local_sq); + atomicMin_float(&s_min, local_min); + atomicMax_float(&s_max, local_max); + for (int i = 0; i < 5; i++) atomicAdd(&s_counts[i], local_counts[i]); + } __syncthreads(); // Thread 0 writes summary diff --git a/crates/ml/src/cuda_pipeline/signal_adapter.rs b/crates/ml/src/cuda_pipeline/signal_adapter.rs index ae1a7d6e6..355c21156 100644 --- a/crates/ml/src/cuda_pipeline/signal_adapter.rs +++ b/crates/ml/src/cuda_pipeline/signal_adapter.rs @@ -300,7 +300,10 @@ pub fn backtest_fitness( #[cfg(test)] mod tests { use super::*; - use candle_core::Device; + + fn cuda_device() -> candle_core::Device { + candle_core::Device::new_cuda(0).expect("CUDA device required") + } // ── ppo_to_exposure_scores ────────────────────────────────────────── @@ -308,17 +311,15 @@ mod tests { fn test_ppo_to_exposure_scores_shape() { let batch = 4; let uniform = vec![1.0_f32 / 45.0; batch * 45]; - let probs = Tensor::from_vec(uniform, &[batch, 45], &Device::Cpu).unwrap(); + let probs = Tensor::from_vec(uniform, &[batch, 45], &cuda_device()).unwrap(); let scores = ppo_to_exposure_scores(&probs).unwrap(); assert_eq!(scores.dims(), &[batch, 5]); - let data = scores.to_vec2::().unwrap(); - for row in &data { - for &v in row { - // Each of 5 bins sums 9 cells of 1/45 ≈ 0.2 - assert!((v - 0.2).abs() < 1e-5, "expected ~0.2, got {v}"); - } - } + // GPU-side: all values should be ~0.2 (each of 5 bins sums 9 cells of 1/45) + let expected = Tensor::new(0.2_f32, scores.device()).unwrap(); + let max_diff = scores.broadcast_sub(&expected).unwrap().abs().unwrap() + .max(0).unwrap().max(0).unwrap().to_scalar::().unwrap(); + assert!(max_diff < 1e-5, "expected ~0.2 everywhere, max_diff={max_diff}"); } #[test] @@ -334,73 +335,67 @@ mod tests { } } } - let probs = Tensor::from_vec(raw, &[batch, 45], &Device::Cpu).unwrap(); + let probs = Tensor::from_vec(raw, &[batch, 45], &cuda_device()).unwrap(); let scores = ppo_to_exposure_scores(&probs).unwrap(); - let data = scores.to_vec2::().unwrap(); - for row in &data { - let argmax = row - .iter() - .enumerate() - .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(i, _)| i) - .unwrap(); - assert_eq!(argmax, 4, "expected Long100 (4), got {argmax}"); - } + // GPU-side argmax: all rows should have argmax at index 4 (Long100) + let argmaxes = scores.argmax(1).unwrap(); + let expected_4 = Tensor::new(&[4_u32, 4_u32], scores.device()).unwrap(); + let all_match = argmaxes.eq(&expected_4).unwrap() + .to_dtype(DType::U32).unwrap() + .sum_all().unwrap().to_scalar::().unwrap(); + assert_eq!(all_match, batch as u32, "all rows should argmax to Long100 (4)"); } #[test] fn test_ppo_to_exposure_scores_rejects_wrong_shape() { - let probs = Tensor::zeros(&[8, 5], DType::F32, &Device::Cpu).unwrap(); + let probs = Tensor::zeros(&[8, 5], DType::F32, &cuda_device()).unwrap(); let err = ppo_to_exposure_scores(&probs); assert!(err.is_err(), "should reject [8,5]"); } // ── signal_to_action_scores ───────────────────────────────────────── - fn argmax_row(t: &Tensor, row: usize) -> usize { - let data = t.to_vec2::().unwrap(); - let r = data.get(row).unwrap(); - r.iter() - .enumerate() - .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(i, _)| i) - .unwrap() + /// GPU-side argmax for a single-row tensor. + fn gpu_argmax(t: &Tensor) -> u32 { + t.argmax(1).unwrap() + .flatten_all().unwrap() + .to_scalar::().unwrap() } #[test] fn test_signal_to_action_scores_strong_long() { - let pred = Tensor::new(&[20.0_f32], &Device::Cpu).unwrap(); + let pred = Tensor::new(&[20.0_f32], &cuda_device()).unwrap(); let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); assert_eq!(scores.dims(), &[1, 5]); - assert_eq!(argmax_row(&scores, 0), 4); // Long100 + assert_eq!(gpu_argmax(&scores), 4); // Long100 } #[test] fn test_signal_to_action_scores_strong_short() { - let pred = Tensor::new(&[-20.0_f32], &Device::Cpu).unwrap(); + let pred = Tensor::new(&[-20.0_f32], &cuda_device()).unwrap(); let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(argmax_row(&scores, 0), 0); // Short100 + assert_eq!(gpu_argmax(&scores), 0); // Short100 } #[test] fn test_signal_to_action_scores_flat() { - let pred = Tensor::new(&[0.0_f32], &Device::Cpu).unwrap(); + let pred = Tensor::new(&[0.0_f32], &cuda_device()).unwrap(); let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(argmax_row(&scores, 0), 2); // Flat + assert_eq!(gpu_argmax(&scores), 2); // Flat } #[test] fn test_signal_to_action_scores_mild_long() { - let pred = Tensor::new(&[7.0_f32], &Device::Cpu).unwrap(); + let pred = Tensor::new(&[7.0_f32], &cuda_device()).unwrap(); let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(argmax_row(&scores, 0), 3); // Long50 + assert_eq!(gpu_argmax(&scores), 3); // Long50 } #[test] fn test_signal_to_action_scores_mild_short() { - let pred = Tensor::new(&[-7.0_f32], &Device::Cpu).unwrap(); + let pred = Tensor::new(&[-7.0_f32], &cuda_device()).unwrap(); let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(argmax_row(&scores, 0), 1); // Short50 + assert_eq!(gpu_argmax(&scores), 1); // Short50 } // ── tft_quantile_to_signal ────────────────────────────────────────── @@ -413,13 +408,15 @@ mod tests { -1.0, 0.5, 2.0, // batch 0: median = 0.5 -3.0, -1.5, 0.0, // batch 1: median = -1.5 ]; - let q = Tensor::from_vec(data, &[2, 1, 3], &Device::Cpu).unwrap(); + let q = Tensor::from_vec(data, &[2, 1, 3], &cuda_device()).unwrap(); let signal = tft_quantile_to_signal(&q).unwrap(); assert_eq!(signal.dims(), &[2]); - let vals = signal.to_vec1::().unwrap(); - assert!((vals.get(0).copied().unwrap_or(0.0) - 0.5).abs() < 1e-6); - assert!((vals.get(1).copied().unwrap_or(0.0) - (-1.5)).abs() < 1e-6); + // GPU-side comparison: signal should equal [0.5, -1.5] + let expected = Tensor::new(&[0.5_f32, -1.5], signal.device()).unwrap(); + let max_diff = signal.sub(&expected).unwrap().abs().unwrap() + .max(0).unwrap().to_scalar::().unwrap(); + assert!(max_diff < 1e-6, "tft signal mismatch: max_diff={max_diff}"); } // ── backtest_fitness ──────────────────────────────────────────────── diff --git a/crates/ml/src/cuda_pipeline/statistics_kernel.cu b/crates/ml/src/cuda_pipeline/statistics_kernel.cu index db3c1d174..e93f8ba85 100644 --- a/crates/ml/src/cuda_pipeline/statistics_kernel.cu +++ b/crates/ml/src/cuda_pipeline/statistics_kernel.cu @@ -70,8 +70,8 @@ extern "C" __global__ void batch_statistics( s_pos_sq_sum[tid] = local_pos_sq_sum; __syncthreads(); - /* Reduction phase */ - for (int s = blockDim.x / 2; s > 0; s >>= 1) { + /* Reduction phase 1: shared memory reduction down to 32 threads */ + for (int s = blockDim.x / 2; s >= 32; s >>= 1) { if (tid < s) { s_sum[tid] += s_sum[tid + s]; s_sq_sum[tid] += s_sq_sum[tid + s]; @@ -84,6 +84,36 @@ extern "C" __global__ void batch_statistics( __syncthreads(); } + /* Reduction phase 2: warp-level reduction (no __syncthreads needed) */ + if (tid < 32) { + float my_sum = s_sum[tid]; + float my_sq_sum = s_sq_sum[tid]; + float my_min = s_min[tid]; + float my_max = s_max[tid]; + int my_win = s_win[tid]; + float my_pos_sum = s_pos_sum[tid]; + float my_pos_sq_sum = s_pos_sq_sum[tid]; + for (int mask = 16; mask >= 1; mask >>= 1) { + my_sum += __shfl_xor_sync(0xFFFFFFFF, my_sum, mask); + my_sq_sum += __shfl_xor_sync(0xFFFFFFFF, my_sq_sum, mask); + my_min = fminf(my_min, __shfl_xor_sync(0xFFFFFFFF, my_min, mask)); + my_max = fmaxf(my_max, __shfl_xor_sync(0xFFFFFFFF, my_max, mask)); + my_win += __shfl_xor_sync(0xFFFFFFFF, my_win, mask); + my_pos_sum += __shfl_xor_sync(0xFFFFFFFF, my_pos_sum, mask); + my_pos_sq_sum += __shfl_xor_sync(0xFFFFFFFF, my_pos_sq_sum, mask); + } + if (tid == 0) { + s_sum[0] = my_sum; + s_sq_sum[0] = my_sq_sum; + s_min[0] = my_min; + s_max[0] = my_max; + s_win[0] = my_win; + s_pos_sum[0] = my_pos_sum; + s_pos_sq_sum[0] = my_pos_sq_sum; + } + } + __syncthreads(); + /* Write results (thread 0 only — single block, no atomics needed) */ if (tid == 0) { output_stats[0] = s_sum[0]; /* reward_sum */ diff --git a/crates/ml/src/cuda_pipeline/training_guard_kernel.cu b/crates/ml/src/cuda_pipeline/training_guard_kernel.cu index bbe8d1d0e..0e31113af 100644 --- a/crates/ml/src/cuda_pipeline/training_guard_kernel.cu +++ b/crates/ml/src/cuda_pipeline/training_guard_kernel.cu @@ -14,12 +14,19 @@ */ /* ------------------------------------------------------------------ */ -/* Kernel 1: training_guard_check */ +/* Kernel 1: training_guard_check_and_accumulate (FUSED) */ /* */ /* Launch: grid=(1,1,1), block=(1,1,1) — single-thread. */ /* */ -/* Reads GPU-resident loss and grad_norm scalars, performs safety */ -/* checks, and writes 7 floats to output (pinned host buffer): */ +/* Combines guard check + accumulate into a single kernel launch. */ +/* Reads GPU-resident loss and grad_norm scalars ONCE (not twice), */ +/* performs safety checks, writes 7 floats to mapped pinned host */ +/* buffer, AND accumulates to device-resident running sums. */ +/* */ +/* Saves one kernel launch per training step (~2-4µs overhead each), */ +/* which compounds to 12-25ms/epoch at 6375 steps. */ +/* */ +/* Mapped output (7 floats, pinned host buffer): */ /* [0] halt_nan — 1.0 if loss or grad_norm is NaN/Inf */ /* [1] halt_loss_clip — 1.0 if loss > clip_threshold */ /* [2] halt_grad_collapse — 1.0 if grad_norm < collapse_threshold */ @@ -28,11 +35,17 @@ /* [4] raw_loss */ /* [5] raw_grad_norm */ /* [6] reserved — 0.0 */ +/* */ +/* acc_buf layout (3 floats, device memory): */ +/* [0] loss_sum */ +/* [1] grad_norm_sum */ +/* [2] step_count (stored as float) */ /* ------------------------------------------------------------------ */ -extern "C" __global__ void training_guard_check( +extern "C" __global__ void training_guard_check_and_accumulate( const float* __restrict__ loss_scalar, /* GPU-resident scalar (1 float) */ const float* __restrict__ grad_norm_scalar, /* GPU-resident scalar (1 float) */ float* output, /* pinned host buffer (7 floats) */ + float* acc_buf, /* device accumulator (3 floats) */ float clip_threshold, /* e.g. 1e6f */ float collapse_threshold, /* e.g. 1e-7f */ int warmup /* 1 = in warmup, skip collapse */ @@ -40,29 +53,25 @@ extern "C" __global__ void training_guard_check( float loss = loss_scalar[0]; float grad_norm = grad_norm_scalar[0]; - /* NaN/Inf guard */ - int nan_loss = (loss != loss) || (loss != loss + 0.0f) - || (loss > 3.402823466e+38f) - || (loss < -3.402823466e+38f); - int nan_grad = (grad_norm != grad_norm) || (grad_norm != grad_norm + 0.0f) - || (grad_norm > 3.402823466e+38f) - || (grad_norm < -3.402823466e+38f); + /* ── Guard check ── */ + + int loss_nan = isnan(loss) || isinf(loss); + int grad_nan = isnan(grad_norm) || isinf(grad_norm); /* Use standard library functions for robust NaN/Inf detection */ - int halt_nan = (isnan(loss) || isinf(loss) || isnan(grad_norm) || isinf(grad_norm)) ? 1 : 0; - (void)nan_loss; (void)nan_grad; /* suppress unused-variable warning */ + int halt_nan = (loss_nan || grad_nan) ? 1 : 0; /* Loss clip */ - int halt_loss_clip = (!isnan(loss) && !isinf(loss) && loss > clip_threshold) ? 1 : 0; + int halt_loss_clip = (!loss_nan && loss > clip_threshold) ? 1 : 0; /* Gradient collapse — only meaningful outside warmup */ int halt_grad_collapse = 0; - if (!warmup && !isnan(grad_norm) && !isinf(grad_norm)) { + if (!warmup && !grad_nan) { halt_grad_collapse = (grad_norm < collapse_threshold) ? 1 : 0; } /* Clipped loss: NaN/Inf → passthrough raw (host will detect via halt_nan) */ - float clipped_loss = (!isnan(loss) && !isinf(loss)) + float clipped_loss = (!loss_nan) ? fminf(loss, clip_threshold) : loss; @@ -74,31 +83,58 @@ extern "C" __global__ void training_guard_check( output[5] = grad_norm; output[6] = 0.0f; /* reserved */ __threadfence_system(); /* ensure mapped writes visible to CPU */ + + /* ── Accumulate ── */ + + /* Skip NaN/Inf to keep running sums meaningful */ + if (!loss_nan) { + acc_buf[0] += loss; + } + if (!grad_nan) { + acc_buf[1] += grad_norm; + } + acc_buf[2] += 1.0f; +} + +/* Legacy entry points — kept for backward compatibility, each delegates + to the fused kernel's logic. New code should use the fused version. */ + +extern "C" __global__ void training_guard_check( + const float* __restrict__ loss_scalar, + const float* __restrict__ grad_norm_scalar, + float* output, + float clip_threshold, + float collapse_threshold, + int warmup +) { + float loss = loss_scalar[0]; + float grad_norm = grad_norm_scalar[0]; + int loss_nan = isnan(loss) || isinf(loss); + int grad_nan = isnan(grad_norm) || isinf(grad_norm); + int halt_nan = (loss_nan || grad_nan) ? 1 : 0; + int halt_loss_clip = (!loss_nan && loss > clip_threshold) ? 1 : 0; + int halt_grad_collapse = 0; + if (!warmup && !grad_nan) { + halt_grad_collapse = (grad_norm < collapse_threshold) ? 1 : 0; + } + float clipped_loss = (!loss_nan) ? fminf(loss, clip_threshold) : loss; + output[0] = (float)halt_nan; + output[1] = (float)halt_loss_clip; + output[2] = (float)halt_grad_collapse; + output[3] = clipped_loss; + output[4] = loss; + output[5] = grad_norm; + output[6] = 0.0f; + __threadfence_system(); } -/* ------------------------------------------------------------------ */ -/* Kernel 2: training_guard_accumulate */ -/* */ -/* Launch: grid=(1,1,1), block=(1,1,1) — single-thread. */ -/* */ -/* Adds loss and grad_norm to device-resident running sums and */ -/* increments the step counter. Called every training step for */ -/* epoch-boundary averaging without a CPU roundtrip. */ -/* */ -/* acc_buf layout (3 floats, device memory): */ -/* [0] loss_sum */ -/* [1] grad_norm_sum */ -/* [2] step_count (stored as float) */ -/* ------------------------------------------------------------------ */ extern "C" __global__ void training_guard_accumulate( - const float* __restrict__ loss_scalar, /* GPU-resident scalar (1 float) */ - const float* __restrict__ grad_norm_scalar, /* GPU-resident scalar (1 float) */ - float* acc_buf /* device accumulator (3 floats) */ + const float* __restrict__ loss_scalar, + const float* __restrict__ grad_norm_scalar, + float* acc_buf ) { float loss = loss_scalar[0]; float grad_norm = grad_norm_scalar[0]; - - /* Skip NaN/Inf to keep running sums meaningful */ if (!isnan(loss) && !isinf(loss)) { acc_buf[0] += loss; } @@ -168,8 +204,8 @@ extern "C" __global__ void qvalue_stats_reduce( s_all_sum[tid] = local_all; __syncthreads(); - /* Parallel reduction with stride halving */ - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + /* Parallel reduction phase 1: shared memory down to 32 threads */ + for (int stride = blockDim.x / 2; stride >= 32; stride >>= 1) { if (tid < stride) { s_min[tid] = fminf(s_min[tid], s_min[tid + stride]); s_max[tid] = fmaxf(s_max[tid], s_max[tid + stride]); @@ -179,6 +215,27 @@ extern "C" __global__ void qvalue_stats_reduce( __syncthreads(); } + /* Parallel reduction phase 2: warp-level reduction (no __syncthreads needed) */ + if (tid < 32) { + float my_min = s_min[tid]; + float my_max = s_max[tid]; + float my_sum = s_sum[tid]; + float my_all_sum = s_all_sum[tid]; + for (int mask = 16; mask >= 1; mask >>= 1) { + my_min = fminf(my_min, __shfl_xor_sync(0xFFFFFFFF, my_min, mask)); + my_max = fmaxf(my_max, __shfl_xor_sync(0xFFFFFFFF, my_max, mask)); + my_sum += __shfl_xor_sync(0xFFFFFFFF, my_sum, mask); + my_all_sum += __shfl_xor_sync(0xFFFFFFFF, my_all_sum, mask); + } + if (tid == 0) { + s_min[0] = my_min; + s_max[0] = my_max; + s_sum[0] = my_sum; + s_all_sum[0] = my_all_sum; + } + } + __syncthreads(); + /* Thread 0 writes final results */ if (tid == 0) { float n = (float)batch_size; diff --git a/scripts/gpu-hotpath-guard.sh b/scripts/gpu-hotpath-guard.sh index 6594cf3ba..818b20d48 100755 --- a/scripts/gpu-hotpath-guard.sh +++ b/scripts/gpu-hotpath-guard.sh @@ -93,13 +93,9 @@ LEAK_PATTERNS=( # Audit those separately with: scripts/cuda-perf-audit.sh ) -# ── Excluded paths (test-only infra that deliberately uses CPU) ──────── +# ── Excluded paths ──────────────────────────────────────────────────── +# NO test exclusions — all DQN/CUDA code must be GPU-only, including tests. EXCLUDE_PATHS=( - "smoke_tests/" - "_test.rs" - "_tests.rs" - "performance_tests.rs" - "performance_validation.rs" "demo_dqn.rs" ) @@ -127,10 +123,6 @@ check_file() { return 0 fi - # Find the line number where #[cfg(test)] starts — everything after is test code - local test_start - test_start=$(grep -n '#\[cfg(test)\]' "$file" 2>/dev/null | head -1 | cut -d: -f1 || echo "999999") - for pattern in "${LEAK_PATTERNS[@]}"; do # Find matches, exclude: doc comments, Storage match arms, regular comments local hits @@ -142,27 +134,14 @@ check_file() { | grep -v 'Storage::Metal' \ || true) - # Filter out lines in test code (line number >= test_start) if [ -n "$hits" ]; then - local filtered="" - while IFS= read -r line; do - local lineno - lineno=$(echo "$line" | cut -d: -f1) - if [ "$lineno" -lt "$test_start" ] 2>/dev/null; then - filtered="${filtered}${line}"$'\n' - fi - done <<< "$hits" - filtered="${filtered%$'\n'}" # trim trailing newline - - if [ -n "$filtered" ]; then - if [ "$found" -eq 0 ]; then - echo "GPU HOT-PATH LEAK: $file" - found=1 - fi - echo "$filtered" | while IFS= read -r line; do - echo " $line" - done + if [ "$found" -eq 0 ]; then + echo "GPU HOT-PATH LEAK: $file" + found=1 fi + echo "$hits" | while IFS= read -r line; do + echo " $line" + done fi done diff --git a/scripts/ptx-cache-invalidate.sh b/scripts/ptx-cache-invalidate.sh new file mode 100755 index 000000000..2192cd526 --- /dev/null +++ b/scripts/ptx-cache-invalidate.sh @@ -0,0 +1,52 @@ +#!/usr/bin/env bash +# ptx-cache-invalidate.sh — Invalidate PTX cache when CUDA kernels change. +# +# The Rust runtime caches compiled PTX in $CARGO_TARGET_DIR/.ptx_cache/ (or +# /tmp/.ptx_cache/). The cache key includes a SHA-256 of the kernel source, +# so changed .cu files naturally miss. However, if the CI PVC preserves the +# cache across runs and the build binary is stale, old PTX can survive. +# +# This script computes a content hash of all .cu/.cuh kernel source files and +# compares it against a stamp file in the cache directory. If the hash differs +# (any kernel source changed), the entire PTX cache is purged, forcing nvcc +# recompilation on the next run. +# +# Usage: scripts/ptx-cache-invalidate.sh [cache_dir] +# cache_dir defaults to $CARGO_TARGET_DIR/.ptx_cache, then /tmp/.ptx_cache + +set -euo pipefail + +KERNEL_DIR="crates/ml/src/cuda_pipeline" +CACHE_DIR="${1:-${CARGO_TARGET_DIR:+${CARGO_TARGET_DIR}/.ptx_cache}}" +CACHE_DIR="${CACHE_DIR:-/tmp/.ptx_cache}" + +# Compute content hash of all CUDA source files +HASH=$(find "$KERNEL_DIR" -type f \( -name '*.cu' -o -name '*.cuh' \) \ + -exec sha256sum {} + 2>/dev/null | sort | sha256sum | cut -d' ' -f1) + +STAMP_FILE="${CACHE_DIR}/.kernel_hash" + +if [ -f "$STAMP_FILE" ]; then + OLD_HASH=$(cat "$STAMP_FILE") + if [ "$HASH" = "$OLD_HASH" ]; then + echo "PTX cache valid — kernel hash unchanged ($HASH)" + exit 0 + fi + echo "PTX cache STALE — kernel hash changed" + echo " old: $OLD_HASH" + echo " new: $HASH" +else + echo "PTX cache stamp missing — first run or cache cleared" +fi + +# Purge stale PTX cache +if [ -d "$CACHE_DIR" ]; then + PTX_COUNT=$(find "$CACHE_DIR" -name '*.ptx' -type f 2>/dev/null | wc -l) + echo "Purging ${PTX_COUNT} cached PTX files from ${CACHE_DIR}" + rm -f "${CACHE_DIR}"/*.ptx +fi + +# Write new stamp +mkdir -p "$CACHE_DIR" +echo "$HASH" > "$STAMP_FILE" +echo "PTX cache stamp updated: $HASH"