gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(cuda): remove hardcoded NUM_ATOMS_MAX cap, fix NVRTC source order

Browse Source 
- 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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-12 21:38:38 +01:00
parent 8494853eb0
commit 2cfbc2633c
1 changed files with 8 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -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,