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

fix: ensemble_aggregate_kernel OOB — buffers sized for total_actions(12) not num_atoms(51)

Browse Source 
ensemble_mean_q_buf and ensemble_var_q_buf were allocated as
batch_size * total_actions (12), but the kernel writes
batch_size * num_atoms (51) elements. 2977 OOB write errors.
Fixed: allocate batch_size * num_atoms. compute-sanitizer: 0 errors.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-15 23:52:24 +02:00
parent 8a54c8a32c
commit bd8b84a2a7
1 changed files with 7 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -244,10 +244,10 @@ pub(crate) struct FusedTrainingCtx {
/// F32 because output layer GemmEx writes f32 (no f32 truncation overflow).
/// None when ensemble_count <= 1.
pub(crate) ensemble_logits_buf: Option<cudarc::driver::CudaSlice<f32>>,
/// Pre-allocated buffer: [B * total_actions] for mean Q-values across K heads.
/// Pre-allocated buffer: [B * num_atoms] for mean Q-values across K heads.
/// None when ensemble_count <= 1.
pub(crate) ensemble_mean_q_buf: Option<cudarc::driver::CudaSlice<f32>>,
/// Pre-allocated buffer: [B * total_actions] for Q-value variance across K heads.
/// Pre-allocated buffer: [B * num_atoms] for Q-value variance across K heads.
/// None when ensemble_count <= 1.
pub(crate) ensemble_var_q_buf: Option<cudarc::driver::CudaSlice<f32>>,
/// Pre-allocated buffer: [1] for accumulated diversity loss scalar.
@@ -596,12 +596,14 @@ impl FusedTrainingCtx {
// Pre-allocate ensemble buffers.
// ensemble_logits_buf: [K * B * num_atoms] — all K heads' value logits
let na = dqn.config.num_atoms;
let total_actions = dqn.config.num_actions + dqn.config.num_order_types + dqn.config.num_urgency_levels;
// total_actions removed — ensemble buffers use num_atoms, not total_actions.
let logits_buf = stream.alloc_zeros::<f32>(k * batch_size * na)
.map_err(|e| anyhow::anyhow!("Alloc ensemble_logits_buf f32: {e}"))?;
let mean_q_buf = stream.alloc_zeros::<f32>(batch_size * total_actions)
// mean_q_buf and var_q_buf: [B * num_atoms] — aggregate kernel outputs
// per-atom mean/variance across K heads (NOT per-action).
let mean_q_buf = stream.alloc_zeros::<f32>(batch_size * na)
.map_err(|e| anyhow::anyhow!("Alloc ensemble_mean_q_buf: {e}"))?;
let var_q_buf = stream.alloc_zeros::<f32>(batch_size * total_actions)
let var_q_buf = stream.alloc_zeros::<f32>(batch_size * na)
.map_err(|e| anyhow::anyhow!("Alloc ensemble_var_q_buf: {e}"))?;
let div_loss_buf = stream.alloc_zeros::<f32>(1)
.map_err(|e| anyhow::anyhow!("Alloc ensemble_diversity_loss_buf: {e}"))?;