diff --git a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs index da6f69952..abd11e5bb 100644 --- a/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs +++ b/crates/ml/src/cuda_pipeline/gpu_iqn_head.rs @@ -658,48 +658,117 @@ impl GpuIqnHead { .map_err(|e| MLError::ModelError(format!("IQN CVaR forward: {e}")))?; } - // 3. Compute CVaR on CPU (small: B scalars from B*N*TBA quantiles) - // Read quantile Q-values and actions to host - let q_size = b * n * tba; - let q_view = self.save_q_online.slice(..q_size); - let mut q_host = vec![0.0_f32; q_size]; - self.stream.memcpy_dtoh(&q_view, &mut q_host) // gpu-exit: CVaR readback (~4KB for B=32) - .map_err(|e| MLError::ModelError(format!("IQN CVaR q readback: {e}")))?; - - let mut act_host = vec![0_i32; b]; - let act_view = actions.slice(..b); - self.stream.memcpy_dtoh(&act_view, &mut act_host) // gpu-exit: action readback - .map_err(|e| MLError::ModelError(format!("IQN CVaR action readback: {e}")))?; - - // Compute CVaR per sample - let alpha_count = ((alpha * n as f32) as usize).max(1); - let mut cvar_scales = vec![1.0_f32; b]; - for i in 0..b { - // Extract quantile values for this sample's exposure action - let exposure_idx = (act_host[i] / 9) as usize; // factored: exposure * 9 + order * 3 + urgency - let exposure_idx = exposure_idx.min(self.config.branch_0_size - 1); - - let mut quantiles: Vec = (0..n) - .map(|t| q_host[i * n * tba + t * tba + exposure_idx]) - .collect(); - quantiles.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); - - // CVaR = mean of lowest alpha_count quantiles - let cvar: f32 = quantiles[..alpha_count].iter().sum::() / alpha_count as f32; - - // Scale: [0.25, 1.0]. CVaR ≥ 0 → full size. CVaR < 0 → reduce. - cvar_scales[i] = if cvar >= 0.0 { - 1.0 - } else { - (1.0 + cvar * 5.0).clamp(0.25, 1.0) - }; - } - - // Upload scales to GPU + // 3. Compute CVaR on GPU — ZERO CPU readback. + // Inline CUDA kernel: one thread per sample, insertion-sort 32 quantiles + // in registers, compute CVaR as mean of lowest alpha_count values. let mut scales_buf = self.stream.alloc_zeros::(b) .map_err(|e| MLError::ModelError(format!("IQN CVaR alloc scales: {e}")))?; - self.stream.memcpy_htod(&cvar_scales, &mut scales_buf) - .map_err(|e| MLError::ModelError(format!("IQN CVaR upload scales: {e}")))?; + let alpha_f = alpha; + let n_i32 = n as i32; + let tba_i32 = tba as i32; + let b0_i32 = self.config.branch_0_size as i32; + let b1b2 = (self.config.branch_1_size * self.config.branch_2_size) as i32; + let cvar_blocks = (b + 255) / 256; + let cvar_config = LaunchConfig { + grid_dim: (cvar_blocks as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + // Compile the CVaR kernel inline (cached after first call) + static CVAR_KERNEL_SRC: &str = r#" +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; +} +"#; + // Compile once, cache in OnceLock + use std::sync::OnceLock; + static CVAR_PTX: OnceLock> = OnceLock::new(); + 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 cvar_kernel = module.load_function("iqn_cvar_kernel") + .map_err(|e| MLError::ModelError(format!("CVaR kernel load: {e}")))?; + + unsafe { + self.stream + .launch_builder(&cvar_kernel) + .arg(&self.save_q_online) + .arg(actions) + .arg(&mut scales_buf) + .arg(&batch_i32) + .arg(&n_i32) + .arg(&tba_i32) + .arg(&b0_i32) + .arg(&b1b2) + .arg(&alpha_f) + .launch(cvar_config) + .map_err(|e| MLError::ModelError(format!("CVaR kernel launch: {e}")))?; + } Ok(scales_buf) }