diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 5fb55d164..8306c384f 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -1519,3 +1519,70 @@ extern "C" __global__ void popart_normalize_kernel( } } } + +/* ══════════════════════════════════════════════════════════════════════ + * SUPERVISED PRE-TRAINING KERNEL (v8 C1) + * + * Cross-entropy logit gradients for the exposure branch against the + * optimal single-bar direction. Runs at epoch 0 to initialize exposure + * weights with directional knowledge before RL begins. + * + * Target: argmax_k(position_k * (raw_next - raw_close)) + * Same softmax + cross-entropy math as exposure_aux_grad_kernel, + * but target computed from raw prices instead of from the CEA loop. + * + * Launch: grid=(ceil(batch_size/256)), block=(256). + * ══════════════════════════════════════════════════════════════════════ */ +extern "C" __global__ void exposure_pretrain_step( + const __nv_bfloat16* __restrict__ targets, /* [total_bars, 4] raw prices */ + const int* __restrict__ bar_indices, /* [batch] which bar each sample maps to */ + const float* __restrict__ adv_logits, /* [batch, b0_size * num_atoms] from forward */ + float* __restrict__ grad_scratch, /* [batch * b0_size * num_atoms] output */ + int b0_size, int num_atoms, int batch_size, + float max_position +) { + int sample = blockIdx.x * blockDim.x + threadIdx.x; + if (sample >= batch_size) return; + + int bar = bar_indices[sample]; + + /* Compute optimal direction from raw prices */ + float raw_close = __bfloat162float(targets[bar * 4 + 2]); + float raw_next = __bfloat162float(targets[bar * 4 + 3]); + float price_delta = raw_next - raw_close; + + /* Find best exposure: which position maximizes single-bar PnL? */ + int best_k = b0_size / 2; /* default: Flat (middle index) */ + float best_pnl = 0.0f; /* Flat PnL = 0 */ + for (int k = 0; k < b0_size; k++) { + float step_size = (b0_size > 1) ? 2.0f / (float)(b0_size - 1) : 0.0f; + float fraction = -1.0f + (float)k * step_size; + float pos = fraction * max_position; + float pnl = pos * price_delta; + if (pnl > best_pnl) { best_pnl = pnl; best_k = k; } + } + + /* Cross-entropy gradient: softmax(logit) - one_hot(target) */ + const float* logits = adv_logits + (long long)sample * b0_size * num_atoms; + float inv_batch = 1.0f / (float)batch_size; + + for (int z = 0; z < num_atoms; z++) { + /* Stable softmax over b0_size for atom z */ + float max_l = -1e30f; + for (int a = 0; a < b0_size; a++) { + float l = logits[a * num_atoms + z]; + if (l > max_l) max_l = l; + } + float sum_exp = 0.0f; + for (int a = 0; a < b0_size; a++) { + sum_exp += expf(logits[a * num_atoms + z] - max_l); + } + float inv_sum = 1.0f / fmaxf(sum_exp, 1e-10f); + + for (int a = 0; a < b0_size; a++) { + float prob = expf(logits[a * num_atoms + z] - max_l) * inv_sum; + float grad = (prob - ((a == best_k) ? 1.0f : 0.0f)) * inv_batch; + grad_scratch[(long long)sample * b0_size * num_atoms + a * num_atoms + z] = grad; + } + } +}