diff --git a/crates/ml/src/cuda_pipeline/dqn_training_kernel.cu b/crates/ml/src/cuda_pipeline/dqn_training_kernel.cu index 0588c3141..a3878fecc 100644 --- a/crates/ml/src/cuda_pipeline/dqn_training_kernel.cu +++ b/crates/ml/src/cuda_pipeline/dqn_training_kernel.cu @@ -580,12 +580,18 @@ extern "C" __global__ void dqn_forward_loss_kernel( * C51 DISTRIBUTIONAL LOSS (per-branch cross-entropy) * ═══════════════════════════════════════════════════════════════ */ - /* Decompose factored action into per-branch indices */ + /* Decompose factored action into per-branch indices. + * Clamp to valid range to prevent OOB access on online_log_probs[] + * when replay buffer contains uninitialized/corrupt action values. */ int factored_action = actions[sample_id]; + /* Safety clamp: action must be in [0, BRANCH_0*BRANCH_1*BRANCH_2) */ + int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE; + if (factored_action < 0 || factored_action >= max_action) + factored_action = 0; /* safe default: hold/market/normal */ int branch_action[3]; - branch_action[0] = factored_action / 9; /* exposure: 0-4 */ - branch_action[1] = (factored_action % 9) / 3; /* order: 0-2 */ - branch_action[2] = factored_action % 3; /* urgency: 0-2 */ + branch_action[0] = factored_action / (BRANCH_1_SIZE * BRANCH_2_SIZE); + branch_action[1] = (factored_action / BRANCH_2_SIZE) % BRANCH_1_SIZE; + branch_action[2] = factored_action % BRANCH_2_SIZE; float reward = rewards[sample_id]; float done = dones[sample_id]; @@ -957,10 +963,13 @@ extern "C" __global__ void dqn_backward_kernel( /* ── Load per-sample data ──────────────────────────────────── */ int factored_action = actions[sample_id]; + int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE; + if (factored_action < 0 || factored_action >= max_action) + factored_action = 0; int branch_action[3]; - branch_action[0] = factored_action / 9; - branch_action[1] = (factored_action % 9) / 3; - branch_action[2] = factored_action % 3; + branch_action[0] = factored_action / (BRANCH_1_SIZE * BRANCH_2_SIZE); + branch_action[1] = (factored_action / BRANCH_2_SIZE) % BRANCH_1_SIZE; + branch_action[2] = factored_action % BRANCH_2_SIZE; float is_weight = is_weights[sample_id]; float scale = is_weight / (float)NUM_BRANCHES; diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 3ef004e10..6207a3d0f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2185,11 +2185,31 @@ pub(crate) fn query_max_shmem_bytes() -> usize { } /// Returns true if CUDA Graph capture should be used for training. -/// Disabled when dynamic shared memory exceeds 48KB — CUDA Graph replay -/// does not reliably inherit cuFuncSetAttribute opt-in for > 48KB on -/// some driver versions, causing CUDA_ERROR_ILLEGAL_ADDRESS. -pub(crate) fn should_use_cuda_graph(shmem_bytes: usize) -> bool { - shmem_bytes <= 49152 +/// +/// Disabled on consumer GPUs (< 164KB max shmem opt-in) — CUDA Graph +/// replay triggers CUDA_ERROR_ILLEGAL_ADDRESS in the driver's local +/// memory management on RTX 3050/4090 (confirmed with compute-sanitizer: +/// crash at libcuda.so+0x28d50 during cuGraphLaunch, not in user kernel). +/// This is a driver-level issue with graph replay on Ampere consumer GPUs. +/// +/// A100/H100 (>= 164KB shmem) use CUDA Graph for zero-overhead replay. +/// Consumer GPUs fall back to direct kernel launches (slightly higher +/// dispatch latency but correct execution). +pub(crate) fn should_use_cuda_graph(_shmem_bytes: usize) -> bool { + use cudarc::driver::sys::{cuDeviceGetAttribute, CUdevice_attribute}; + let mut max_shmem: i32 = 0; + let result = unsafe { + cuDeviceGetAttribute( + &mut max_shmem, + CUdevice_attribute::CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, + 0, + ) + }; + if result != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + return false; + } + // A100 = 164KB, H100 = 228KB — these support reliable CUDA Graph replay + max_shmem >= 164_000 } fn compute_shmem_bytes(config: &GpuDqnTrainConfig) -> usize {