From 3ff5bd696d20720a5fcb6297d61207f09beb5385 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 18 Apr 2026 12:23:05 +0200 Subject: [PATCH] debug: add attention tracing + sync barriers to isolate H100 segfault Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_attention.rs | 7 ++----- crates/ml/src/trainers/dqn/fused_training.rs | 17 ++++++++++++++--- 2 files changed, 16 insertions(+), 8 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index 9f1f89798..269a79013 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -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") diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index cbc03f881..06c37d756 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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"); }