diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index be5f71b6a..5a983f617 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -974,21 +974,14 @@ impl GpuExperienceCollector { let exp_h_b2_f32 = stream.alloc_zeros::(alloc_episodes * network_dims.3) .map_err(|e| MLError::ModelError(format!("alloc exp_h_b2_f32: {e}")))?; - // GPU-native domain randomization kernels - let domain_rand_starts_kernel = { - let context = stream.context(); - let module = context.load_cubin(EXPERIENCE_KERNELS_CUBIN.to_vec()) - .map_err(|e| MLError::ModelError(format!("dr cubin: {e}")))?; - module.load_function("domain_rand_episode_starts") - .map_err(|e| MLError::ModelError(format!("domain_rand_starts load: {e}")))? - }; - let domain_rand_params_kernel = { - let context = stream.context(); - let module = context.load_cubin(EXPERIENCE_KERNELS_CUBIN.to_vec()) - .map_err(|e| MLError::ModelError(format!("dr cubin2: {e}")))?; - module.load_function("domain_rand_sim_params") - .map_err(|e| MLError::ModelError(format!("domain_rand_params load: {e}")))? - }; + // GPU-native domain randomization + saboteur kernels (ONE cubin load) + let exp_module_extra = stream.context() + .load_cubin(EXPERIENCE_KERNELS_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("exp cubin extra: {e}")))?; + let domain_rand_starts_kernel = exp_module_extra.load_function("domain_rand_episode_starts") + .map_err(|e| MLError::ModelError(format!("domain_rand_starts load: {e}")))?; + let domain_rand_params_kernel = exp_module_extra.load_function("domain_rand_sim_params") + .map_err(|e| MLError::ModelError(format!("domain_rand_params load: {e}")))?; let domain_rand_params_buf = stream.alloc_zeros::(alloc_episodes * 5) .map_err(|e| MLError::ModelError(format!("alloc dr_params: {e}")))?; @@ -1004,21 +997,11 @@ impl GpuExperienceCollector { .map_err(|e| MLError::ModelError(format!("alloc saboteur_best: {e}")))?; let saboteur_best_return_buf = stream.alloc_zeros::(1) .map_err(|e| MLError::ModelError(format!("alloc saboteur_best_return: {e}")))?; - // Load saboteur kernels from experience cubin - let saboteur_generate_kernel = { - let context = stream.context(); - let module = context.load_cubin(EXPERIENCE_KERNELS_CUBIN.to_vec()) - .map_err(|e| MLError::ModelError(format!("saboteur cubin: {e}")))?; - module.load_function("saboteur_generate_params") - .map_err(|e| MLError::ModelError(format!("saboteur_generate load: {e}")))? - }; - let saboteur_select_kernel = { - let context = stream.context(); - let module = context.load_cubin(EXPERIENCE_KERNELS_CUBIN.to_vec()) - .map_err(|e| MLError::ModelError(format!("saboteur cubin2: {e}")))?; - module.load_function("saboteur_select_best") - .map_err(|e| MLError::ModelError(format!("saboteur_select load: {e}")))? - }; + // Saboteur kernels from same module (no extra cubin load) + let saboteur_generate_kernel = exp_module_extra.load_function("saboteur_generate_params") + .map_err(|e| MLError::ModelError(format!("saboteur_generate load: {e}")))?; + let saboteur_select_kernel = exp_module_extra.load_function("saboteur_select_best") + .map_err(|e| MLError::ModelError(format!("saboteur_select load: {e}")))?; // #31 Load bottleneck tanh+concat kernel from utility cubin (if active) let bn_tanh_concat_fn = if bn_dim_from_params > 0 { diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 2506132e4..15fc3e80d 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -1107,7 +1107,7 @@ mod tests { let cfg = PpoCollectorConfig::default(); assert_eq!(cfg.n_episodes, 128); assert_eq!(cfg.timesteps_per_episode, 500); - assert_eq!(cfg.total_experiences(), 128_000); // 2x counterfactual + assert_eq!(cfg.total_experiences(), 64_000); // PPO: no counterfactual doubling assert!((cfg.gamma - 0.99).abs() < f32::EPSILON); assert!((cfg.gae_lambda - 0.95).abs() < f32::EPSILON); }