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

fix: attention kernel load from wrong cubin + build.rs bf16 label

Browse Source 
attn_sdp_fwd and attn_layer_norm_fwd were loaded from attention_kernel.cubin
(which only contains the legacy multihead_feature_attention) instead of
attention_backward_kernel.cubin where they actually live. This caused
CUDA_ERROR_NOT_FOUND on H100, silently disabling attention.

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

2
crates/ml/build.rs
View File

@@ -93,7 +93,7 @@ fn main() {
}
}
let passed = kernels_with_common.len() - failed.len();
eprintln!(" Precompiled {passed}/{} CUDA kernels ({arch}) — all BF16",
eprintln!(" Precompiled {passed}/{} CUDA kernels ({arch}) — f32/TF32",
kernels_with_common.len());
if !failed.is_empty() {
eprintln!(" FAILED: {}", failed.join(", "));

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

@@ -217,17 +217,17 @@ impl GpuAttention {
// ── Load cubin kernels ──────────────────────────────────────────
// Forward cubin — SDP kernel lives here
let module = context.load_cubin(ATTENTION_CUBIN.to_vec())
// 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}")))?;
let sdp_fwd_kernel = module.load_function("attn_sdp_fwd")
.map_err(|e| MLError::ModelError(format!("attn_sdp_fwd kernel load: {e}")))?;
let layer_norm_fwd_kernel = module.load_function("attn_layer_norm_fwd")
.map_err(|e| MLError::ModelError(format!("attn_layer_norm_fwd kernel load: {e}")))?;
// Backward cubin — SDP bwd, LayerNorm bwd, bias grad, grad_norm, adam
// Backward cubin — contains ALL element-wise kernels (fwd + bwd) for cuBLAS attention
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")
.map_err(|e| MLError::ModelError(format!("attn_sdp_fwd kernel load: {e}")))?;
let layer_norm_fwd_kernel = bwd_module.load_function("attn_layer_norm_fwd")
.map_err(|e| MLError::ModelError(format!("attn_layer_norm_fwd kernel load: {e}")))?;
let sdp_bwd_kernel = bwd_module.load_function("attn_sdp_bwd")
.map_err(|e| MLError::ModelError(format!("attn_sdp_bwd kernel load: {e}")))?;
let layer_norm_bwd_kernel = bwd_module.load_function("attn_layer_norm_bwd")