From 2cfbc2633c84e09eb3eca2a8b145fed40ee11c43 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 12 Mar 2026 21:38:38 +0100 Subject: [PATCH] fix(cuda): remove hardcoded NUM_ATOMS_MAX cap, fix NVRTC source order - Remove hardcoded cap of NUM_ATOMS_MAX=51 on sm_90+. Now that TMA uses proper mbarrier sync, the INVALID_PTX was caused by wrong instructions, not kernel size. Hyperparams control atom count directly. - Fix NVRTC source concatenation order: dim_overrides now comes BEFORE common_src and kernel_src so the #ifndef guards in .cuh/.cu files properly skip defaults. Previously overrides came after, causing macro redefinition warnings. - Make PORTFOLIO_DIM dynamic (was hardcoded "3" in format string). Co-Authored-By: Claude Opus 4.6 --- .../cuda_pipeline/gpu_experience_collector.rs | 27 ++++++------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index d5e5139ed..dbd08a8f1 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -381,37 +381,26 @@ impl GpuExperienceCollector { let pow2 = (max_tile as u32).next_power_of_two() >> 1; pow2.clamp(16, 256) as usize }; - // Cap NUM_ATOMS_MAX at 51 on Hopper (sm_90+) to reduce per-lane stack - // pressure in the warp kernel. The C51 distributional path allocates - // adv_atoms[NUM_ACTIONS * NUM_ATOMS_MAX] per lane — at 100 atoms this is - // 2000 bytes/lane which can cause CUDA_ERROR_INVALID_PTX during JIT. - // 51 atoms (original C51 paper) → 1020 bytes/lane, well within limits. - let kernel_atoms_max = if sm_major >= 9 && num_atoms_max > 51 { - info!( - "Capping NUM_ATOMS_MAX from {num_atoms_max} to 51 for sm_{sm_major}0 warp kernel \ - (runtime num_atoms will be clamped inside the kernel)" - ); - 51 - } else { - num_atoms_max - }; + let portfolio_dim = 3usize; // cash, position_size, unrealized_pnl let dim_overrides = format!( "#define NOISY_MAX_DIM {noisy_max_dim}\n\ #define STATE_DIM {state_dim}\n\ #define MARKET_DIM {market_dim}\n\ - #define PORTFOLIO_DIM 3\n\ + #define PORTFOLIO_DIM {portfolio_dim}\n\ #define SHARED_H1 {shared_h1}\n\ #define SHARED_H2 {shared_h2}\n\ #define VALUE_H {value_h}\n\ #define ADV_H {adv_h}\n\ - #define NUM_ATOMS_MAX {kernel_atoms_max}\n\ + #define NUM_ATOMS_MAX {num_atoms_max}\n\ #define SHMEM_MAX_IN_DIM {shmem_max_in_dim}\n\ #define SHMEM_TILE_ROWS {shmem_tile_rows}\n" ); - let full_source = format!("{common_src}\n{dim_overrides}\n{kernel_src}"); + // dim_overrides BEFORE common_src so #ifndef guards in .cuh/.cu + // skip their defaults — no macro redefinition warnings from NVRTC. + let full_source = format!("{dim_overrides}\n{common_src}\n{kernel_src}"); info!( "GPU experience collector: compiling kernel with dims state={state_dim} market={market_dim} \ - shared=[{shared_h1},{shared_h2}] value={value_h} adv={adv_h} atoms_max={kernel_atoms_max} \ + shared=[{shared_h1},{shared_h2}] value={value_h} adv={adv_h} atoms_max={num_atoms_max} \ shmem_max_in={shmem_max_in_dim} shmem_tile_rows={shmem_tile_rows}" ); let context = stream.context(); @@ -463,7 +452,7 @@ impl GpuExperienceCollector { ); let disable_tma = "#define DISABLE_TMA 1\n"; let full_source_no_tma = format!( - "{common_src}\n{disable_tma}{dim_overrides}\n{kernel_src}" + "{disable_tma}{dim_overrides}\n{common_src}\n{kernel_src}" ); let ptx_no_tma = crate::cuda_pipeline::compile_ptx_for_device( &full_source_no_tma, &context,