feat: GPU-native CVaR kernel — eliminated CPU round-trip bottleneck

The compute_cvar_scales() was doing GPU→CPU→sort→CPU→GPU for quantile
CVaR computation. With batch_size=307 × 32 quantiles × 15 actions,
this was ~150KB of DtoH + sort + HtoD per epoch — likely the source
of the 10s/epoch overhead vs 3s baseline.

New inline CUDA kernel (iqn_cvar_kernel):
- One thread per sample (256 threads/block)
- Insertion sort of alpha_count smallest quantiles in registers
- Fast path for alpha_count=1 (just find minimum)
- Zero CPU readback, zero CPU allocation

Compiled once via OnceLock, cached for all subsequent calls.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

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<f32> = (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::<f32>() / 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::<f32>(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<Result<cudarc::nvrtc::Ptx, String>> = 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)
     }