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

fix: NUM_ATOMS 51→201 in c51/mse kernels — array overflow when num_atoms=101

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-26 08:19:19 +01:00
parent 500d8f18cd
commit 9f63193351
21 changed files with 897 additions and 891 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
crates/ml/build.rs
View File

@@ -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.

35
crates/ml/src/cuda_pipeline/backward_kernels.cu Normal file
View File

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

49
crates/ml/src/cuda_pipeline/batched_backward.rs
View File

@@ -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<CudaStream>,
) -> 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")

45
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -687,51 +687,18 @@ fn raw_f32_ptr_mut(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> 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<CudaStream>,
) -> 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")

32
crates/ml/src/cuda_pipeline/bias_kernels.cu Normal file
View File

@@ -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];
}

85
crates/ml/src/cuda_pipeline/c51_grad_kernel.cu Normal file
View File

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

158
crates/ml/src/cuda_pipeline/cql_grad_kernel.cu Normal file
View File

@@ -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];
}
}

43
crates/ml/src/cuda_pipeline/decision_transformer.rs
View File

@@ -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<CudaStream>,
config: &DecisionTransformerConfig,
_config: &DecisionTransformerConfig,
) -> Result<DtKernels, MLError> {
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<CudaFunction, MLError> {
module.load_function(name)
@@ -396,21 +382,14 @@ struct AdamKernels {
adam_update: CudaFunction,
}
fn compile_adam_kernels(stream: &Arc<CudaStream>, state_dim: usize) -> Result<AdamKernels, 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!("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<CudaStream>, _state_dim: usize) -> Result<AdamKernels, MLError> {
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}")))?;

19
crates/ml/src/cuda_pipeline/ema_kernel.cu Normal file
View File

@@ -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];
}
}

78
crates/ml/src/cuda_pipeline/expected_q_kernel.cu Normal file
View File

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

779
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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<CudaStream>,
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<CudaStream>) -> Result<CudaFunction, MLError> {
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<CudaStream>) -> Result<CudaFunction, MLError> {
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<CudaStream>) -> Result<CudaFunction, MLError> {
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<CudaStream>,
config: &GpuDqnTrainConfig,
_config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
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<CudaStream>,
_config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
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<CudaStream>,
config: &GpuDqnTrainConfig,
_config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
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<CudaStream>,
config: &GpuDqnTrainConfig,
_config: &GpuDqnTrainConfig,
) -> Result<CudaFunction, MLError> {
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<CudaStream>,
) -> Result<CudaFunction, MLError> {
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<CudaStream>,
) -> Result<CudaFunction, MLError> {
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::<u16>(), 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<CudaStream>,
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<CudaStream>,
) -> Result<CudaFunction, MLError> {
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}")))
}

68
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -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<CudaStream>,
) -> Result<CudaFunction, MLError> {
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}")))

14
crates/ml/src/cuda_pipeline/gpu_iqn_head.rs
View File

@@ -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<Result<cudarc::nvrtc::Ptx, String>> = 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}")))?;

12
crates/ml/src/cuda_pipeline/gpu_portfolio.rs
View File

@@ -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<CudaContext>) -> Result<CudaFunction, MLError> {
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}"))

76
crates/ml/src/cuda_pipeline/iqn_cvar_kernel.cu Normal file
View File

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

89
crates/ml/src/cuda_pipeline/mse_grad_kernel.cu Normal file
View File

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

39
crates/ml/src/cuda_pipeline/per_update_kernel.cu Normal file
View File

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

52
crates/ml/src/cuda_pipeline/q_stats_kernel.cu Normal file
View File

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

17
crates/ml/src/cuda_pipeline/relu_mask_kernel.cu Normal file
View File

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

60
crates/ml/src/cuda_pipeline/trade_stats_kernel.cu Normal file
View File

@@ -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];
}
}
}

12
crates/ml/src/trainers/dqn/fused_training.rs
View File

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