fix(gpu): dynamic SHMEM_MAX_IN_DIM prevents shared memory overflow on H100

Root cause: SHMEM_MAX_IN_DIM was hardcoded to 256, but on H100 with
hidden_dim_base=2048, SHARED_H1=SHARED_H2=512. The cooperative tile
loader wrote 64×512=32768 floats into shmem allocated for 64×256=16384,
causing CUDA_ERROR_INVALID_VALUE on kernel launch.

Fixes:
- Make SHMEM_MAX_IN_DIM injectable via NVRTC #define (was hardcoded)
- Compute max(state_dim, shared_h1, shared_h2) and inject at compile time
- Use dynamic value for host-side shmem_bytes allocation
- Shrink eps_out[NOISY_MAX_DIM] → eps_out[SHMEM_TILE_ROWS] (saves 768B/thread)
- Force smoke tests to Device::Cpu (CUDA driver sensitivity to binary layout)

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

crates/ml/src/cuda_pipeline/common_device_functions.cuh

@@ -153,7 +153,9 @@ __device__ void noisy_matvec_leaky_relu(
 #ifndef SHMEM_TILE_ROWS
 #define SHMEM_TILE_ROWS 64
 #endif
+#ifndef SHMEM_MAX_IN_DIM
 #define SHMEM_MAX_IN_DIM 256  /* max(STATE_DIM, SHARED_H1, SHARED_H2) */
+#endif
 
 /**
  * Cooperatively load a weight tile from global to shared memory.
@@ -229,19 +231,20 @@ __device__ void noisy_matvec_leaky_relu_shmem(
     (void)out_dim;
     float sigma_scale = sigma_init / sqrtf((float)in_dim);
 
-    /* Generate factorized noise vectors for this tile */
+    /* Generate factorized noise vectors for this tile.
+     * eps_out only needs tile_rows elements (≤ SHMEM_TILE_ROWS=64). */
     float eps_in[NOISY_MAX_DIM];
-    float eps_out[NOISY_MAX_DIM];
+    float eps_out[SHMEM_TILE_ROWS];
 
     for (int i = 0; i < in_dim && i < NOISY_MAX_DIM; i++) {
         eps_in[i] = factorized_noise_fn(gpu_random_gaussian(rng));
     }
-    for (int j = 0; j < tile_rows && j < NOISY_MAX_DIM; j++) {
+    for (int j = 0; j < tile_rows && j < SHMEM_TILE_ROWS; j++) {
         eps_out[j] = factorized_noise_fn(gpu_random_gaussian(rng));
     }
 
     for (int j = 0; j < tile_rows; j++) {
-        float ej = (j < NOISY_MAX_DIM) ? eps_out[j] : 0.0f;
+        float ej = (j < SHMEM_TILE_ROWS) ? eps_out[j] : 0.0f;
         float acc = shmem_b[j] + sigma_scale * ej;
         const float* row = shmem_W + j * in_dim;
         for (int i = 0; i < in_dim; i++) {

crates/ml/src/cuda_pipeline/gpu_experience_collector.rs

@@ -195,6 +195,8 @@ pub struct GpuExperienceCollector {
     kernel_func: CudaFunction,
     /// Actual STATE_DIM injected into the kernel (matches network input_dim).
     state_dim: usize,
+    /// Network hidden dims — needed for shared memory tile sizing at launch time.
+    network_dims: (usize, usize, usize, usize),
     /// Number of episodes buffers were allocated for (from config, not MAX constant).
     alloc_episodes: usize,
     /// Number of timesteps buffers were allocated for (from config, not MAX constant).
@@ -273,6 +275,9 @@ impl GpuExperienceCollector {
         // feature buffer indexing, and scratch buffer sizes match the real data layout.
         let common_src = include_str!("common_device_functions.cuh");
         let kernel_src = include_str!("dqn_experience_kernel.cu");
+        // Shared memory tile must fit the widest input dimension across all layers.
+        // s1: in=STATE_DIM, s2: in=SHARED_H1, v1/a1: in=SHARED_H2 — take the max.
+        let shmem_max_in_dim = state_dim.max(shared_h1).max(shared_h2);
         let dim_overrides = format!(
             "#define STATE_DIM {state_dim}\n\
              #define MARKET_DIM {market_dim}\n\
@@ -281,7 +286,8 @@ impl GpuExperienceCollector {
              #define SHARED_H2 {shared_h2}\n\
              #define VALUE_H {value_h}\n\
              #define ADV_H {adv_h}\n\
-             #define NUM_ATOMS_MAX {num_atoms_max}\n"
+             #define NUM_ATOMS_MAX {num_atoms_max}\n\
+             #define SHMEM_MAX_IN_DIM {shmem_max_in_dim}\n"
         );
         let full_source = format!("{common_src}\n{dim_overrides}\n{kernel_src}");
         info!(
@@ -459,6 +465,7 @@ impl GpuExperienceCollector {
             stream,
             kernel_func,
             state_dim,
+            network_dims,
             alloc_episodes,
             alloc_timesteps,
             online_weights,
@@ -543,10 +550,11 @@ impl GpuExperienceCollector {
 
         // ---- Step 4: Launch config ----
         // Shared memory for weight tiling: tile_rows * max_in_dim floats (weights)
-        // + tile_rows floats (bias). Default: (64 * 256 + 64) * 4 = 65,792 bytes.
+        // + tile_rows floats (bias). max_in_dim = max(state_dim, shared_h1, shared_h2).
         const SHMEM_TILE_ROWS: u32 = 64;
-        const SHMEM_MAX_IN_DIM: u32 = 256;
-        let shmem_bytes = (SHMEM_TILE_ROWS * SHMEM_MAX_IN_DIM + SHMEM_TILE_ROWS)
+        let (sh1, sh2, _, _) = self.network_dims;
+        let shmem_max_in = self.state_dim.max(sh1).max(sh2) as u32;
+        let shmem_bytes = (SHMEM_TILE_ROWS * shmem_max_in + SHMEM_TILE_ROWS)
             * std::mem::size_of::<f32>() as u32;
 
         let n = n_episodes as u32;

crates/ml/src/trainers/dqn/smoke_tests/helpers.rs

@@ -64,14 +64,14 @@ pub(super) fn safe_device() -> Device {
     Device::Cpu
 }
 
-/// Create trainer on a safe device (CPU fallback if CUDA broken)
+/// Create trainer on CPU (smoke tests must be deterministic and GPU-independent).
 pub(super) fn cpu_trainer() -> anyhow::Result<DQNTrainer> {
-    DQNTrainer::new_with_device(smoke_params(), safe_device())
+    DQNTrainer::new_with_device(smoke_params(), Device::Cpu)
 }
 
-/// Create trainer with custom params on safe device
+/// Create trainer with custom params on CPU
 pub(super) fn cpu_trainer_with(params: DQNHyperparameters) -> anyhow::Result<DQNTrainer> {
-    DQNTrainer::new_with_device(params, safe_device())
+    DQNTrainer::new_with_device(params, Device::Cpu)
 }
 
 /// Assert a value is finite (not NaN or Inf)