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

debug: add attention tracing + sync barriers to isolate H100 segfault

Browse Source 
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-18 12:23:05 +02:00
parent f8f1147de7
commit 3ff5bd696d
2 changed files with 16 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/ml/src/cuda_pipeline/gpu_attention.rs
View File

@@ -217,11 +217,8 @@ impl GpuAttention {
// ── Load cubin kernels ──────────────────────────────────────────
// Forward cubin — old multihead_feature_attention (legacy, not used by cuBLAS path)
let _module = context.load_cubin(ATTENTION_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("attention module: {e}")))?;
// Backward cubin — contains ALL element-wise kernels (fwd + bwd) for cuBLAS attention
// All element-wise kernels (fwd SDP + LayerNorm, bwd, bias, grad_norm, adam)
// live in the backward cubin after the cuBLAS rewrite.
let bwd_module = context.load_cubin(ATTENTION_BACKWARD_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("attention backward module: {e}")))?;
let sdp_fwd_kernel = bwd_module.load_function("attn_sdp_fwd")

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

@@ -1234,16 +1234,27 @@ impl FusedTrainingCtx {
// Attention forward + backward + Adam.
if let Some(ref mut attn) = self.gpu_attention {
tracing::info!("ATTN_DEBUG: starting attention forward (batch={})", self.batch_size);
// Sync stream to catch GPU errors before the crash
self.trainer.sync_all_streams()
.map_err(|e| anyhow::anyhow!("pre-attn sync: {e}"))?;
self.trainer.apply_attention_forward(attn, self.batch_size)
.map_err(|e| anyhow::anyhow!("Attention forward: {e}"))?;
let d_h_s2_bf16 = self.trainer.bw_d_h_s2()
.map_err(|e| anyhow::anyhow!("bw_d_h_s2 f32 cast: {e}"))?;
attn.backward(d_h_s2_bf16, self.batch_size)
tracing::info!("ATTN_DEBUG: forward done, starting backward");
self.trainer.sync_all_streams()
.map_err(|e| anyhow::anyhow!("post-attn-fwd sync: {e}"))?;
let d_h_s2_ref = self.trainer.bw_d_h_s2()
.map_err(|e| anyhow::anyhow!("bw_d_h_s2 ref: {e}"))?;
attn.backward(d_h_s2_ref, self.batch_size)
.map_err(|e| anyhow::anyhow!("Attention backward: {e}"))?;
tracing::info!("ATTN_DEBUG: backward done, starting adam");
self.trainer.sync_all_streams()
.map_err(|e| anyhow::anyhow!("post-attn-bwd sync: {e}"))?;
let lr = self.trainer.config().lr;
let mgn = self.trainer.config().max_grad_norm;
attn.adam_step(lr, mgn)
.map_err(|e| anyhow::anyhow!("Attention Adam: {e}"))?;
tracing::info!("ATTN_DEBUG: attention complete");
}