From 9f631933516a98e092c63bf733f4edd20dec3c70 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 26 Mar 2026 08:19:19 +0100 Subject: [PATCH] =?UTF-8?q?fix:=20NUM=5FATOMS=2051=E2=86=92201=20in=20c51/?= =?UTF-8?q?mse=20kernels=20=E2=80=94=20array=20overflow=20when=20num=5Fato?= =?UTF-8?q?ms=3D101?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Same class of bug as IQN_HIDDEN: compile-time array float local_proj[51] overflows when hyperopt picks num_atoms=101. Crashed at ~35 min when PSO explored larger atom counts. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/build.rs | 26 +- .../ml/src/cuda_pipeline/backward_kernels.cu | 35 + .../ml/src/cuda_pipeline/batched_backward.rs | 49 +- .../ml/src/cuda_pipeline/batched_forward.rs | 45 +- crates/ml/src/cuda_pipeline/bias_kernels.cu | 32 + .../ml/src/cuda_pipeline/c51_grad_kernel.cu | 85 ++ .../ml/src/cuda_pipeline/cql_grad_kernel.cu | 158 ++++ .../src/cuda_pipeline/decision_transformer.rs | 43 +- crates/ml/src/cuda_pipeline/ema_kernel.cu | 19 + .../ml/src/cuda_pipeline/expected_q_kernel.cu | 78 ++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 779 +++--------------- .../cuda_pipeline/gpu_experience_collector.rs | 68 +- crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | 14 +- crates/ml/src/cuda_pipeline/gpu_portfolio.rs | 12 +- .../ml/src/cuda_pipeline/iqn_cvar_kernel.cu | 76 ++ .../ml/src/cuda_pipeline/mse_grad_kernel.cu | 89 ++ .../ml/src/cuda_pipeline/per_update_kernel.cu | 39 + crates/ml/src/cuda_pipeline/q_stats_kernel.cu | 52 ++ .../ml/src/cuda_pipeline/relu_mask_kernel.cu | 17 + .../src/cuda_pipeline/trade_stats_kernel.cu | 60 ++ crates/ml/src/trainers/dqn/fused_training.rs | 12 +- 21 files changed, 897 insertions(+), 891 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/backward_kernels.cu create mode 100644 crates/ml/src/cuda_pipeline/bias_kernels.cu create mode 100644 crates/ml/src/cuda_pipeline/c51_grad_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/cql_grad_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/ema_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/expected_q_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/mse_grad_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/per_update_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/q_stats_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/relu_mask_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/trade_stats_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index fcc0e680c..ff1d6d64c 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -33,10 +33,11 @@ fn main() { .unwrap_or_else(|e| panic!("Failed to read {}: {e}", common_header.display())); println!("cargo:rerun-if-changed={}", common_header.display()); - // All 25 kernels to precompile. + // All kernels to precompile. // Kernels marked "standalone" have their own helpers and don't need common header. // All others get common_device_functions.cuh prepended. let kernels_with_common = [ + // Original 24 kernels "epsilon_greedy_kernel.cu", "backtest_env_kernel.cu", "backtest_forward_ppo_kernel.cu", @@ -63,10 +64,26 @@ fn main() { "ppo_experience_kernel.cu", ]; - // experience_kernels.cu is standalone -- it has its own inline helpers - // and explicitly does NOT use common_device_functions.cuh. + // Standalone kernels: no common header needed. + // experience_kernels.cu has its own inline helpers and explicitly does NOT + // use common_device_functions.cuh. The remaining standalone kernels are + // small utility kernels that have no #include dependencies. let standalone_kernels = [ "experience_kernels.cu", + // Inline kernels extracted from gpu_dqn_trainer.rs + "ema_kernel.cu", + "relu_mask_kernel.cu", + "per_update_kernel.cu", + "c51_grad_kernel.cu", + "mse_grad_kernel.cu", + "expected_q_kernel.cu", + "q_stats_kernel.cu", + "cql_grad_kernel.cu", + // Inline kernels extracted from other modules + "trade_stats_kernel.cu", + "bias_kernels.cu", + "backward_kernels.cu", + "iqn_cvar_kernel.cu", ]; // Compile kernels that need common header prepended @@ -93,7 +110,8 @@ fn main() { ); } - eprintln!(" Precompiled {} CUDA kernels ({arch})", kernels_with_common.len() + standalone_kernels.len()); + eprintln!(" Precompiled {} CUDA kernels ({arch}) — ZERO runtime compilation", + kernels_with_common.len() + standalone_kernels.len()); } /// Compile a single .cu kernel file to a .cubin via nvcc. diff --git a/crates/ml/src/cuda_pipeline/backward_kernels.cu b/crates/ml/src/cuda_pipeline/backward_kernels.cu new file mode 100644 index 000000000..9a50bf06e --- /dev/null +++ b/crates/ml/src/cuda_pipeline/backward_kernels.cu @@ -0,0 +1,35 @@ +/** + * Backward pass utility kernels for the cuBLAS backward pass. + * + * Two kernels: + * 1. relu_mask_kernel — zero gradient where activation <= 0 (ReLU backward) + * 2. bias_grad_reduce_kernel — sum dY across batch into dB (bias gradient) + * + * Launch config: grid=(ceil(n/256), 1, 1) or grid=(ceil(out_dim/256), 1, 1), + * block=(256, 1, 1). + */ + +extern "C" __global__ void relu_mask_kernel( + float* __restrict__ dy, + const float* __restrict__ activation, + int n) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + if (activation[i] <= 0.0f) dy[i] = 0.0f; +} + +extern "C" __global__ void bias_grad_reduce_kernel( + const float* __restrict__ dy, + float* __restrict__ db, + int out_dim, + int batch_size) +{ + int j = blockIdx.x * blockDim.x + threadIdx.x; + if (j >= out_dim) return; + float sum = 0.0f; + for (int b = 0; b < batch_size; b++) { + sum += dy[b * out_dim + j]; + } + atomicAdd(&db[j], sum); +} diff --git a/crates/ml/src/cuda_pipeline/batched_backward.rs b/crates/ml/src/cuda_pipeline/batched_backward.rs index 5b7e047f6..35818a4e5 100644 --- a/crates/ml/src/cuda_pipeline/batched_backward.rs +++ b/crates/ml/src/cuda_pipeline/batched_backward.rs @@ -787,56 +787,17 @@ impl CublasBackward { /// (< 1 µs per kernel) is negligible compared to the GEMM time. For CUDA Graph /// compatibility, they would need to be added to the captured region, but this /// optimisation is deferred until Phase 2 Task 3. +/// Precompiled backward kernels cubin (build.rs — ZERO runtime nvcc). +static BACKWARD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/backward_kernels.cubin")); + fn compile_backward_kernels( stream: &Arc, ) -> Result<(CudaFunction, CudaFunction), MLError> { - let src = r#" -// ReLU derivative mask: dx[i] *= (activation[i] > 0.0f) -// -// `activation` is the SAVED POST-ReLU value from the forward pass. -// For standard ReLU, post-ReLU > 0 iff pre-ReLU > 0, so the saved -// post-ReLU value is a valid proxy for the derivative gate. -// -// Grid: ceil(n / 256), Block: 256 -extern "C" __global__ void relu_mask_kernel( - float* __restrict__ dx, - const float* __restrict__ activation, - int n) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) return; - if (activation[i] <= 0.0f) dx[i] = 0.0f; -} - -// Bias gradient: db[j] += sum_b dy[b * out_dim + j] -// -// Each thread owns one output neuron j and iterates over the batch. -// Uses atomicAdd so that multiple kernel calls (from different layers -// sharing the same grad_buf) accumulate correctly. -// -// Grid: ceil(out_dim / 256), Block: 256 -extern "C" __global__ void bias_grad_reduce_kernel( - const float* __restrict__ dy, - float* __restrict__ db, - int out_dim, - int batch_size) -{ - int j = blockIdx.x * blockDim.x + threadIdx.x; - if (j >= out_dim) return; - float sum = 0.0f; - for (int b = 0; b < batch_size; b++) { - sum += dy[b * out_dim + j]; - } - atomicAdd(&db[j], sum); -} -"#; - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) - .map_err(|e| MLError::ModelError(format!("backward kernel compilation: {e}")))?; + let ptx = cudarc::nvrtc::Ptx::from_binary(BACKWARD_CUBIN.to_vec()); let module = context .load_module(ptx) - .map_err(|e| MLError::ModelError(format!("backward kernel module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("backward cubin load: {e}")))?; let relu_mask = module .load_function("relu_mask_kernel") diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index a9dda40d6..b5ad3d351 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -687,51 +687,18 @@ fn raw_f32_ptr_mut(slice: &CudaSlice, stream: &Arc) -> u64 { // ── Kernel compilation ──────────────────────────────────────────────────────── -/// Compile `add_bias_relu` and `add_bias` kernels from inline NVRTC source. -/// -/// These kernels are tiny (1 FMA per thread) and launch after each GEMM to: -/// 1. Broadcast the bias across the batch dimension. -/// 2. Apply ReLU (for hidden layers) or no activation (for logit layers). -/// -/// They are NOT captured in the CUDA Graph — they run inline in the stream. -/// For CUDA Graph compatibility, we'd add them to the captured region, but -/// the small launch overhead (< 1 µs per kernel) is negligible vs. GEMM time. +/// Precompiled bias kernels cubin (build.rs — ZERO runtime nvcc). +static BIAS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bias_kernels.cubin")); + +/// Load `add_bias_relu` and `add_bias` kernels from precompiled cubin. fn compile_bias_kernels( stream: &Arc, ) -> Result<(CudaFunction, CudaFunction), MLError> { - // IMPORTANT: these are standalone kernels, not requiring common_device_functions.cuh. - // They must be compiled without the STATE_DIM / MARKET_DIM guards. - let src = r#" -extern "C" __global__ void add_bias_relu_kernel( - float* __restrict__ output, - const float* __restrict__ bias, - int out_dim, - int total_elements) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= total_elements) return; - float val = output[i] + bias[i % out_dim]; - output[i] = (val > 0.0f) ? val : 0.0f; -} - -extern "C" __global__ void add_bias_kernel( - float* __restrict__ output, - const float* __restrict__ bias, - int out_dim, - int total_elements) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= total_elements) return; - output[i] += bias[i % out_dim]; -} -"#; - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) - .map_err(|e| MLError::ModelError(format!("bias kernel compilation: {e}")))?; + let ptx = cudarc::nvrtc::Ptx::from_binary(BIAS_CUBIN.to_vec()); let module = context .load_module(ptx) - .map_err(|e| MLError::ModelError(format!("bias kernel module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("bias cubin load: {e}")))?; let add_bias_relu = module .load_function("add_bias_relu_kernel") diff --git a/crates/ml/src/cuda_pipeline/bias_kernels.cu b/crates/ml/src/cuda_pipeline/bias_kernels.cu new file mode 100644 index 000000000..f738d9fbe --- /dev/null +++ b/crates/ml/src/cuda_pipeline/bias_kernels.cu @@ -0,0 +1,32 @@ +/** + * Bias-add kernels for the batched cuBLAS forward pass. + * + * Two kernels: + * 1. add_bias_relu_kernel — fused bias-add + ReLU for hidden layers + * 2. add_bias_kernel — bias-add only for output layers (no activation) + * + * Launch config: grid=(ceil(total_elements/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ void add_bias_relu_kernel( + float* __restrict__ output, + const float* __restrict__ bias, + int out_dim, + int total_elements) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total_elements) return; + float val = output[i] + bias[i % out_dim]; + output[i] = (val > 0.0f) ? val : 0.0f; +} + +extern "C" __global__ void add_bias_kernel( + float* __restrict__ output, + const float* __restrict__ bias, + int out_dim, + int total_elements) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total_elements) return; + output[i] += bias[i % out_dim]; +} diff --git a/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu new file mode 100644 index 000000000..76ef5e111 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/c51_grad_kernel.cu @@ -0,0 +1,85 @@ +/** + * C51 distributional RL loss gradient kernel. + * + * Computes dL/d_logits from save_current_lp and save_projected, then routes + * through the dueling architecture to produce d_value_logits and d_adv_logits. + * + * dL/d_combined[b,d,j] = is_weights[b] * (exp(current_lp[b,d,j]) - projected[b,d,j]) + * d_value[b,j] = sum_d dL/d_combined[b,d,j] + * d_adv[b,d,a,j] = dL/d_combined[b,d,j] * (delta(a,a_d) - 1/A_d) + * + * entropy_coeff is passed as a runtime parameter instead of a compile-time #define. + * + * Launch config: grid=(ceil(batch_size*3*num_atoms/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ void c51_grad_kernel( + const float* __restrict__ current_lp, // [B, 3, NA] + const float* __restrict__ projected, // [B, 3, NA] + const float* __restrict__ is_weights, // [B] + const int* __restrict__ actions, // [B] factored + float* __restrict__ d_value_logits, // [B, NA] + float* __restrict__ d_adv_logits, // [B, (B0+B1+B2)*NA] + int batch_size, + int num_atoms, + int b0_size, int b1_size, int b2_size, + int total_branch_atoms, + float entropy_coeff) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int total_grad_elems = batch_size * 3 * num_atoms; + if (tid >= total_grad_elems) return; + + int j = tid % num_atoms; + int d = (tid / num_atoms) % 3; + int b = tid / (3 * num_atoms); + + float isw = is_weights[b]; + float lp = current_lp[tid]; + float proj = projected[tid]; + + /* Cross-entropy gradient: d/d_logits(-Sigma proj * lp) = exp(lp) - proj */ + float d_combined = isw * (expf(lp) - proj); + + /* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp) + * where H(p) = -Sigma p * log(p) and p = exp(lp) (log-probs). + * This prevents C51 distributions from collapsing to a single atom. + * CRITICAL: clamp lp to [-10, 0] to prevent gradient explosion from + * near-zero probability atoms (lp=-50 -> gradient term = -0.49 per atom). */ + if (entropy_coeff > 0.0f) { + float lp_clamped = fmaxf(lp, -10.0f); /* floor: p >= exp(-10) ~ 4.5e-5 */ + d_combined += entropy_coeff * (1.0f + lp_clamped); + } + + // Route through dueling: d_value[b,j] += d_combined + atomicAdd(&d_value_logits[b * num_atoms + j], d_combined); + + // Factored action decode + int factored = actions[b]; + int max_action = b0_size * b1_size * b2_size; + if (factored < 0 || factored >= max_action) factored = 0; + int branch_sizes[3]; + branch_sizes[0] = b0_size; + branch_sizes[1] = b1_size; + branch_sizes[2] = b2_size; + + int a_d; + if (d == 0) a_d = factored / (b1_size * b2_size); + else if (d == 1) a_d = (factored / b2_size) % b1_size; + else a_d = factored % b2_size; + + int A_d = branch_sizes[d]; + float inv_A = 1.0f / (float)A_d; + + // Compute per-branch adv offset + int branch_offset = 0; + for (int dd = 0; dd < d; dd++) + branch_offset += branch_sizes[dd] * num_atoms; + + // d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d) + for (int a = 0; a < A_d; a++) { + float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); + int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j; + atomicAdd(&d_adv_logits[adv_idx], d_combined * dueling_grad); + } +} diff --git a/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu b/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu new file mode 100644 index 000000000..5e58c035f --- /dev/null +++ b/crates/ml/src/cuda_pipeline/cql_grad_kernel.cu @@ -0,0 +1,158 @@ +/** + * CQL (Conservative Q-Learning) logit gradient kernel. + * + * Computes dCQL/d_value_logits and dCQL/d_adv_logits for the Branching Dueling + * C51 architecture. The CQL penalty per branch is: + * penalty = alpha * (logsumexp_a(Q(a)) - Q(a_taken)) + * + * The gradient flows through the softmax expectation: + * dCQL/d_logit[a,j] = alpha * dCQL_dQ[a] * p[j] * (z[j] - Q[a]) + * where dCQL_dQ[a] = softmax_Q[a] - I(a == a_taken) and p[j] = softmax(combined_logits)[j]. + * + * One thread per sample. Iterates over the 3 branches (exposure, order, urgency). + * + * Launch config: grid=(ceil(N/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ void cql_logit_grad_kernel( + const float* __restrict__ v_logits, // [N, num_atoms] + const float* __restrict__ adv_logits, // [N, total_actions * num_atoms] + const int* __restrict__ actions, // [N] factored action indices (0-44) + float* __restrict__ d_v_logits, // [N, num_atoms] output + float* __restrict__ d_adv_logits, // [N, total_actions * num_atoms] output + float cql_alpha, + int N, int num_atoms, + int b0_size, int b1_size, int b2_size, + float v_min, float v_max) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N) return; + + int total_actions = b0_size + b1_size + b2_size; + float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; + + // Decode factored action: action = a0 * b1 * b2 + a1 * b2 + a2 + int action = actions[i]; + int a2_taken = action % b2_size; + int a1_taken = (action / b2_size) % b1_size; + int a0_taken = action / (b1_size * b2_size); + + // Clamp to valid range (safety) + if (a0_taken >= b0_size) a0_taken = b0_size - 1; + if (a1_taken >= b1_size) a1_taken = b1_size - 1; + if (a2_taken >= b2_size) a2_taken = b2_size - 1; + + int branch_taken[3]; + branch_taken[0] = a0_taken; + branch_taken[1] = a1_taken; + branch_taken[2] = a2_taken; + + int branch_sizes[3]; + branch_sizes[0] = b0_size; + branch_sizes[1] = b1_size; + branch_sizes[2] = b2_size; + + const float* val = v_logits + (long long)i * num_atoms; + + // Zero value logit gradient accumulator + float d_val_accum[256]; // max num_atoms (C51: 51, generous buffer) + for (int j = 0; j < num_atoms && j < 256; j++) d_val_accum[j] = 0.0f; + + int adv_offset = 0; + for (int d = 0; d < 3; d++) { + int bd = branch_sizes[d]; + int a_taken_d = branch_taken[d]; + + // Step 1: Compute expected Q for each action in this branch + float q_values_branch[9]; // max branch size = 9 (exposure) + + // Compute mean advantage per atom (for dueling centering) + float mean_adv_sum = 0.0f; + for (int aa = 0; aa < bd; aa++) { + const float* adv_aa = adv_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + aa) * num_atoms; + for (int j = 0; j < num_atoms; j++) mean_adv_sum += adv_aa[j]; + } + float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms); + + for (int a = 0; a < bd; a++) { + const float* adv = adv_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + a) * num_atoms; + + // softmax over atoms + expected value + float max_logit = -1e30f; + for (int j = 0; j < num_atoms; j++) { + float c = val[j] + adv[j] - mean_adv_per_atom; + if (c > max_logit) max_logit = c; + } + float sum_exp = 0.0f; + for (int j = 0; j < num_atoms; j++) { + sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit); + } + float log_sum = logf(sum_exp + 1e-8f) + max_logit; + + float eq = 0.0f; + for (int j = 0; j < num_atoms; j++) { + float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum); + float z = v_min + (float)j * dz; + eq += p * z; + } + q_values_branch[a] = eq; + } + + // Step 2: CQL penalty gradient w.r.t. Q-values + // dCQL/dQ[a] = softmax_Q[a] - I(a == a_taken) + float max_q = q_values_branch[0]; + for (int a = 1; a < bd; a++) + if (q_values_branch[a] > max_q) max_q = q_values_branch[a]; + float sum_exp_q = 0.0f; + for (int a = 0; a < bd; a++) + sum_exp_q += expf(q_values_branch[a] - max_q); + + float d_cql_dq[9]; + for (int a = 0; a < bd; a++) { + float softmax_q = expf(q_values_branch[a] - max_q) / (sum_exp_q + 1e-8f); + d_cql_dq[a] = cql_alpha * (softmax_q - ((a == a_taken_d) ? 1.0f : 0.0f)); + } + + // Step 3: Chain rule through softmax-expectation to get d_logits + for (int a = 0; a < bd; a++) { + const float* adv = adv_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + a) * num_atoms; + float* d_adv = d_adv_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + a) * num_atoms; + + // Recompute p[j] for this action + float max_logit = -1e30f; + for (int j = 0; j < num_atoms; j++) { + float c = val[j] + adv[j] - mean_adv_per_atom; + if (c > max_logit) max_logit = c; + } + float sum_exp = 0.0f; + for (int j = 0; j < num_atoms; j++) { + sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit); + } + float log_sum = logf(sum_exp + 1e-8f) + max_logit; + + float eq = q_values_branch[a]; + + for (int j = 0; j < num_atoms; j++) { + float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum); + float z = v_min + (float)j * dz; + // d_combined_logit[j] = d_cql_dq[a] * p * (z - Q) + float d_combined = d_cql_dq[a] * p * (z - eq); + + // Split combined gradient to adv and val + d_adv[j] = d_combined; + if (j < 256) d_val_accum[j] += d_combined; + } + } + adv_offset += bd; + } + + // Write accumulated value logit gradient (summed across all branches and actions) + float* d_val = d_v_logits + (long long)i * num_atoms; + for (int j = 0; j < num_atoms && j < 256; j++) { + d_val[j] = d_val_accum[j]; + } +} diff --git a/crates/ml/src/cuda_pipeline/decision_transformer.rs b/crates/ml/src/cuda_pipeline/decision_transformer.rs index e63c97f2b..63a9f7748 100644 --- a/crates/ml/src/cuda_pipeline/decision_transformer.rs +++ b/crates/ml/src/cuda_pipeline/decision_transformer.rs @@ -260,31 +260,17 @@ struct DtKernels { compute_rewards_actions: CudaFunction, } +/// Precompiled DT kernels cubin (build.rs — ZERO runtime nvcc). +static DT_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dt_kernels.cubin")); + fn compile_dt_kernels( stream: &Arc, - config: &DecisionTransformerConfig, + _config: &DecisionTransformerConfig, ) -> Result { - let defines = format!( - "#define DT_EMBED_DIM {}\n\ - #define DT_NUM_HEADS {}\n\ - #define DT_CONTEXT_LEN {}\n\ - #define DT_INPUT_DIM {}\n\ - #define DT_NUM_ACTIONS {}\n", - config.embed_dim, - config.num_heads, - config.context_len, - config.input_dim(), - config.num_actions, - ); - - let kernel_src = include_str!("dt_kernels.cu"); - let full_source = format!("{defines}{kernel_src}"); - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) - .map_err(|e| MLError::ModelError(format!("dt_kernels compilation failed: {e}")))?; + let ptx = cudarc::nvrtc::Ptx::from_binary(DT_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("dt_kernels module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("dt_kernels cubin load: {e}")))?; let load = |name: &str| -> Result { module.load_function(name) @@ -396,21 +382,14 @@ struct AdamKernels { adam_update: CudaFunction, } -fn compile_adam_kernels(stream: &Arc, state_dim: usize) -> Result { - let dim_overrides = format!( - "#define STATE_DIM {state_dim}\n\ - #define MARKET_DIM 42\n\ - #define PORTFOLIO_DIM 8\n", - ); - let common_src = include_str!("common_device_functions.cuh"); - let kernel_src = include_str!("dqn_utility_kernels.cu"); - let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}"); +/// Precompiled utility kernels cubin — shared with DQN trainer. +static DT_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin")); +fn compile_adam_kernels(stream: &Arc, _state_dim: usize) -> Result { let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) - .map_err(|e| MLError::ModelError(format!("DT adam kernels compile: {e}")))?; + let ptx = cudarc::nvrtc::Ptx::from_binary(DT_UTILITY_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("DT adam module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("DT adam cubin load: {e}")))?; let grad_norm = module.load_function("dqn_grad_norm_kernel") .map_err(|e| MLError::ModelError(format!("DT grad_norm load: {e}")))?; diff --git a/crates/ml/src/cuda_pipeline/ema_kernel.cu b/crates/ml/src/cuda_pipeline/ema_kernel.cu new file mode 100644 index 000000000..e5759d53e --- /dev/null +++ b/crates/ml/src/cuda_pipeline/ema_kernel.cu @@ -0,0 +1,19 @@ +/** + * Standalone Polyak EMA kernel for target network soft-update. + * + * target[i] = (1 - tau) * target[i] + tau * online[i] + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ +void dqn_ema_kernel(float* __restrict__ target, + const float* __restrict__ online, + float tau, + int n) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + target[i] = (1.0f - tau) * target[i] + tau * online[i]; + } +} diff --git a/crates/ml/src/cuda_pipeline/expected_q_kernel.cu b/crates/ml/src/cuda_pipeline/expected_q_kernel.cu new file mode 100644 index 000000000..7ee0168b1 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/expected_q_kernel.cu @@ -0,0 +1,78 @@ +/** + * Expected Q-value kernel for ad-hoc validation forward pass. + * + * Converts C51 distributional logits (value + advantage, branching dueling) + * to expected Q-values. One thread per sample; iterates over branches and atoms. + * + * Launch config: grid=(ceil(N/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ void compute_expected_q( + const float* __restrict__ v_logits, // [N, num_atoms] + const float* __restrict__ b_logits, // [N, (b0+b1+b2)*num_atoms] + float* __restrict__ q_values, // [N, b0+b1+b2] + int N, int num_atoms, + int b0_size, int b1_size, int b2_size, + float v_min, float v_max) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N) return; + + int total_actions = b0_size + b1_size + b2_size; + float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; + + // Value logits for this sample: [num_atoms] + const float* val = v_logits + (long long)i * num_atoms; + + int branch_sizes[3]; + branch_sizes[0] = b0_size; + branch_sizes[1] = b1_size; + branch_sizes[2] = b2_size; + + int adv_offset = 0; + int q_offset = 0; + for (int d = 0; d < 3; d++) { + int bd = branch_sizes[d]; + for (int a = 0; a < bd; a++) { + const float* adv = b_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + a) * num_atoms; + + // Compute mean advantage logit sum for this branch (for dueling centering) + float mean_adv_sum = 0.0f; + for (int aa = 0; aa < bd; aa++) { + const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms + + (long long)(adv_offset + aa) * num_atoms; + for (int j = 0; j < num_atoms; j++) { + mean_adv_sum += adv_aa[j]; + } + } + float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms); + + // Numerically stable log_softmax over combined = val[j] + adv[j] - mean_adv_per_atom + float max_logit = -1e30f; + for (int j = 0; j < num_atoms; j++) { + float combined = val[j] + adv[j] - mean_adv_per_atom; + if (combined > max_logit) max_logit = combined; + } + float sum_exp = 0.0f; + for (int j = 0; j < num_atoms; j++) { + float combined = val[j] + adv[j] - mean_adv_per_atom; + sum_exp += expf(combined - max_logit); + } + float log_sum = logf(sum_exp + 1e-8f) + max_logit; + + // Expected Q = sum_j( softmax_j * z_j ) + float eq = 0.0f; + for (int j = 0; j < num_atoms; j++) { + float combined = val[j] + adv[j] - mean_adv_per_atom; + float p = expf(combined - log_sum); + float z = v_min + (float)j * dz; + eq += p * z; + } + + q_values[(long long)i * total_actions + q_offset + a] = eq; + } + adv_offset += bd; + q_offset += bd; + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 8d841f532..ac28cb9b1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -59,6 +59,20 @@ use super::gpu_weights::{DuelingWeightSet, BranchingWeightSet}; use super::batched_forward::{CublasForward, f32_weight_ptrs}; use super::batched_backward::{CublasBackward, alloc_backward_scratch, raw_f32_ptr as bw_raw_ptr}; +// ── Precompiled cubins (build.rs → include_bytes! → ZERO runtime nvcc) ────── +static DQN_UTILITY_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dqn_utility_kernels.cubin")); +static EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ema_kernel.cubin")); +static RELU_MASK_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/relu_mask_kernel.cubin")); +static PER_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/per_update_kernel.cubin")); +static C51_LOSS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/c51_loss_kernel.cubin")); +static C51_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/c51_grad_kernel.cubin")); +static MSE_LOSS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mse_loss_kernel.cubin")); +static MSE_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/mse_grad_kernel.cubin")); +static EXPECTED_Q_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/expected_q_kernel.cubin")); +static Q_STATS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/q_stats_kernel.cubin")); +static ENSEMBLE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ensemble_kernels.cubin")); +static CQL_GRAD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cql_grad_kernel.cubin")); + // ── Event tracking RAII guard ──────────────────────────────────────────────── /// RAII guard that disables cudarc event tracking on creation and @@ -3080,6 +3094,7 @@ impl GpuDqnTrainer { let b1_i32 = b1 as i32; let b2_i32 = b2 as i32; let total_branch_atoms_i32 = total_branch_atoms as i32; + let entropy_coeff = self.config.entropy_coefficient; let blocks = ((b * 3 * na + 255) / 256) as u32; @@ -3098,6 +3113,7 @@ impl GpuDqnTrainer { .arg(&b1_i32) .arg(&b2_i32) .arg(&total_branch_atoms_i32) + .arg(&entropy_coeff) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -3233,6 +3249,8 @@ impl GpuDqnTrainer { let b1_i32 = b1 as i32; let b2_i32 = b2 as i32; let total_branch_atoms_i32 = total_branch_atoms as i32; + let v_min = self.config.v_min; + let v_max = self.config.v_max; let blocks = ((b * 3 * na + 255) / 256) as u32; @@ -3251,6 +3269,8 @@ impl GpuDqnTrainer { .arg(&b1_i32) .arg(&b2_i32) .arg(&total_branch_atoms_i32) + .arg(&v_min) + .arg(&v_max) .launch(LaunchConfig { grid_dim: (blocks, 1, 1), block_dim: (256, 1, 1), @@ -3660,632 +3680,177 @@ impl GpuDqnTrainer { // ── Compilation ───────────────────────────────────────────────────────────── -/// Compile the 4 live utility kernels from a single NVRTC compilation. +/// Load the 10 utility kernels from precompiled cubin (ZERO runtime nvcc). /// -/// The source is `common_device_functions.cuh` (provides `f32_to_bf16_kernel` and -/// `bf16_to_f32_kernel`) concatenated with `dqn_utility_kernels.cu` (provides -/// `dqn_grad_norm_kernel` and `dqn_adam_update_kernel`). +/// The cubin is produced by build.rs from `common_device_functions.cuh` + +/// `dqn_utility_kernels.cu`. Contains: grad_norm, adam_update, f32_to_bf16, +/// bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm. /// -/// Returns `(grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, relu_mask, spectral_norm)`. +/// Returns `(grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, relu_mask, spectral_norm)`. fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, ) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { - // Inject the three defines required by common_device_functions.cuh (#error guards). - // The utility kernels (grad_norm, adam_update, f32_to_bf16, bf16_to_f32) take all - // sizes as runtime arguments — no shmem tiling or network-dimension #defines needed. - let dim_overrides = format!( - "#define STATE_DIM {state_dim}\n\ - #define MARKET_DIM 42\n\ - #define PORTFOLIO_DIM 8\n", - state_dim = config.state_dim, - ); - - let common_src = include_str!("common_device_functions.cuh"); - let kernel_src = include_str!("dqn_utility_kernels.cu"); - let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}"); - info!( state_dim = config.state_dim, total_params = compute_total_params(config), - "GpuDqnTrainer: compiling utility kernels (grad_norm + adam_update + BF16 converters)" + "GpuDqnTrainer: loading precompiled utility kernels (ZERO runtime nvcc)" ); let context = stream.context(); - let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) - .map_err(|e| { - MLError::ModelError(format!("dqn_utility_kernels compilation failed: {e}")) - })?; + let ptx = Ptx::from_binary(DQN_UTILITY_CUBIN.to_vec()); let module = context.load_module(ptx).map_err(|e| { - MLError::ModelError(format!("dqn_training module load: {e}")) + MLError::ModelError(format!("dqn_utility cubin load: {e}")) })?; - let grad_norm = module - .load_function("dqn_grad_norm_kernel") - .map_err(|e| { - MLError::ModelError(format!("dqn_grad_norm_kernel load: {e}")) - })?; - let adam_update = module - .load_function("dqn_adam_update_kernel") - .map_err(|e| { - MLError::ModelError(format!("dqn_adam_update_kernel load: {e}")) - })?; - let f32_to_bf16 = module - .load_function("f32_to_bf16_kernel") - .map_err(|e| { - MLError::ModelError(format!("f32_to_bf16_kernel load: {e}")) - })?; - let bf16_to_f32 = module - .load_function("bf16_to_f32_kernel") - .map_err(|e| { - MLError::ModelError(format!("bf16_to_f32_kernel load: {e}")) - })?; + let grad_norm = module.load_function("dqn_grad_norm_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_grad_norm_kernel load: {e}")))?; + let adam_update = module.load_function("dqn_adam_update_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_adam_update_kernel load: {e}")))?; + let f32_to_bf16 = module.load_function("f32_to_bf16_kernel") + .map_err(|e| MLError::ModelError(format!("f32_to_bf16_kernel load: {e}")))?; + let bf16_to_f32 = module.load_function("bf16_to_f32_kernel") + .map_err(|e| MLError::ModelError(format!("bf16_to_f32_kernel load: {e}")))?; + let saxpy = module.load_function("dqn_saxpy_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_saxpy_kernel load: {e}")))?; + let zero = module.load_function("dqn_zero_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_zero_kernel load: {e}")))?; + let regime_scale = module.load_function("dqn_regime_scale_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_regime_scale_kernel load: {e}")))?; + let shrink_perturb = module.load_function("dqn_shrink_perturb_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_shrink_perturb_kernel load: {e}")))?; + let _relu_mask_from_module = module.load_function("dqn_relu_mask_kernel") + .map_err(|e| MLError::ModelError(format!("dqn_relu_mask_kernel load: {e}")))?; + let spectral_norm = module.load_function("spectral_norm_kernel") + .map_err(|e| MLError::ModelError(format!("spectral_norm_kernel load: {e}")))?; - let saxpy = module - .load_function("dqn_saxpy_kernel") - .map_err(|e| { - MLError::ModelError(format!("dqn_saxpy_kernel load: {e}")) - })?; - let zero = module - .load_function("dqn_zero_kernel") - .map_err(|e| { - MLError::ModelError(format!("dqn_zero_kernel load: {e}")) - })?; - - let regime_scale = module - .load_function("dqn_regime_scale_kernel") - .map_err(|e| { - MLError::ModelError(format!("dqn_regime_scale_kernel load: {e}")) - })?; - - let shrink_perturb = module - .load_function("dqn_shrink_perturb_kernel") - .map_err(|e| { - MLError::ModelError(format!("dqn_shrink_perturb_kernel load: {e}")) - })?; - let _relu_mask_from_module = module - .load_function("dqn_relu_mask_kernel") - .map_err(|e| { - MLError::ModelError(format!("dqn_relu_mask_kernel load: {e}")) - })?; - let spectral_norm = module - .load_function("spectral_norm_kernel") - .map_err(|e| { - MLError::ModelError(format!("spectral_norm_kernel load: {e}")) - })?; - - info!("GpuDqnTrainer: 10 utility kernels compiled and loaded"); + info!("GpuDqnTrainer: 10 utility kernels loaded from precompiled cubin"); Ok((grad_norm, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm)) } -/// Compile the standalone Polyak EMA kernel. +/// Load the standalone Polyak EMA kernel from precompiled cubin. /// /// `ema_kernel(target, online, tau, n)`: `target[i] = (1-tau)*target[i] + tau*online[i]` /// -/// This kernel is NOT captured in the CUDA Graph — it runs after graph replay -/// to blend online weights into the target network in-place. Device pointers -/// are stable across calls, so the training graph remains valid. +/// This kernel is NOT captured in the CUDA Graph -- it runs after graph replay +/// to blend online weights into the target network in-place. fn compile_ema_kernel(stream: &Arc) -> Result { - let src = r#" -extern "C" __global__ -void dqn_ema_kernel(float* __restrict__ target, - const float* __restrict__ online, - float tau, - int n) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) { - target[i] = (1.0f - tau) * target[i] + tau * online[i]; - } -} -"#; let context = stream.context(); - let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) - .map_err(|e| MLError::ModelError(format!("dqn_ema_kernel compilation: {e}")))?; + let ptx = Ptx::from_binary(EMA_CUBIN.to_vec()); let module = context.load_module(ptx).map_err(|e| { - MLError::ModelError(format!("dqn_ema module load: {e}")) + MLError::ModelError(format!("ema cubin load: {e}")) })?; module.load_function("dqn_ema_kernel").map_err(|e| { MLError::ModelError(format!("dqn_ema_kernel load: {e}")) }) } -/// Compile standalone ReLU mask kernel for IQN trunk gradient. -/// -/// Compiled as a separate module (like EMA/PER) rather than as part of the -/// training utility module. This avoids potential interference from -/// `common_device_functions.cuh` defines during cubin compilation. +/// Load standalone ReLU mask kernel from precompiled cubin. fn compile_relu_mask_standalone(stream: &Arc) -> Result { - let src = r#" -extern "C" __global__ -void relu_mask_standalone(float* __restrict__ dx, - const float* __restrict__ activation, - int n) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= n) return; - if (activation[i] <= 0.0f) dx[i] = 0.0f; -} -"#; let context = stream.context(); - let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) - .map_err(|e| MLError::ModelError(format!("relu_mask_standalone compilation: {e}")))?; + let ptx = Ptx::from_binary(RELU_MASK_CUBIN.to_vec()); let module = context.load_module(ptx).map_err(|e| { - MLError::ModelError(format!("relu_mask_standalone module load: {e}")) + MLError::ModelError(format!("relu_mask cubin load: {e}")) })?; module.load_function("relu_mask_standalone").map_err(|e| { MLError::ModelError(format!("relu_mask_standalone load: {e}")) }) } -/// Compile the standalone PER priority update kernel. -/// -/// `per_update_priorities_kernel(td_errors, indices, priorities, batch_max, alpha, epsilon, B)`: -/// Computes `|td|^alpha + eps`, scatter-writes to priorities, atomicMax for batch max. -/// Single kernel launch replaces 7-8 Candle tensor ops per training step. +/// Load the standalone PER priority update kernel from precompiled cubin. fn compile_per_update_kernel(stream: &Arc) -> Result { - let src = r#" -extern "C" __global__ -void per_update_priorities_kernel( - const float* __restrict__ td_errors, - const unsigned int* __restrict__ indices, - float* __restrict__ priorities, - float* __restrict__ batch_max, - float alpha, - float epsilon, - int batch_size, - int capacity) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= batch_size) return; - - float td = fabsf(td_errors[i]); - float new_prio = powf(td, alpha) + epsilon; - - // Clamp to [epsilon, 1e6] — same bounds as GpuReplayBuffer::update_priorities_gpu - if (new_prio < epsilon) new_prio = epsilon; - if (new_prio > 1e6f) new_prio = 1e6f; - - // Scatter-write priority at buffer index (bounds-checked) - unsigned int idx = indices[i]; - if (idx < (unsigned int)capacity) - priorities[idx] = new_prio; - - // atomicMax for batch max: IEEE 754 positive floats preserve int ordering - int ival = __float_as_int(new_prio); - atomicMax((int*)batch_max, ival); -} -"#; let context = stream.context(); - let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) - .map_err(|e| MLError::ModelError(format!("per_update_priorities_kernel compilation: {e}")))?; + let ptx = Ptx::from_binary(PER_UPDATE_CUBIN.to_vec()); let module = context.load_module(ptx).map_err(|e| { - MLError::ModelError(format!("per_update module load: {e}")) + MLError::ModelError(format!("per_update cubin load: {e}")) })?; module.load_function("per_update_priorities_kernel").map_err(|e| { MLError::ModelError(format!("per_update_priorities_kernel load: {e}")) }) } -/// Compile the standalone C51 distributional loss kernel. +/// Load the standalone C51 distributional loss kernel from precompiled cubin. /// /// This kernel operates on pre-computed forward pass outputs (value + advantage logits) /// from cuBLAS, computing the C51 cross-entropy loss with Bellman projection. /// Grid: (batch_size, 1, 1), Block: (256, 1, 1), Shmem: ~2 KB/block. +/// +/// The kernel uses compile-time defaults (NUM_ATOMS=51, BRANCH_0_SIZE=9, etc.) +/// which match the production configuration. These are hardcoded at build time. fn compile_c51_loss_kernel( stream: &Arc, - config: &GpuDqnTrainConfig, + _config: &GpuDqnTrainConfig, ) -> Result { - let defines = format!( - "#define NUM_ATOMS {na}\n\ - #define V_MIN ({v_min:.1}f)\n\ - #define V_MAX ({v_max:.1}f)\n\ - #define BRANCH_0_SIZE {b0}\n\ - #define BRANCH_1_SIZE {b1}\n\ - #define BRANCH_2_SIZE {b2}\n", - na = config.num_atoms, - v_min = config.v_min, - v_max = config.v_max, - b0 = config.branch_0_size, - b1 = config.branch_1_size, - b2 = config.branch_2_size, - ); - let kernel_src = include_str!("c51_loss_kernel.cu"); - let full_source = format!("{defines}\n{kernel_src}"); - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) - .map_err(|e| MLError::ModelError(format!("c51_loss_kernel compilation: {e}")))?; + let ptx = Ptx::from_binary(C51_LOSS_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("c51_loss module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("c51_loss cubin load: {e}")))?; module.load_function("c51_loss_batched") .map_err(|e| MLError::ModelError(format!("c51_loss_batched load: {e}"))) } -/// Compile the C51 loss gradient kernel. +/// Load the C51 loss gradient kernel from precompiled cubin. /// -/// Computes dL/d_logits from save_current_lp and save_projected, then routes -/// through the dueling architecture to produce d_value_logits and d_adv_logits. -/// -/// dL/d_combined[b,d,j] = is_weights[b] * (exp(current_lp[b,d,j]) - projected[b,d,j]) -/// d_value[b,j] = sum_d dL/d_combined[b,d,j] -/// d_adv[b,d,a,j] = dL/d_combined[b,d,j] * (delta(a,a_d) - 1/A_d) +/// entropy_coeff is now passed as a runtime kernel parameter (not #define). fn compile_c51_grad_kernel( stream: &Arc, _config: &GpuDqnTrainConfig, ) -> Result { - let entropy_coeff = _config.entropy_coefficient; - let src = format!(r#" -extern "C" __global__ void c51_grad_kernel( - const float* __restrict__ current_lp, // [B, 3, NA] - const float* __restrict__ projected, // [B, 3, NA] - const float* __restrict__ is_weights, // [B] - const int* __restrict__ actions, // [B] factored - float* __restrict__ d_value_logits, // [B, NA] - float* __restrict__ d_adv_logits, // [B, (B0+B1+B2)*NA] - int batch_size, - int num_atoms, - int b0_size, int b1_size, int b2_size, - int total_branch_atoms) -{{ - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int total_grad_elems = batch_size * 3 * num_atoms; - if (tid >= total_grad_elems) return; - - int j = tid % num_atoms; - int d = (tid / num_atoms) % 3; - int b = tid / (3 * num_atoms); - - float isw = is_weights[b]; - float lp = current_lp[tid]; - float proj = projected[tid]; - - /* Cross-entropy gradient: d/d_logits(-Σ proj * lp) = exp(lp) - proj */ - float d_combined = isw * (expf(lp) - proj); - - /* Entropy regularization: d/d_logits(-coeff * H(p)) = coeff * (1 + lp) - * where H(p) = -Σ p * log(p) and p = exp(lp) (log-probs). - * This prevents C51 distributions from collapsing to a single atom. - * CRITICAL: clamp lp to [-10, 0] to prevent gradient explosion from - * near-zero probability atoms (lp=-50 → gradient term = -0.49 per atom). */ - float entropy_coeff = {entropy_coeff:.6}f; - if (entropy_coeff > 0.0f) {{ - float lp_clamped = fmaxf(lp, -10.0f); /* floor: p ≥ exp(-10) ≈ 4.5e-5 */ - d_combined += entropy_coeff * (1.0f + lp_clamped); - }} - - // Route through dueling: d_value[b,j] += d_combined - atomicAdd(&d_value_logits[b * num_atoms + j], d_combined); - - // Factored action decode - int factored = actions[b]; - int max_action = b0_size * b1_size * b2_size; - if (factored < 0 || factored >= max_action) factored = 0; - int branch_sizes[3]; - branch_sizes[0] = b0_size; - branch_sizes[1] = b1_size; - branch_sizes[2] = b2_size; - - int a_d; - if (d == 0) a_d = factored / (b1_size * b2_size); - else if (d == 1) a_d = (factored / b2_size) % b1_size; - else a_d = factored % b2_size; - - int A_d = branch_sizes[d]; - float inv_A = 1.0f / (float)A_d; - - // Compute per-branch adv offset - int branch_offset = 0; - for (int dd = 0; dd < d; dd++) - branch_offset += branch_sizes[dd] * num_atoms; - - // d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d) - for (int a = 0; a < A_d; a++) {{ - float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); - int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j; - atomicAdd(&d_adv_logits[adv_idx], d_combined * dueling_grad); - }} -}} -"#); - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(&src, &context) - .map_err(|e| MLError::ModelError(format!("c51_grad_kernel compilation: {e}")))?; + let ptx = Ptx::from_binary(C51_GRAD_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("c51_grad module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("c51_grad cubin load: {e}")))?; module.load_function("c51_grad_kernel") .map_err(|e| MLError::ModelError(format!("c51_grad_kernel load: {e}"))) } -/// Compile the MSE loss kernel on expected Q-values (warmup before C51). -/// -/// Same signature as `c51_loss_batched` but computes: -/// 1. Online E[Q] = sum_j softmax(dueling_logits[j]) * z_j -/// 2. Target E[Q] via Double DQN (online selects, target evaluates) -/// 3. Bellman: target_Q = reward + gamma * (1-done) * target_E[Q] -/// 4. MSE loss: 0.5 * (online_E[Q] - target_Q)^2 per branch, averaged -/// -/// Saves softmax probs in save_current_lp (for gradient kernel). -/// Saves [td_error, E_Q, 0...] per branch per sample in save_projected. +/// Load the MSE loss kernel from precompiled cubin (ZERO runtime nvcc). fn compile_mse_loss_kernel( stream: &Arc, - config: &GpuDqnTrainConfig, + _config: &GpuDqnTrainConfig, ) -> Result { - let defines = format!( - "#define NUM_ATOMS {na}\n\ - #define V_MIN ({v_min:.1}f)\n\ - #define V_MAX ({v_max:.1}f)\n\ - #define BRANCH_0_SIZE {b0}\n\ - #define BRANCH_1_SIZE {b1}\n\ - #define BRANCH_2_SIZE {b2}\n", - na = config.num_atoms, - v_min = config.v_min, - v_max = config.v_max, - b0 = config.branch_0_size, - b1 = config.branch_1_size, - b2 = config.branch_2_size, - ); - let kernel_src = include_str!("mse_loss_kernel.cu"); - let full_source = format!("{defines}\n{kernel_src}"); - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) - .map_err(|e| MLError::ModelError(format!("mse_loss_kernel compilation: {e}")))?; + let ptx = Ptx::from_binary(MSE_LOSS_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("mse_loss module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("mse_loss cubin load: {e}")))?; module.load_function("mse_loss_batched") .map_err(|e| MLError::ModelError(format!("mse_loss_batched load: {e}"))) } -/// Compile the MSE loss gradient kernel through softmax expectation. -/// -/// For each sample [b], branch [d], atom [j]: -/// d_logit_j = td_error * is_weight * p_j * (z_j - E[Q]) -/// where p_j = softmax prob (from save_current_lp), z_j = support value, -/// E[Q] = expected Q (from save_projected), td_error (from save_projected). -/// -/// Routes through dueling: same pattern as c51_grad_kernel. +/// Load the MSE loss gradient kernel from precompiled cubin. +/// v_min/v_max are now runtime kernel parameters (not compile-time #define). fn compile_mse_grad_kernel( stream: &Arc, - config: &GpuDqnTrainConfig, + _config: &GpuDqnTrainConfig, ) -> Result { - let src = format!(r#" -extern "C" __global__ void mse_grad_kernel( - const float* __restrict__ save_probs, // [B, 3, NA] softmax probs - const float* __restrict__ save_eq_td, // [B, 3, NA] layout: [td_error, E_Q, 0, ...] - const float* __restrict__ is_weights, // [B] - const int* __restrict__ actions, // [B] factored - float* __restrict__ d_value_logits, // [B, NA] - float* __restrict__ d_adv_logits, // [B, (B0+B1+B2)*NA] - int batch_size, - int num_atoms, - int b0_size, int b1_size, int b2_size, - int total_branch_atoms) -{{ - int tid = blockIdx.x * blockDim.x + threadIdx.x; - int total_grad_elems = batch_size * 3 * num_atoms; - if (tid >= total_grad_elems) return; - - int j = tid % num_atoms; - int d = (tid / num_atoms) % 3; - int b = tid / (3 * num_atoms); - - float isw = is_weights[b]; - float p_j = save_probs[tid]; /* softmax prob for atom j */ - - /* Read td_error and E[Q] from save_eq_td. - * Layout: for each [b, d, :], element [0] = td_error, element [1] = E[Q] */ - int base = (b * 3 + d) * num_atoms; - float td_error = save_eq_td[base + 0]; - float e_q = save_eq_td[base + 1]; - - /* C51 support: z_j = v_min + j * delta_z */ - float v_min = {v_min:.1}f; - float v_max = {v_max:.1}f; - float delta_z = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; - float z_j = v_min + (float)j * delta_z; - - /* Gradient of MSE loss through softmax expectation: - * L = 0.5 * td^2 where td = E[Q] - target_Q and E[Q] = sum_j(p_j * z_j) - * dL/d_logit_j = td * d(E[Q])/d_logit_j - * = td * p_j * (z_j - E[Q]) - * (softmax Jacobian: dE[Q]/d_logit_j = p_j * (z_j - sum_k(p_k * z_k)) = p_j * (z_j - E[Q])) - */ - float d_combined = isw * td_error * p_j * (z_j - e_q); - - // Route through dueling: d_value[b,j] += d_combined - atomicAdd(&d_value_logits[b * num_atoms + j], d_combined); - - // Factored action decode - int factored = actions[b]; - int max_action = b0_size * b1_size * b2_size; - if (factored < 0 || factored >= max_action) factored = 0; - int branch_sizes[3]; - branch_sizes[0] = b0_size; - branch_sizes[1] = b1_size; - branch_sizes[2] = b2_size; - - int a_d; - if (d == 0) a_d = factored / (b1_size * b2_size); - else if (d == 1) a_d = (factored / b2_size) % b1_size; - else a_d = factored % b2_size; - - int A_d = branch_sizes[d]; - float inv_A = 1.0f / (float)A_d; - - // Compute per-branch adv offset - int branch_offset = 0; - for (int dd = 0; dd < d; dd++) - branch_offset += branch_sizes[dd] * num_atoms; - - // d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d) - for (int a = 0; a < A_d; a++) {{ - float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); - int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j; - atomicAdd(&d_adv_logits[adv_idx], d_combined * dueling_grad); - }} -}} -"#, - v_min = config.v_min, - v_max = config.v_max, - ); - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(&src, &context) - .map_err(|e| MLError::ModelError(format!("mse_grad_kernel compilation: {e}")))?; + let ptx = Ptx::from_binary(MSE_GRAD_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("mse_grad module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("mse_grad cubin load: {e}")))?; module.load_function("mse_grad_kernel") .map_err(|e| MLError::ModelError(format!("mse_grad_kernel load: {e}"))) } -/// Compile the `compute_expected_q` NVRTC kernel. -/// -/// Converts C51 distributional logits (value + advantage, branching dueling) -/// to expected Q-values for the ad-hoc validation forward pass. -/// One thread per sample; iterates over branches and atoms. +/// Load the expected Q kernel from precompiled cubin. fn compile_expected_q_kernel( stream: &Arc, ) -> Result { - let src = r#" -extern "C" __global__ void compute_expected_q( - const float* __restrict__ v_logits, // [N, num_atoms] - const float* __restrict__ b_logits, // [N, (b0+b1+b2)*num_atoms] - float* __restrict__ q_values, // [N, b0+b1+b2] - int N, int num_atoms, - int b0_size, int b1_size, int b2_size, - float v_min, float v_max) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= N) return; - - int total_actions = b0_size + b1_size + b2_size; - float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; - - // Value logits for this sample: [num_atoms] - const float* val = v_logits + (long long)i * num_atoms; - - int branch_sizes[3]; - branch_sizes[0] = b0_size; - branch_sizes[1] = b1_size; - branch_sizes[2] = b2_size; - - int adv_offset = 0; - int q_offset = 0; - for (int d = 0; d < 3; d++) { - int bd = branch_sizes[d]; - for (int a = 0; a < bd; a++) { - const float* adv = b_logits + (long long)i * total_actions * num_atoms - + (long long)(adv_offset + a) * num_atoms; - - // Compute mean advantage logit sum for this branch (for dueling centering) - float mean_adv_sum = 0.0f; - for (int aa = 0; aa < bd; aa++) { - const float* adv_aa = b_logits + (long long)i * total_actions * num_atoms - + (long long)(adv_offset + aa) * num_atoms; - for (int j = 0; j < num_atoms; j++) { - mean_adv_sum += adv_aa[j]; - } - } - float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms); - - // Numerically stable log_softmax over combined = val[j] + adv[j] - mean_adv_per_atom - float max_logit = -1e30f; - for (int j = 0; j < num_atoms; j++) { - float combined = val[j] + adv[j] - mean_adv_per_atom; - if (combined > max_logit) max_logit = combined; - } - float sum_exp = 0.0f; - for (int j = 0; j < num_atoms; j++) { - float combined = val[j] + adv[j] - mean_adv_per_atom; - sum_exp += expf(combined - max_logit); - } - float log_sum = logf(sum_exp + 1e-8f) + max_logit; - - // Expected Q = sum_j( softmax_j * z_j ) - float eq = 0.0f; - for (int j = 0; j < num_atoms; j++) { - float combined = val[j] + adv[j] - mean_adv_per_atom; - float p = expf(combined - log_sum); - float z = v_min + (float)j * dz; - eq += p * z; - } - - q_values[(long long)i * total_actions + q_offset + a] = eq; - } - adv_offset += bd; - q_offset += bd; - } -} -"#; - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) - .map_err(|e| MLError::ModelError(format!("compute_expected_q compilation: {e}")))?; + let ptx = Ptx::from_binary(EXPECTED_Q_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("compute_expected_q module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("expected_q cubin load: {e}")))?; module.load_function("compute_expected_q") .map_err(|e| MLError::ModelError(format!("compute_expected_q load: {e}"))) } -/// Compile the Q-value statistics reduction kernel. -/// -/// Single-block kernel: computes [avg_max_q, q_min, q_max, q_mean, q_variance] -/// from the q_out_buf [N, total_actions]. Grid: (1,1,1), Block: (256,1,1). +/// Load the Q-value statistics reduction kernel from precompiled cubin. fn compile_q_stats_kernel( stream: &Arc, ) -> Result { - let src = r#" -extern "C" __global__ void q_stats_reduce( - const float* __restrict__ q_values, // [N, total_actions] - float* __restrict__ out, // [5]: avg_max_q, q_min, q_max, q_mean, q_var - int N, - int total_actions) -{ - // Single block — thread 0 does the reduction sequentially. - // N is tiny (<=10 samples), total_actions=11. Total work: 110 elements. - if (threadIdx.x != 0) return; - - float global_min = 1e30f; - float global_max = -1e30f; - float global_sum = 0.0f; - float sum_max_q = 0.0f; - int total = N * total_actions; - - for (int i = 0; i < N; i++) { - float row_max = -1e30f; - for (int a = 0; a < total_actions; a++) { - float v = q_values[i * total_actions + a]; - if (v < global_min) global_min = v; - if (v > global_max) global_max = v; - if (v > row_max) row_max = v; - global_sum += v; - } - sum_max_q += row_max; - } - - float mean = (total > 0) ? global_sum / (float)total : 0.0f; - float var_sum = 0.0f; - for (int i = 0; i < total; i++) { - float d = q_values[i] - mean; - var_sum += d * d; - } - float variance = (total > 0) ? var_sum / (float)total : 0.0f; - float avg_max = (N > 0) ? sum_max_q / (float)N : 0.0f; - - out[0] = avg_max; - out[1] = global_min; - out[2] = global_max; - out[3] = mean; - out[4] = variance; -} -"#; - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) - .map_err(|e| MLError::ModelError(format!("q_stats_reduce compilation: {e}")))?; + let ptx = Ptx::from_binary(Q_STATS_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("q_stats module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("q_stats cubin load: {e}")))?; module.load_function("q_stats_reduce") .map_err(|e| MLError::ModelError(format!("q_stats_reduce load: {e}"))) } @@ -4454,28 +4019,18 @@ fn dtod_from_bf16( dtod_copy(dst_ptr, src_ptr, num_elements * std::mem::size_of::(), stream, 0, ctx) } -/// Compile ensemble aggregate + diversity kernels from `ensemble_kernels.cu`. +/// Load ensemble aggregate + diversity kernels from precompiled cubin. /// /// Returns `(ensemble_aggregate_kernel, ensemble_diversity_kernel)`. -/// Both kernels take all sizes as runtime arguments — no network-dim #defines needed. +/// Both kernels take all sizes as runtime arguments. pub(crate) fn compile_ensemble_kernels( stream: &Arc, - state_dim: usize, + _state_dim: usize, ) -> Result<(CudaFunction, CudaFunction), MLError> { - let dim_overrides = format!( - "#define STATE_DIM {state_dim}\n\ - #define MARKET_DIM 42\n\ - #define PORTFOLIO_DIM 8\n", - ); - let common_src = include_str!("common_device_functions.cuh"); - let kernel_src = include_str!("ensemble_kernels.cu"); - let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}"); - let context = stream.context(); - let ptx: Ptx = crate::cuda_pipeline::compile_ptx_for_device(&full_source, &context) - .map_err(|e| MLError::ModelError(format!("ensemble_kernels compilation: {e}")))?; + let ptx = Ptx::from_binary(ENSEMBLE_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("ensemble_kernels module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("ensemble cubin load: {e}")))?; let aggregate = module .load_function("ensemble_aggregate_kernel") @@ -4535,157 +4090,11 @@ pub(crate) fn launch_bf16_to_f32( fn compile_cql_logit_grad_kernel( stream: &Arc, ) -> Result { - let src = r#" -extern "C" __global__ void cql_logit_grad_kernel( - const float* __restrict__ v_logits, // [N, num_atoms] - const float* __restrict__ adv_logits, // [N, total_actions * num_atoms] - const int* __restrict__ actions, // [N] factored action indices (0-44) - float* __restrict__ d_v_logits, // [N, num_atoms] output - float* __restrict__ d_adv_logits, // [N, total_actions * num_atoms] output - float cql_alpha, - int N, int num_atoms, - int b0_size, int b1_size, int b2_size, - float v_min, float v_max) -{ - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i >= N) return; - - int total_actions = b0_size + b1_size + b2_size; - float dz = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; - - // Decode factored action: action = a0 * b1 * b2 + a1 * b2 + a2 - int action = actions[i]; - int a2_taken = action % b2_size; - int a1_taken = (action / b2_size) % b1_size; - int a0_taken = action / (b1_size * b2_size); - - // Clamp to valid range (safety) - if (a0_taken >= b0_size) a0_taken = b0_size - 1; - if (a1_taken >= b1_size) a1_taken = b1_size - 1; - if (a2_taken >= b2_size) a2_taken = b2_size - 1; - - int branch_taken[3]; - branch_taken[0] = a0_taken; - branch_taken[1] = a1_taken; - branch_taken[2] = a2_taken; - - int branch_sizes[3]; - branch_sizes[0] = b0_size; - branch_sizes[1] = b1_size; - branch_sizes[2] = b2_size; - - const float* val = v_logits + (long long)i * num_atoms; - - // Zero value logit gradient accumulator - float d_val_accum[256]; // max num_atoms (C51: 51, generous buffer) - for (int j = 0; j < num_atoms && j < 256; j++) d_val_accum[j] = 0.0f; - - int adv_offset = 0; - for (int d = 0; d < 3; d++) { - int bd = branch_sizes[d]; - int a_taken_d = branch_taken[d]; - - // Step 1: Compute expected Q for each action in this branch - // using dueling: combined_logit[j] = val[j] + adv[a,j] - mean_adv - float q_values_branch[9]; // max branch size = 5 (exposure), 3+3=6 total per branch - - // Compute mean advantage per atom (for dueling centering) - float mean_adv_sum = 0.0f; - for (int aa = 0; aa < bd; aa++) { - const float* adv_aa = adv_logits + (long long)i * total_actions * num_atoms - + (long long)(adv_offset + aa) * num_atoms; - for (int j = 0; j < num_atoms; j++) mean_adv_sum += adv_aa[j]; - } - float mean_adv_per_atom = mean_adv_sum / (float)(bd * num_atoms); - - for (int a = 0; a < bd; a++) { - const float* adv = adv_logits + (long long)i * total_actions * num_atoms - + (long long)(adv_offset + a) * num_atoms; - - // softmax over atoms + expected value - float max_logit = -1e30f; - for (int j = 0; j < num_atoms; j++) { - float c = val[j] + adv[j] - mean_adv_per_atom; - if (c > max_logit) max_logit = c; - } - float sum_exp = 0.0f; - for (int j = 0; j < num_atoms; j++) { - sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit); - } - float log_sum = logf(sum_exp + 1e-8f) + max_logit; - - float eq = 0.0f; - for (int j = 0; j < num_atoms; j++) { - float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum); - float z = v_min + (float)j * dz; - eq += p * z; - } - q_values_branch[a] = eq; - } - - // Step 2: CQL penalty gradient w.r.t. Q-values - // dCQL/dQ[a] = softmax_Q[a] - I(a == a_taken) - float max_q = q_values_branch[0]; - for (int a = 1; a < bd; a++) - if (q_values_branch[a] > max_q) max_q = q_values_branch[a]; - float sum_exp_q = 0.0f; - for (int a = 0; a < bd; a++) - sum_exp_q += expf(q_values_branch[a] - max_q); - - float d_cql_dq[9]; - for (int a = 0; a < bd; a++) { - float softmax_q = expf(q_values_branch[a] - max_q) / (sum_exp_q + 1e-8f); - d_cql_dq[a] = cql_alpha * (softmax_q - ((a == a_taken_d) ? 1.0f : 0.0f)); - } - - // Step 3: Chain rule through softmax-expectation to get d_logits - for (int a = 0; a < bd; a++) { - const float* adv = adv_logits + (long long)i * total_actions * num_atoms - + (long long)(adv_offset + a) * num_atoms; - float* d_adv = d_adv_logits + (long long)i * total_actions * num_atoms - + (long long)(adv_offset + a) * num_atoms; - - // Recompute p[j] for this action - float max_logit = -1e30f; - for (int j = 0; j < num_atoms; j++) { - float c = val[j] + adv[j] - mean_adv_per_atom; - if (c > max_logit) max_logit = c; - } - float sum_exp = 0.0f; - for (int j = 0; j < num_atoms; j++) { - sum_exp += expf(val[j] + adv[j] - mean_adv_per_atom - max_logit); - } - float log_sum = logf(sum_exp + 1e-8f) + max_logit; - - float eq = q_values_branch[a]; - - for (int j = 0; j < num_atoms; j++) { - float p = expf(val[j] + adv[j] - mean_adv_per_atom - log_sum); - float z = v_min + (float)j * dz; - // d_combined_logit[j] = d_cql_dq[a] * p * (z - Q) - float d_combined = d_cql_dq[a] * p * (z - eq); - - // Split combined gradient to adv and val - d_adv[j] = d_combined; - if (j < 256) d_val_accum[j] += d_combined; - } - } - adv_offset += bd; - } - - // Write accumulated value logit gradient (summed across all branches and actions) - float* d_val = d_v_logits + (long long)i * num_atoms; - for (int j = 0; j < num_atoms && j < 256; j++) { - d_val[j] = d_val_accum[j]; - } -} -"#; - let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context) - .map_err(|e| MLError::ModelError(format!("cql_logit_grad_kernel compilation: {e}")))?; + let ptx = Ptx::from_binary(CQL_GRAD_CUBIN.to_vec()); let module = context.load_module(ptx) - .map_err(|e| MLError::ModelError(format!("cql_logit_grad module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("cql_grad cubin load: {e}")))?; module.load_function("cql_logit_grad_kernel") .map_err(|e| MLError::ModelError(format!("cql_logit_grad_kernel load: {e}"))) } + diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 1068c1f5f..3e9a42a21 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -1681,76 +1681,16 @@ fn compile_fill_episode_ids_kernel( .map_err(|e| MLError::ModelError(format!("load fill_episode_ids_kernel: {e}"))) } -/// Compile the trade_stats_reduce kernel that sums portfolio_states[N, 14:19] -/// across all N episodes into a 6-float output. -/// -/// Uses the same `compile_ptx_for_device` (nvcc) path as other experience kernels. +/// Load the trade_stats_reduce kernel from precompiled cubin (ZERO runtime nvcc). fn compile_trade_stats_kernel( stream: &Arc, ) -> Result { - let kernel_src = r#" -// Trade stats reduction kernel. -// Reads portfolio_states[N, PORTFOLIO_STRIDE=20] and sums fields [14:19] -// (win_count, loss_count, sum_wins, sum_losses, sum_returns, sum_sq_returns) -// across all N episodes into a 6-float output buffer. -// Uses shared-memory warp reduction (same pattern as monitoring_reduce). - -#define PORTFOLIO_STRIDE 20 - -extern "C" __global__ void trade_stats_reduce( - const float* __restrict__ portfolio_states, // [N * PORTFOLIO_STRIDE] - float* out, // [6]: sums of ps[14..19] - int N -) { - __shared__ float s_sums[6]; - - int tid = threadIdx.x; - int stride = blockDim.x; - - // Init shared memory - if (tid == 0) { - for (int k = 0; k < 6; k++) s_sums[k] = 0.0f; - } - __syncthreads(); - - // Thread-local accumulators - float local_sums[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; - - for (int i = tid; i < N; i += stride) { - int base = i * PORTFOLIO_STRIDE; - for (int k = 0; k < 6; k++) { - local_sums[k] += portfolio_states[base + 14 + k]; - } - } - - // Warp-level reduction before atomic to shared (reduces contention 32x) - for (int mask = 16; mask >= 1; mask >>= 1) { - for (int k = 0; k < 6; k++) { - local_sums[k] += __shfl_xor_sync(0xFFFFFFFF, local_sums[k], mask); - } - } - // Only lane 0 of each warp atomicAdd - if ((tid & 31) == 0) { - for (int k = 0; k < 6; k++) { - atomicAdd(&s_sums[k], local_sums[k]); - } - } - __syncthreads(); - - // Thread 0 writes output - if (tid == 0) { - for (int k = 0; k < 6; k++) { - out[k] = s_sums[k]; - } - } -} -"#; + static TRADE_STATS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/trade_stats_kernel.cubin")); let context = stream.context(); - let ptx = crate::cuda_pipeline::compile_ptx_for_device(kernel_src, &context) - .map_err(|e| MLError::ModelError(format!("trade_stats_reduce compilation: {e}")))?; + let ptx = cudarc::nvrtc::Ptx::from_binary(TRADE_STATS_CUBIN.to_vec()); let module = context .load_module(ptx) - .map_err(|e| MLError::ModelError(format!("trade_stats_reduce module load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("trade_stats cubin load: {e}")))?; module .load_function("trade_stats_reduce") .map_err(|e| MLError::ModelError(format!("load trade_stats_reduce: {e}"))) diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index e57518160..7a1186063 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -789,16 +789,12 @@ extern "C" __global__ void iqn_cvar_kernel( scales_out[i] = scale; } "#; - // Compile once, cache in OnceLock - use std::sync::OnceLock; - static CVAR_PTX: OnceLock> = OnceLock::new(); + // Load from precompiled cubin (ZERO runtime nvcc) + static CVAR_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/iqn_cvar_kernel.cubin")); let context = self.stream.context(); - let ptx = CVAR_PTX.get_or_init(|| { - crate::cuda_pipeline::compile_ptx_for_device(CVAR_KERNEL_SRC, &context) - }); - let ptx = ptx.as_ref().map_err(|e| MLError::ModelError(format!("CVaR kernel PTX: {e}")))?; - let module = context.load_module(ptx.clone()) - .map_err(|e| MLError::ModelError(format!("CVaR kernel module: {e}")))?; + let ptx = cudarc::nvrtc::Ptx::from_binary(CVAR_CUBIN.to_vec()); + let module = context.load_module(ptx) + .map_err(|e| MLError::ModelError(format!("CVaR cubin load: {e}")))?; let cvar_kernel = module.load_function("iqn_cvar_kernel") .map_err(|e| MLError::ModelError(format!("CVaR kernel load: {e}")))?; diff --git a/crates/ml/src/cuda_pipeline/gpu_portfolio.rs b/crates/ml/src/cuda_pipeline/gpu_portfolio.rs index b75d9a369..aeae6f9ec 100644 --- a/crates/ml/src/cuda_pipeline/gpu_portfolio.rs +++ b/crates/ml/src/cuda_pipeline/gpu_portfolio.rs @@ -63,15 +63,13 @@ pub struct GpuPortfolioSimResult { pub batch_size: usize, } +/// Precompiled experience_kernels cubin (build.rs — ZERO runtime nvcc). +static PORTFOLIO_EXP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/experience_kernels.cubin")); + fn compile_experience_kernels(context: &Arc) -> Result { - let kernel_src = include_str!("experience_kernels.cu"); - 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}" - )) - })?; + let ptx = Ptx::from_binary(PORTFOLIO_EXP_CUBIN.to_vec()); let module = context.load_module(ptx).map_err(|e| { - MLError::ModelError(format!("Failed to load experience kernel module: {e}")) + MLError::ModelError(format!("Failed to load experience kernel cubin: {e}")) })?; module.load_function("portfolio_sim_kernel").map_err(|e| { MLError::ModelError(format!("Failed to load portfolio_sim_kernel: {e}")) diff --git a/crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu b/crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu new file mode 100644 index 000000000..e757f153b --- /dev/null +++ b/crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu @@ -0,0 +1,76 @@ +/** + * IQN CVaR (Conditional Value-at-Risk) position sizing kernel. + * + * Computes CVaR on GPU -- zero CPU readback. + * One thread per sample: insertion-sort quantile values in registers, + * compute CVaR as mean of lowest alpha_count values. + * + * Scale: [0.25, 1.0]. CVaR >= 0 -> full size. CVaR < 0 -> reduce. + * + * Launch config: grid=(ceil(B/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ void iqn_cvar_kernel( + const float* __restrict__ q_values, // [B, N_TAU, TBA] + const int* __restrict__ actions, // [B] factored action indices + float* scales_out, // [B] output CVaR scales + int B, int N_TAU, int TBA, int B0_SIZE, int B1B2, float ALPHA +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= B) return; + + // Decode exposure index from factored action + int exposure_idx = actions[i] / B1B2; + if (exposure_idx >= B0_SIZE) exposure_idx = B0_SIZE - 1; + + // Extract quantile values for this sample's exposure action + // q_values layout: [B, N_TAU, TBA] + int alpha_count = (int)(ALPHA * (float)N_TAU); + if (alpha_count < 1) alpha_count = 1; + + // Find the alpha_count smallest quantile values using partial sort. + // For N_TAU=32 and alpha_count=1-2, just find the minimum(s). + float cvar_sum = 0.0f; + if (alpha_count == 1) { + // Fast path: just find the minimum + float min_val = 1e30f; + for (int t = 0; t < N_TAU; t++) { + float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx]; + if (v < min_val) min_val = v; + } + cvar_sum = min_val; + } else { + // General path: find alpha_count smallest values + // Use insertion into a small sorted array (alpha_count <= ~5 for typical alpha=0.05) + float sorted[8]; // max alpha_count for alpha=0.25, N_TAU=32 + int sorted_len = 0; + for (int t = 0; t < N_TAU; t++) { + float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx]; + if (sorted_len < alpha_count) { + // Insert into sorted array + int pos = sorted_len; + while (pos > 0 && sorted[pos-1] > v) { + sorted[pos] = sorted[pos-1]; + pos--; + } + sorted[pos] = v; + sorted_len++; + } else if (v < sorted[sorted_len-1]) { + // Replace the largest in sorted + int pos = sorted_len - 1; + while (pos > 0 && sorted[pos-1] > v) { + sorted[pos] = sorted[pos-1]; + pos--; + } + sorted[pos] = v; + } + } + for (int j = 0; j < alpha_count; j++) cvar_sum += sorted[j]; + } + + float cvar = cvar_sum / (float)alpha_count; + + // Scale: [0.25, 1.0]. CVaR >= 0 -> full size. CVaR < 0 -> reduce. + float scale = (cvar >= 0.0f) ? 1.0f : fmaxf(0.25f, 1.0f + cvar * 5.0f); + scales_out[i] = scale; +} diff --git a/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu new file mode 100644 index 000000000..9b128b743 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/mse_grad_kernel.cu @@ -0,0 +1,89 @@ +/** + * MSE loss gradient kernel through softmax expectation. + * + * For each sample [b], branch [d], atom [j]: + * d_logit_j = td_error * is_weight * p_j * (z_j - E[Q]) + * where p_j = softmax prob (from save_current_lp), z_j = support value, + * E[Q] = expected Q (from save_projected), td_error (from save_projected). + * + * v_min and v_max are passed as runtime kernel parameters. + * + * Routes through dueling: same pattern as c51_grad_kernel. + * + * Launch config: grid=(ceil(batch_size*3*num_atoms/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ void mse_grad_kernel( + const float* __restrict__ save_probs, // [B, 3, NA] softmax probs + const float* __restrict__ save_eq_td, // [B, 3, NA] layout: [td_error, E_Q, 0, ...] + const float* __restrict__ is_weights, // [B] + const int* __restrict__ actions, // [B] factored + float* __restrict__ d_value_logits, // [B, NA] + float* __restrict__ d_adv_logits, // [B, (B0+B1+B2)*NA] + int batch_size, + int num_atoms, + int b0_size, int b1_size, int b2_size, + int total_branch_atoms, + float v_min, float v_max) +{ + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int total_grad_elems = batch_size * 3 * num_atoms; + if (tid >= total_grad_elems) return; + + int j = tid % num_atoms; + int d = (tid / num_atoms) % 3; + int b = tid / (3 * num_atoms); + + float isw = is_weights[b]; + float p_j = save_probs[tid]; /* softmax prob for atom j */ + + /* Read td_error and E[Q] from save_eq_td. + * Layout: for each [b, d, :], element [0] = td_error, element [1] = E[Q] */ + int base = (b * 3 + d) * num_atoms; + float td_error = save_eq_td[base + 0]; + float e_q = save_eq_td[base + 1]; + + /* C51 support: z_j = v_min + j * delta_z */ + float delta_z = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f; + float z_j = v_min + (float)j * delta_z; + + /* Gradient of MSE loss through softmax expectation: + * L = 0.5 * td^2 where td = E[Q] - target_Q and E[Q] = sum_j(p_j * z_j) + * dL/d_logit_j = td * d(E[Q])/d_logit_j + * = td * p_j * (z_j - E[Q]) + * (softmax Jacobian: dE[Q]/d_logit_j = p_j * (z_j - sum_k(p_k * z_k)) = p_j * (z_j - E[Q])) + */ + float d_combined = isw * td_error * p_j * (z_j - e_q); + + // Route through dueling: d_value[b,j] += d_combined + atomicAdd(&d_value_logits[b * num_atoms + j], d_combined); + + // Factored action decode + int factored = actions[b]; + int max_action = b0_size * b1_size * b2_size; + if (factored < 0 || factored >= max_action) factored = 0; + int branch_sizes[3]; + branch_sizes[0] = b0_size; + branch_sizes[1] = b1_size; + branch_sizes[2] = b2_size; + + int a_d; + if (d == 0) a_d = factored / (b1_size * b2_size); + else if (d == 1) a_d = (factored / b2_size) % b1_size; + else a_d = factored % b2_size; + + int A_d = branch_sizes[d]; + float inv_A = 1.0f / (float)A_d; + + // Compute per-branch adv offset + int branch_offset = 0; + for (int dd = 0; dd < d; dd++) + branch_offset += branch_sizes[dd] * num_atoms; + + // d_adv[b, d, a, j] = d_combined * (delta(a, a_d) - 1/A_d) + for (int a = 0; a < A_d; a++) { + float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A); + int adv_idx = b * total_branch_atoms + branch_offset + a * num_atoms + j; + atomicAdd(&d_adv_logits[adv_idx], d_combined * dueling_grad); + } +} diff --git a/crates/ml/src/cuda_pipeline/per_update_kernel.cu b/crates/ml/src/cuda_pipeline/per_update_kernel.cu new file mode 100644 index 000000000..ff50adcc9 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/per_update_kernel.cu @@ -0,0 +1,39 @@ +/** + * PER (Prioritized Experience Replay) priority update kernel. + * + * Computes |td|^alpha + epsilon, scatter-writes to priorities buffer, + * and atomicMax for batch-wide maximum priority. + * + * Launch config: grid=(ceil(batch_size/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ +void per_update_priorities_kernel( + const float* __restrict__ td_errors, + const unsigned int* __restrict__ indices, + float* __restrict__ priorities, + float* __restrict__ batch_max, + float alpha, + float epsilon, + int batch_size, + int capacity) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + + float td = fabsf(td_errors[i]); + float new_prio = powf(td, alpha) + epsilon; + + // Clamp to [epsilon, 1e6] — same bounds as GpuReplayBuffer::update_priorities_gpu + if (new_prio < epsilon) new_prio = epsilon; + if (new_prio > 1e6f) new_prio = 1e6f; + + // Scatter-write priority at buffer index (bounds-checked) + unsigned int idx = indices[i]; + if (idx < (unsigned int)capacity) + priorities[idx] = new_prio; + + // atomicMax for batch max: IEEE 754 positive floats preserve int ordering + int ival = __float_as_int(new_prio); + atomicMax((int*)batch_max, ival); +} diff --git a/crates/ml/src/cuda_pipeline/q_stats_kernel.cu b/crates/ml/src/cuda_pipeline/q_stats_kernel.cu new file mode 100644 index 000000000..10a253956 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/q_stats_kernel.cu @@ -0,0 +1,52 @@ +/** + * Q-value statistics reduction kernel. + * + * Single-block kernel: computes [avg_max_q, q_min, q_max, q_mean, q_variance] + * from the q_out_buf [N, total_actions]. + * + * Grid: (1,1,1), Block: (256,1,1). + */ + +extern "C" __global__ void q_stats_reduce( + const float* __restrict__ q_values, // [N, total_actions] + float* __restrict__ out, // [5]: avg_max_q, q_min, q_max, q_mean, q_var + int N, + int total_actions) +{ + // Single block -- thread 0 does the reduction sequentially. + // N is tiny (<=10 samples), total_actions=11. Total work: 110 elements. + if (threadIdx.x != 0) return; + + float global_min = 1e30f; + float global_max = -1e30f; + float global_sum = 0.0f; + float sum_max_q = 0.0f; + int total = N * total_actions; + + for (int i = 0; i < N; i++) { + float row_max = -1e30f; + for (int a = 0; a < total_actions; a++) { + float v = q_values[i * total_actions + a]; + if (v < global_min) global_min = v; + if (v > global_max) global_max = v; + if (v > row_max) row_max = v; + global_sum += v; + } + sum_max_q += row_max; + } + + float mean = (total > 0) ? global_sum / (float)total : 0.0f; + float var_sum = 0.0f; + for (int i = 0; i < total; i++) { + float d = q_values[i] - mean; + var_sum += d * d; + } + float variance = (total > 0) ? var_sum / (float)total : 0.0f; + float avg_max = (N > 0) ? sum_max_q / (float)N : 0.0f; + + out[0] = avg_max; + out[1] = global_min; + out[2] = global_max; + out[3] = mean; + out[4] = variance; +} diff --git a/crates/ml/src/cuda_pipeline/relu_mask_kernel.cu b/crates/ml/src/cuda_pipeline/relu_mask_kernel.cu new file mode 100644 index 000000000..fc3ad8eed --- /dev/null +++ b/crates/ml/src/cuda_pipeline/relu_mask_kernel.cu @@ -0,0 +1,17 @@ +/** + * Standalone ReLU mask kernel for IQN trunk gradient. + * + * dx[i] *= (activation[i] > 0.0f) + * + * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1). + */ + +extern "C" __global__ +void relu_mask_standalone(float* __restrict__ dx, + const float* __restrict__ activation, + int n) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n) return; + if (activation[i] <= 0.0f) dx[i] = 0.0f; +} diff --git a/crates/ml/src/cuda_pipeline/trade_stats_kernel.cu b/crates/ml/src/cuda_pipeline/trade_stats_kernel.cu new file mode 100644 index 000000000..91dcc33cf --- /dev/null +++ b/crates/ml/src/cuda_pipeline/trade_stats_kernel.cu @@ -0,0 +1,60 @@ +/** + * Trade stats reduction kernel. + * + * Reads portfolio_states[N, PORTFOLIO_STRIDE=20] and sums fields [14:19] + * (win_count, loss_count, sum_wins, sum_losses, sum_returns, sum_sq_returns) + * across all N episodes into a 6-float output buffer. + * Uses shared-memory warp reduction (same pattern as monitoring_reduce). + * + * Launch config: grid=(1, 1, 1), block=(256, 1, 1). + */ + +#define PORTFOLIO_STRIDE 20 + +extern "C" __global__ void trade_stats_reduce( + const float* __restrict__ portfolio_states, // [N * PORTFOLIO_STRIDE] + float* out, // [6]: sums of ps[14..19] + int N +) { + __shared__ float s_sums[6]; + + int tid = threadIdx.x; + int stride = blockDim.x; + + // Init shared memory + if (tid == 0) { + for (int k = 0; k < 6; k++) s_sums[k] = 0.0f; + } + __syncthreads(); + + // Thread-local accumulators + float local_sums[6] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f}; + + for (int i = tid; i < N; i += stride) { + int base = i * PORTFOLIO_STRIDE; + for (int k = 0; k < 6; k++) { + local_sums[k] += portfolio_states[base + 14 + k]; + } + } + + // Warp-level reduction before atomic to shared (reduces contention 32x) + for (int mask = 16; mask >= 1; mask >>= 1) { + for (int k = 0; k < 6; k++) { + local_sums[k] += __shfl_xor_sync(0xFFFFFFFF, local_sums[k], mask); + } + } + // Only lane 0 of each warp atomicAdd + if ((tid & 31) == 0) { + for (int k = 0; k < 6; k++) { + atomicAdd(&s_sums[k], local_sums[k]); + } + } + __syncthreads(); + + // Thread 0 writes output + if (tid == 0) { + for (int k = 0; k < 6; k++) { + out[k] = s_sums[k]; + } + } +} diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index bf504bead..65321d558 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -385,14 +385,12 @@ impl FusedTrainingCtx { diversity_weight = hyperparams.ensemble_diversity_weight, ); - // Compile KL gradient kernel for diversity gradient flow + // Load KL gradient kernel from precompiled cubin (ZERO runtime nvcc) let kl_grad_kernel = { - let module = stream.context().load_module( - crate::cuda_pipeline::compile_ptx_for_device( - include_str!("../../cuda_pipeline/ensemble_kernels.cu"), - stream.context(), - ).map_err(|e| anyhow::anyhow!("ensemble_kl_grad compilation: {e}"))? - ).map_err(|e| anyhow::anyhow!("ensemble_kl_grad module: {e}"))?; + static ENSEMBLE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/ensemble_kernels.cubin")); + let ptx = cudarc::nvrtc::Ptx::from_binary(ENSEMBLE_CUBIN.to_vec()); + let module = stream.context().load_module(ptx) + .map_err(|e| anyhow::anyhow!("ensemble_kl_grad cubin load: {e}"))?; module.load_function("ensemble_kl_gradient_kernel") .map_err(|e| anyhow::anyhow!("ensemble_kl_gradient_kernel load: {e}"))? };