diff --git a/crates/ml/src/cuda_pipeline/gpu_attention.rs b/crates/ml/src/cuda_pipeline/gpu_attention.rs index 03a35a649..afb8ead0b 100644 --- a/crates/ml/src/cuda_pipeline/gpu_attention.rs +++ b/crates/ml/src/cuda_pipeline/gpu_attention.rs @@ -10,20 +10,62 @@ //! Residual connection ensures gradient flow even with frozen attention weights. //! //! Phase B: Full backward pass with gradient flow through attention weights. -//! The backward kernel computes gradients for all attention parameters +//! The backward pass computes gradients for all attention parameters //! (W_Q, W_K, W_V, W_O, biases, LayerNorm gamma/beta) and propagates -//! gradients to the input. Attention weights are updated via a dedicated -//! Adam optimizer (separate from the DQN trunk Adam). +//! gradients to the input. +//! +//! cuBLAS rewrite: Q/K/V/O projections use cublasLtMatmul via SharedCublasHandle. +//! SDP (scaled dot-product attention) and LayerNorm remain as cubin kernels. +//! Zero atomicAdd — all gradient accumulation uses cuBLAS GEMMs with +//! deterministic reduction. use std::sync::Arc; + +use cudarc::cublaslt::result as cublaslt_result; +use cudarc::cublaslt::sys as cublaslt_sys; use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; use tracing::info; + use crate::MLError; +use super::shared_cublas_handle::SharedCublasHandle; /// Precompiled attention kernel cubins, embedded at compile time by build.rs. static ATTENTION_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_kernel.cubin")); static ATTENTION_BACKWARD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/attention_backward_kernel.cubin")); +// ── GEMM descriptor ───────────────────────────────────────────────────────── + +/// Cached cublasLt GEMM descriptor for attention forward/backward projections. +/// Same structure as CachedGemmDesc in batched_forward.rs / CachedBwdGemmDesc +/// in batched_backward.rs. +struct AttnGemmDesc { + matmul_desc: cublaslt_sys::cublasLtMatmulDesc_t, + a_layout: cublaslt_sys::cublasLtMatrixLayout_t, + b_layout: cublaslt_sys::cublasLtMatrixLayout_t, + c_layout: cublaslt_sys::cublasLtMatrixLayout_t, + d_layout: cublaslt_sys::cublasLtMatrixLayout_t, + algo: cublaslt_sys::cublasLtMatmulAlgo_t, +} + +/// SAFETY: cublasLt descriptors are opaque pointers owned exclusively by +/// `GpuAttention`, which is single-threaded in practice (owned by `FusedTrainingCtx`). +unsafe impl Send for AttnGemmDesc {} +unsafe impl Sync for AttnGemmDesc {} + +impl Drop for AttnGemmDesc { + fn drop(&mut self) { + unsafe { + let _ = cublaslt_result::destroy_matrix_layout(self.d_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.c_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.b_layout); + let _ = cublaslt_result::destroy_matrix_layout(self.a_layout); + let _ = cublaslt_result::destroy_matmul_desc(self.matmul_desc); + } + } +} + +// ── Configuration ─────────────────────────────────────────────────────────── + /// Configuration for the GPU attention layer. #[derive(Debug, Clone)] pub struct GpuAttentionConfig { @@ -44,46 +86,110 @@ impl Default for GpuAttentionConfig { impl GpuAttentionConfig { /// Total number of attention parameters. - /// Layout: 3*D*D (W_QKV) + 3*D (b_QKV) + D*D (W_O) + D (b_O) + D (gamma) + D (beta) + /// Layout: W_Q(D*D) + b_Q(D) + W_K(D*D) + b_K(D) + W_V(D*D) + b_V(D) + /// + W_O(D*D) + b_O(D) + gamma(D) + beta(D) pub fn total_params(&self) -> usize { let d = self.state_dim; - // W_Q(D²) + W_K(D²) + W_V(D²) + W_O(D²) + b_Q(D) + b_K(D) + b_V(D) + b_O(D) + ln_gamma(D) + ln_beta(D) 4 * d * d + 6 * d } } +// ── GpuAttention ──────────────────────────────────────────────────────────── + /// GPU multi-head feature attention layer with full backward pass. +/// +/// Q/K/V/O projections use batched cuBLAS (cublasLtMatmul) via SharedCublasHandle. +/// SDP and LayerNorm use precompiled cubin kernels. +/// Zero atomicAdd — all cross-sample gradient reduction is handled by cuBLAS +/// GEMM (dW = dY @ X^T) with alpha=1/B for mean reduction. #[allow(missing_debug_implementations)] // CudaSlice does not implement Debug pub struct GpuAttention { config: GpuAttentionConfig, stream: Arc, - forward_kernel: CudaFunction, - backward_kernel: CudaFunction, - weight_grad_reduce_kernel: CudaFunction, + shared_handle: Arc, + + // ── cuBLAS GEMM descriptors (forward: 4, backward: 8) ───────────── + // Forward: Q[D,B] = W_Q^T @ states[D,B], same for K, V, O + gemm_fwd_q: AttnGemmDesc, + gemm_fwd_k: AttnGemmDesc, + gemm_fwd_v: AttnGemmDesc, + gemm_fwd_o: AttnGemmDesc, + // Backward dW: dW[D,D] = dY[D,B] @ X^T[B,D] (TRANSA=N, TRANSB=T) + gemm_bwd_dw_o: AttnGemmDesc, + gemm_bwd_dw_q: AttnGemmDesc, + gemm_bwd_dw_k: AttnGemmDesc, + gemm_bwd_dw_v: AttnGemmDesc, + // Backward dX: dX[D,B] = W[D,D] @ dY[D,B] (TRANSA=N, TRANSB=N) + gemm_bwd_dx_o: AttnGemmDesc, + gemm_bwd_dx_q: AttnGemmDesc, + gemm_bwd_dx_k: AttnGemmDesc, + gemm_bwd_dx_v: AttnGemmDesc, + + // ── cubin kernels (SDP, LayerNorm, bias, grad_norm, adam) ────────── + sdp_fwd_kernel: CudaFunction, + sdp_bwd_kernel: CudaFunction, + layer_norm_fwd_kernel: CudaFunction, + layer_norm_bwd_kernel: CudaFunction, + bias_grad_reduce_kernel: CudaFunction, grad_norm_phase1_kernel: CudaFunction, grad_norm_phase2_kernel: CudaFunction, adam_kernel: CudaFunction, + + // ── Weight parameters ───────────────────────────────────────────── /// Attention parameters (Xavier-initialized, trainable). + /// Layout: W_Q[D*D] + b_Q[D] + W_K[D*D] + b_K[D] + W_V[D*D] + b_V[D] + /// + W_O[D*D] + b_O[D] + gamma[D] + beta[D] params: CudaSlice, - /// Output buffer for attended states [batch_size, state_dim]. + + // ── Intermediate buffers (forward) ──────────────────────────────── + /// Q projection output [D, B] — saved for backward SDP. + proj_q_buf: CudaSlice, + /// K projection output [D, B] — saved for backward SDP. + proj_k_buf: CudaSlice, + /// V projection output [D, B] — saved for backward SDP. + proj_v_buf: CudaSlice, + /// SDP output [D, B] — pre-output-projection. Saved for backward dW_O. + attn_out_buf: CudaSlice, + /// SDP scores [num_heads, B] — saved for backward SDP. + /// Feature-level attention: each head has head_dim features, attention is + /// computed per-sample across the feature dimension (not across sequence). + /// With B=16384, num_heads=4: 4*16384 = 65536 floats. + sdp_scores_buf: CudaSlice, + /// Output buffer for attended states [D, B] — final output after LayerNorm. output_buf: CudaSlice, - /// Saved input states from forward pass (for backward). + /// Saved input states from forward pass (for backward) [D, B]. saved_input: CudaSlice, - /// Saved Q, K, V projections from forward pass [batch_size, 3, state_dim]. - /// Eliminates recomputation in backward kernel (~14MB on H100). - saved_qkv: CudaSlice, - /// Scratch buffer for input gradient [batch_size, state_dim]. - /// Written by backward kernel via atomicAdd. Not propagated further - /// (attention is the first layer — no upstream to backprop into). + /// LayerNorm saved mean [B]. + ln_save_mean: CudaSlice, + /// LayerNorm saved reciprocal std [B]. + ln_save_rstd: CudaSlice, + /// LayerNorm normalized values [D, B] — saved for backward. + ln_save_xnorm: CudaSlice, + /// Post-projection pre-LayerNorm output [D, B] — saved for backward. + projected_buf: CudaSlice, + + // ── Gradient buffers (backward) ─────────────────────────────────── + /// Gradient for Q projection [D, B]. + d_proj_q_buf: CudaSlice, + /// Gradient for K projection [D, B]. + d_proj_k_buf: CudaSlice, + /// Gradient for V projection [D, B]. + d_proj_v_buf: CudaSlice, + /// Gradient for SDP output [D, B] (before output projection backward). + d_attn_out_buf: CudaSlice, + /// Gradient for projected output [D, B] (after LayerNorm backward). + d_projected_buf: CudaSlice, + /// Scratch buffer for input gradient accumulation [D, B]. + /// Accumulates W_Q^T @ d_Q + W_K^T @ d_K + W_V^T @ d_V. d_input_scratch: CudaSlice, - /// Per-sample gradient accumulator [batch_size * total_params] for deterministic training. - d_params_per_sample: CudaSlice, /// Reduced gradient accumulator for attention parameters [total_params]. d_params: CudaSlice, /// Gradient norm output buffer [1]. grad_norm_buf: CudaSlice, /// Per-block partial sums for two-phase grad norm [num_norm_blocks]. grad_norm_block_sums: CudaSlice, + + // ── Adam optimizer state ────────────────────────────────────────── /// Adam first moment (m) for attention parameters [total_params]. attn_m: CudaSlice, /// Adam second moment (v) for attention parameters [total_params]. @@ -98,36 +204,70 @@ pub struct GpuAttention { impl GpuAttention { pub fn new( - stream: Arc, + shared_handle: Arc, config: GpuAttentionConfig, ) -> Result { + let stream = Arc::clone(&shared_handle.stream); let context = stream.context(); let d = config.state_dim; let total_params = config.total_params(); let b = config.batch_size; - // Load precompiled forward attention cubin + // ── Load cubin kernels ────────────────────────────────────────── + + // Forward cubin — SDP kernel lives here let module = context.load_cubin(ATTENTION_CUBIN.to_vec()) .map_err(|e| MLError::ModelError(format!("attention module: {e}")))?; - let forward_kernel = module.load_function("multihead_feature_attention") - .map_err(|e| MLError::ModelError(format!("attention kernel load: {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}")))?; - // Load precompiled backward attention cubin - + // Backward cubin — SDP bwd, LayerNorm bwd, bias grad, grad_norm, adam let bwd_module = context.load_cubin(ATTENTION_BACKWARD_CUBIN.to_vec()) .map_err(|e| MLError::ModelError(format!("attention backward module: {e}")))?; - let backward_kernel = bwd_module.load_function("attention_backward_kernel") - .map_err(|e| MLError::ModelError(format!("attention backward kernel load: {e}")))?; - let weight_grad_reduce_kernel = bwd_module.load_function("attn_weight_grad_reduce") - .map_err(|e| MLError::ModelError(format!("attention weight_grad_reduce 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") + .map_err(|e| MLError::ModelError(format!("attn_layer_norm_bwd kernel load: {e}")))?; + let bias_grad_reduce_kernel = bwd_module.load_function("attn_bias_grad_reduce") + .map_err(|e| MLError::ModelError(format!("attn_bias_grad_reduce load: {e}")))?; let grad_norm_phase1_kernel = bwd_module.load_function("attn_grad_norm_phase1") - .map_err(|e| MLError::ModelError(format!("attention grad_norm_phase1 load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("attn_grad_norm_phase1 load: {e}")))?; let grad_norm_phase2_kernel = bwd_module.load_function("attn_grad_norm_phase2") - .map_err(|e| MLError::ModelError(format!("attention grad_norm_phase2 load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("attn_grad_norm_phase2 load: {e}")))?; let adam_kernel = bwd_module.load_function("attn_adam_kernel") - .map_err(|e| MLError::ModelError(format!("attention adam kernel load: {e}")))?; + .map_err(|e| MLError::ModelError(format!("attn_adam_kernel load: {e}")))?; + + // ── Create GEMM descriptors ───────────────────────────────────── + + let lt_handle = shared_handle.lt_handle.0; + let ws_size = shared_handle.lt_workspace_size; + + // Forward GEMMs: C[D,B] = W^T[D,D] @ X[D,B] + // TRANSA=T (transpose W), TRANSB=N (X as-is) + // M=D, N=B, K=D, lda=D, ldb=D, ldc=D + let gemm_fwd_q = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "fwd_q")?; + let gemm_fwd_k = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "fwd_k")?; + let gemm_fwd_v = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "fwd_v")?; + let gemm_fwd_o = create_attn_gemm_desc(lt_handle, 1, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "fwd_o")?; + + // Backward dW GEMMs: dW[D,D] = dY[D,B] @ X^T[B,D] + // TRANSA=N, TRANSB=T: M=D, N=D, K=B, lda=D, ldb=D, ldc=D + let gemm_bwd_dw_o = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, d as i32, b as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dw_o")?; + let gemm_bwd_dw_q = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, d as i32, b as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dw_q")?; + let gemm_bwd_dw_k = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, d as i32, b as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dw_k")?; + let gemm_bwd_dw_v = create_attn_gemm_desc(lt_handle, 0, 1, d as i32, d as i32, b as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dw_v")?; + + // Backward dX GEMMs: dX[D,B] = W[D,D] @ dY[D,B] + // TRANSA=N, TRANSB=N: M=D, N=B, K=D, lda=D, ldb=D, ldc=D + let gemm_bwd_dx_o = create_attn_gemm_desc(lt_handle, 0, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dx_o")?; + let gemm_bwd_dx_q = create_attn_gemm_desc(lt_handle, 0, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dx_q")?; + let gemm_bwd_dx_k = create_attn_gemm_desc(lt_handle, 0, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dx_k")?; + let gemm_bwd_dx_v = create_attn_gemm_desc(lt_handle, 0, 0, d as i32, b as i32, d as i32, d as i32, d as i32, d as i32, ws_size, "bwd_dx_v")?; + + // ── Xavier initialization for attention weights ───────────────── - // Xavier initialization for attention weights let mut host_params = vec![0.0_f32; total_params]; let fan_in = d as f32; let fan_out = d as f32; @@ -141,46 +281,64 @@ impl GpuAttention { // Initialize W_Q, W_K, W_V, W_O with Xavier let weight_count = 4 * d * d; for i in 0..weight_count { - // Simple LCG PRNG for initialization rng_state = rng_state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); - let u = (rng_state >> 33) as f32 / (1u64 << 31) as f32; // [0, 1) - // Box-Muller approximation - let normal = (u - 0.5) * 3.46; // ~ uniform -> normal approximation + let u = (rng_state >> 33) as f32 / (1u64 << 31) as f32; + let normal = (u - 0.5) * 3.46; host_params[i] = normal as f32 * xavier_std; } // Biases = 0, LayerNorm gamma = 1, beta = 0 - // Layout: 3*D*D (W_QKV) + 3*D (b_QKV) + D*D (W_O) + D (b_O) + D (gamma) + D (beta) - // weight_count = 4*D*D covers W_QKV + W_O. Biases = 3*D + D = 4*D. - let ln_gamma_start = 4 * d * d + 4 * d; // after all weights + biases + let ln_gamma_start = 4 * d * d + 4 * d; for i in 0..d { - host_params[ln_gamma_start + i] = 1.0; // gamma = 1 + host_params[ln_gamma_start + i] = 1.0; } - // beta stays 0 let mut params = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("attention params alloc: {e}")))?; super::htod_f32(&stream, &host_params, &mut params)?; - let output_buf = stream.alloc_zeros::(b * d) + // ── Allocate intermediate buffers ──────────────────────────────── + + let db = d * b; + let output_buf = stream.alloc_zeros::(db) .map_err(|e| MLError::ModelError(format!("attention output alloc: {e}")))?; - - // Saved state buffer for backward pass (input states) - let saved_input = stream.alloc_zeros::(b * d) + let saved_input = stream.alloc_zeros::(db) .map_err(|e| MLError::ModelError(format!("attention saved_input alloc: {e}")))?; + let proj_q_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention proj_q alloc: {e}")))?; + let proj_k_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention proj_k alloc: {e}")))?; + let proj_v_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention proj_v alloc: {e}")))?; + let attn_out_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention attn_out alloc: {e}")))?; + // SDP scores: for feature-level attention, each head produces one score per sample + // (dot product of head_dim-sized Q and K vectors). Shape: [num_heads, B]. + let sdp_scores_buf = stream.alloc_zeros::(config.num_heads * b) + .map_err(|e| MLError::ModelError(format!("attention sdp_scores alloc: {e}")))?; + let ln_save_mean = stream.alloc_zeros::(b) + .map_err(|e| MLError::ModelError(format!("attention ln_mean alloc: {e}")))?; + let ln_save_rstd = stream.alloc_zeros::(b) + .map_err(|e| MLError::ModelError(format!("attention ln_rstd alloc: {e}")))?; + let ln_save_xnorm = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention ln_xnorm alloc: {e}")))?; + let projected_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention projected alloc: {e}")))?; - // Saved Q, K, V projections from forward pass [B, 3, D] - let saved_qkv = stream.alloc_zeros::(b * 3 * d) - .map_err(|e| MLError::ModelError(format!("attention saved_qkv alloc: {e}")))?; + // ── Gradient buffers ──────────────────────────────────────────── - // Gradient and optimizer buffers - let d_input_scratch = stream.alloc_zeros::(b * d) + let d_proj_q_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention d_proj_q alloc: {e}")))?; + let d_proj_k_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention d_proj_k alloc: {e}")))?; + let d_proj_v_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention d_proj_v alloc: {e}")))?; + let d_attn_out_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention d_attn_out alloc: {e}")))?; + let d_projected_buf = stream.alloc_zeros::(db) + .map_err(|e| MLError::ModelError(format!("attention d_projected alloc: {e}")))?; + let d_input_scratch = stream.alloc_zeros::(db) .map_err(|e| MLError::ModelError(format!("attention d_input_scratch alloc: {e}")))?; - // Tiled per-sample gradients: process min(B,256) samples per tile. - // At B=16384, P~26K: full buffer = 1.6GB. Tiled at 256: 26MB. - let attn_tile_size = b.min(256); - let d_params_per_sample = stream.alloc_zeros::(attn_tile_size * total_params) - .map_err(|e| MLError::ModelError(format!("attention d_params_per_sample alloc: {e}")))?; let d_params = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("attention d_params alloc: {e}")))?; let grad_norm_buf = stream.alloc_zeros::(1) @@ -188,40 +346,82 @@ impl GpuAttention { let norm_blocks = (total_params + 255) / 256; let grad_norm_block_sums = stream.alloc_zeros::(norm_blocks) .map_err(|e| MLError::ModelError(format!("attention grad_norm_block_sums alloc: {e}")))?; + + // ── Adam state ────────────────────────────────────────────────── + let attn_m = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("attention adam_m alloc: {e}")))?; let attn_v = stream.alloc_zeros::(total_params) .map_err(|e| MLError::ModelError(format!("attention adam_v alloc: {e}")))?; + let t_buf = stream.alloc_zeros::(1) + .map_err(|e| MLError::ModelError(format!("attn_t_buf alloc: {e}")))?; - // VRAM: params(4) + d_params + m + v + output + saved_input + saved_qkv + d_input_scratch + grad_norm - let vram_bytes = (total_params * 4 + b * d * 3 + b * 3 * d + 1) * 4; + // ── VRAM accounting ───────────────────────────────────────────── + + // Forward intermediates: proj_q + proj_k + proj_v + attn_out + output + saved_input + // + sdp_scores + ln_mean + ln_rstd + ln_xnorm + projected = 9*DB + num_heads*B + 2*B + // Backward intermediates: d_proj_q + d_proj_k + d_proj_v + d_attn_out + d_projected + // + d_input_scratch = 6*DB + // Params + optimizer: params + d_params + m + v = 4*total_params + // Total: 15*DB + num_heads*B + 2*B + 4*total_params + norm_blocks + 2 + let fwd_bufs = 9 * db + config.num_heads * b + 2 * b; + let bwd_bufs = 6 * db; + let opt_bufs = 4 * total_params + norm_blocks + 2; + let vram_bytes = (fwd_bufs + bwd_bufs + opt_bufs) * 4; let vram_kb = vram_bytes / 1024; + info!( state_dim = d, num_heads = config.num_heads, head_dim = d / config.num_heads, total_params, vram_kb, - "GpuAttention initialized: multi-head feature attention with backward pass" + gemm_descs = 12, + "GpuAttention initialized: cuBLAS GEMM projections, cubin SDP+LayerNorm" ); - let t_buf = stream.alloc_zeros::(1) - .map_err(|e| MLError::ModelError(format!("attn_t_buf alloc: {e}")))?; Ok(Self { config, stream, - forward_kernel, - backward_kernel, - weight_grad_reduce_kernel, + shared_handle, + gemm_fwd_q, + gemm_fwd_k, + gemm_fwd_v, + gemm_fwd_o, + gemm_bwd_dw_o, + gemm_bwd_dw_q, + gemm_bwd_dw_k, + gemm_bwd_dw_v, + gemm_bwd_dx_o, + gemm_bwd_dx_q, + gemm_bwd_dx_k, + gemm_bwd_dx_v, + sdp_fwd_kernel, + sdp_bwd_kernel, + layer_norm_fwd_kernel, + layer_norm_bwd_kernel, + bias_grad_reduce_kernel, grad_norm_phase1_kernel, grad_norm_phase2_kernel, adam_kernel, params, + proj_q_buf, + proj_k_buf, + proj_v_buf, + attn_out_buf, + sdp_scores_buf, output_buf, saved_input, - saved_qkv, + ln_save_mean, + ln_save_rstd, + ln_save_xnorm, + projected_buf, + d_proj_q_buf, + d_proj_k_buf, + d_proj_v_buf, + d_attn_out_buf, + d_projected_buf, d_input_scratch, - d_params_per_sample, d_params, grad_norm_buf, grad_norm_block_sums, @@ -235,14 +435,23 @@ impl GpuAttention { /// Apply multi-head attention to batch of states. /// - /// Saves input states and pre-layernorm output for the backward pass. + /// Pipeline: + /// 1. Q = W_Q^T @ states (cuBLAS) + bias_add + /// 2. K = W_K^T @ states (cuBLAS) + bias_add + /// 3. V = W_V^T @ states (cuBLAS) + bias_add + /// 4. attn_out = SDP(Q, K, V) (cubin kernel) + /// 5. projected = W_O^T @ attn_out (cuBLAS) + bias_add + /// 6. output = LayerNorm(projected + states) (cubin kernel) + /// /// Returns a reference to the output buffer (attended states). pub fn stream_handle(&self) -> u64 { self.stream.cu_stream() as u64 } + pub fn forward(&mut self, states: &CudaSlice, batch_size: usize) -> Result<&CudaSlice, MLError> { let b = batch_size; let d = self.config.state_dim; + let num_heads = self.config.num_heads; - // Save input states for backward pass (DtoD copy, raw ptrs) + // Save input states for backward pass (DtoD copy) let n_bytes = b * d * std::mem::size_of::(); unsafe { cudarc::driver::result::memcpy_dtod_async( @@ -250,50 +459,119 @@ impl GpuAttention { ).map_err(|e| MLError::ModelError(format!("attention save input DtoD: {e}")))?; } - // Shared memory: shared_concat[D] + warp_reduce[8] in f32 - let fwd_shmem = (d + 8) * std::mem::size_of::(); - let launch_cfg = LaunchConfig { - grid_dim: (b as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: fwd_shmem as u32, - }; - let b_i32 = b as i32; - let state_dim_i32 = d as i32; - let num_heads_i32 = self.config.num_heads as i32; + // Weight layout offsets in params_buf + let w_q_off = 0; + let b_q_off = d * d; + let w_k_off = d * d + d; + let b_k_off = 2 * d * d + d; + let w_v_off = 2 * d * d + 2 * d; + let b_v_off = 3 * d * d + 2 * d; + let w_o_off = 3 * d * d + 3 * d; + let b_o_off = 4 * d * d + 3 * d; + let f32_sz = std::mem::size_of::(); + let params_base = self.params.raw_ptr(); + // 1. Q = W_Q^T @ states + self.lt_matmul(&self.gemm_fwd_q, + params_base + (w_q_off * f32_sz) as u64, + states.raw_ptr(), + self.proj_q_buf.raw_ptr(), + 1.0, 0.0, "fwd_q")?; + // + bias_add Q + self.launch_bias_add(self.proj_q_buf.raw_ptr(), params_base + (b_q_off * f32_sz) as u64, d, b)?; + + // 2. K = W_K^T @ states + self.lt_matmul(&self.gemm_fwd_k, + params_base + (w_k_off * f32_sz) as u64, + states.raw_ptr(), + self.proj_k_buf.raw_ptr(), + 1.0, 0.0, "fwd_k")?; + // + bias_add K + self.launch_bias_add(self.proj_k_buf.raw_ptr(), params_base + (b_k_off * f32_sz) as u64, d, b)?; + + // 3. V = W_V^T @ states + self.lt_matmul(&self.gemm_fwd_v, + params_base + (w_v_off * f32_sz) as u64, + states.raw_ptr(), + self.proj_v_buf.raw_ptr(), + 1.0, 0.0, "fwd_v")?; + // + bias_add V + self.launch_bias_add(self.proj_v_buf.raw_ptr(), params_base + (b_v_off * f32_sz) as u64, d, b)?; + + // 4. SDP: attn_out = SDP_fwd(Q, K, V) + let b_i32 = b as i32; + let d_i32 = d as i32; + let num_heads_i32 = num_heads as i32; unsafe { self.stream - .launch_builder(&self.forward_kernel) - .arg(states) - .arg(&self.params) - .arg(&mut self.output_buf) - .arg(&mut self.saved_qkv) + .launch_builder(&self.sdp_fwd_kernel) + .arg(&self.proj_q_buf) + .arg(&self.proj_k_buf) + .arg(&self.proj_v_buf) + .arg(&mut self.attn_out_buf) + .arg(&mut self.sdp_scores_buf) .arg(&b_i32) - .arg(&state_dim_i32) + .arg(&d_i32) .arg(&num_heads_i32) - .launch(launch_cfg) - .map_err(|e| MLError::ModelError(format!("attention forward: {e}")))?; + .launch(LaunchConfig { + grid_dim: (b as u32, num_heads as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("attn_sdp_fwd: {e}")))?; } - // Forward kernel saves Q, K, V projections to saved_qkv [B, 3, D]. - // The backward kernel reads them directly, eliminating recomputation. + // 5. projected = W_O^T @ attn_out + self.lt_matmul(&self.gemm_fwd_o, + params_base + (w_o_off * f32_sz) as u64, + self.attn_out_buf.raw_ptr(), + self.projected_buf.raw_ptr(), + 1.0, 0.0, "fwd_o")?; + // + bias_add O + self.launch_bias_add(self.projected_buf.raw_ptr(), params_base + (b_o_off * f32_sz) as u64, d, b)?; + + // 6. LayerNorm(projected + residual=states) -> output + let gamma_off = 4 * d * d + 4 * d; + let beta_off = gamma_off + d; + unsafe { + self.stream + .launch_builder(&self.layer_norm_fwd_kernel) + .arg(&self.projected_buf) + .arg(states) + .arg(&(params_base + (gamma_off * f32_sz) as u64)) + .arg(&(params_base + (beta_off * f32_sz) as u64)) + .arg(&mut self.output_buf) + .arg(&mut self.ln_save_mean) + .arg(&mut self.ln_save_rstd) + .arg(&mut self.ln_save_xnorm) + .arg(&d_i32) + .arg(&b_i32) + .launch(LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("attn_layer_norm_fwd: {e}")))?; + } Ok(&self.output_buf) } /// Full backward pass: computes gradients for all attention parameters. /// - /// The backward kernel reads saved Q, K, V projections from the forward pass - /// to recover attention scores via softmax, eliminating expensive weight - /// matrix recomputation. Only softmax is recomputed (cheap element-wise ops). + /// Pipeline (reverse of forward): + /// 1. LayerNorm backward -> d_projected, d_gamma, d_beta + /// 2. dW_O = d_projected @ attn_out^T (cuBLAS, alpha=1/B) + /// 3. bias_grad_reduce for b_O + /// 4. d_attn_out = W_O @ d_projected (cuBLAS) + /// 5. SDP backward -> d_Q, d_K, d_V (cubin) + /// 6. dW_Q = d_Q @ states^T, dW_K = d_K @ states^T, dW_V = d_V @ states^T + /// (cuBLAS, alpha=1/B) + /// 7. bias_grad_reduce for b_Q, b_K, b_V + /// 8. d_input = W_Q @ d_Q + W_K @ d_K + W_V @ d_V (cuBLAS, accumulated) /// - /// `d_output` is the gradient w.r.t. the attended states [B, D] from the - /// DQN trunk backward pass (bw_d_h_s2). - /// - /// After this call, `d_params` contains accumulated weight gradients. - /// Call `adam_step()` after this to update the attention weights. - /// Input gradients are written to an internal scratch buffer (not propagated - /// further since attention is the first layer). + /// Zero atomicAdd. Weight gradients via GEMM (dW = dY @ X^T / B). + /// Bias gradients via deterministic sequential reduction kernel. pub fn backward( &mut self, d_output: &CudaSlice, @@ -301,77 +579,151 @@ impl GpuAttention { ) -> Result<(), MLError> { let b = batch_size; let d = self.config.state_dim; - let f32_size = std::mem::size_of::(); - let tile = b.min(256); + let num_heads = self.config.num_heads; + let f32_sz = std::mem::size_of::(); + let inv_batch = 1.0_f32 / b as f32; - // Zero d_input_scratch for the global output. - unsafe { - cudarc::driver::result::memset_d8_async( - self.d_input_scratch.raw_ptr(), 0u8, b * d * f32_size, self.stream.cu_stream() - ).map_err(|e| MLError::ModelError(format!("attention d_input zero: {e}")))?; - } - - // Zero d_params before tiled accumulation + // Zero grad_buf before accumulation self.stream.memset_zeros(&mut self.d_params) .map_err(|e| MLError::ModelError(format!("attention zero d_params: {e}")))?; - let shared_mem = ((6 * d + 8) * std::mem::size_of::()) as u32; - let state_dim_i32 = d as i32; - let num_heads_i32 = self.config.num_heads as i32; - let total_params_i32 = self.total_params as i32; - let reduce_blocks = (self.total_params + 255) / 256; - let d_params_tile_ptr = self.d_params_per_sample.raw_ptr(); - let d_params_ptr = self.d_params.raw_ptr(); + // Weight layout offsets + let w_q_off = 0; + let b_q_off = d * d; + let w_k_off = d * d + d; + let b_k_off = 2 * d * d + d; + let w_v_off = 2 * d * d + 2 * d; + let b_v_off = 3 * d * d + 2 * d; + let w_o_off = 3 * d * d + 3 * d; + let b_o_off = 4 * d * d + 3 * d; + let gamma_off = 4 * d * d + 4 * d; + let beta_off = gamma_off + d; + let params_base = self.params.raw_ptr(); + let grad_base = self.d_params.raw_ptr(); - // Tiled backward + reduce: process `tile` samples at a time - for tile_start in (0..b).step_by(tile) { - let tile_b = (b - tile_start).min(tile); - let tile_b_i32 = tile_b as i32; + let b_i32 = b as i32; + let d_i32 = d as i32; + let num_heads_i32 = num_heads as i32; - // Offset pointers into full-B buffers - let d_out_off = d_output.raw_ptr() + (tile_start * d * f32_size) as u64; - let saved_in_off = self.saved_input.raw_ptr() + (tile_start * d * f32_size) as u64; - let saved_qkv_off = self.saved_qkv.raw_ptr() + (tile_start * 3 * d * f32_size) as u64; - let d_input_off = self.d_input_scratch.raw_ptr() + (tile_start * d * f32_size) as u64; - - unsafe { - self.stream - .launch_builder(&self.backward_kernel) - .arg(&d_out_off) - .arg(&saved_in_off) - .arg(&self.params) - .arg(&saved_qkv_off) - .arg(&d_input_off) - .arg(&d_params_tile_ptr) - .arg(&tile_b_i32) - .arg(&state_dim_i32) - .arg(&num_heads_i32) - .arg(&total_params_i32) - .launch(LaunchConfig { - grid_dim: (tile_b as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: shared_mem, - }) - .map_err(|e| MLError::ModelError(format!("attention backward tile {tile_start}: {e}")))?; - } - - // Reduce tile into d_params (accumulates via +=) - unsafe { - self.stream - .launch_builder(&self.weight_grad_reduce_kernel) - .arg(&d_params_tile_ptr) - .arg(&d_params_ptr) - .arg(&tile_b_i32) - .arg(&total_params_i32) - .launch(LaunchConfig { - grid_dim: (reduce_blocks as u32, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: 0, - }) - .map_err(|e| MLError::ModelError(format!("attention reduce tile {tile_start}: {e}")))?; - } + // 1. LayerNorm backward -> d_projected (also writes d_gamma, d_beta into grad_buf) + unsafe { + self.stream + .launch_builder(&self.layer_norm_bwd_kernel) + .arg(d_output) + .arg(&self.ln_save_xnorm) + .arg(&(params_base + (gamma_off * f32_sz) as u64)) + .arg(&self.ln_save_rstd) + .arg(&mut self.d_projected_buf) + .arg(&(grad_base + (gamma_off * f32_sz) as u64)) + .arg(&(grad_base + (beta_off * f32_sz) as u64)) + .arg(&d_i32) + .arg(&b_i32) + .launch(LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("attn_layer_norm_bwd: {e}")))?; } + // 2. dW_O = d_projected[D,B] @ attn_out^T[B,D] -> [D,D], alpha=1/B for mean + self.lt_matmul(&self.gemm_bwd_dw_o, + self.d_projected_buf.raw_ptr(), + self.attn_out_buf.raw_ptr(), + grad_base + (w_o_off * f32_sz) as u64, + inv_batch, 0.0, "bwd_dw_o")?; + + // 3. Bias gradient reduce for b_O + self.launch_bias_grad_reduce( + self.d_projected_buf.raw_ptr(), + grad_base + (b_o_off * f32_sz) as u64, + d, b, inv_batch)?; + + // 4. d_attn_out = W_O[D,D] @ d_projected[D,B] -> [D,B] + self.lt_matmul(&self.gemm_bwd_dx_o, + params_base + (w_o_off * f32_sz) as u64, + self.d_projected_buf.raw_ptr(), + self.d_attn_out_buf.raw_ptr(), + 1.0, 0.0, "bwd_dx_o")?; + + // 5. SDP backward: d_attn_out -> d_Q, d_K, d_V + unsafe { + self.stream + .launch_builder(&self.sdp_bwd_kernel) + .arg(&self.d_attn_out_buf) + .arg(&self.proj_q_buf) + .arg(&self.proj_k_buf) + .arg(&self.proj_v_buf) + .arg(&self.sdp_scores_buf) + .arg(&mut self.d_proj_q_buf) + .arg(&mut self.d_proj_k_buf) + .arg(&mut self.d_proj_v_buf) + .arg(&b_i32) + .arg(&d_i32) + .arg(&num_heads_i32) + .launch(LaunchConfig { + grid_dim: (b as u32, num_heads as u32, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("attn_sdp_bwd: {e}")))?; + } + + // 6. dW_Q = d_Q[D,B] @ states^T[B,D], alpha=1/B + self.lt_matmul(&self.gemm_bwd_dw_q, + self.d_proj_q_buf.raw_ptr(), + self.saved_input.raw_ptr(), + grad_base + (w_q_off * f32_sz) as u64, + inv_batch, 0.0, "bwd_dw_q")?; + + // dW_K = d_K[D,B] @ states^T[B,D], alpha=1/B + self.lt_matmul(&self.gemm_bwd_dw_k, + self.d_proj_k_buf.raw_ptr(), + self.saved_input.raw_ptr(), + grad_base + (w_k_off * f32_sz) as u64, + inv_batch, 0.0, "bwd_dw_k")?; + + // dW_V = d_V[D,B] @ states^T[B,D], alpha=1/B + self.lt_matmul(&self.gemm_bwd_dw_v, + self.d_proj_v_buf.raw_ptr(), + self.saved_input.raw_ptr(), + grad_base + (w_v_off * f32_sz) as u64, + inv_batch, 0.0, "bwd_dw_v")?; + + // 7. Bias gradient reduce for b_Q, b_K, b_V + self.launch_bias_grad_reduce( + self.d_proj_q_buf.raw_ptr(), + grad_base + (b_q_off * f32_sz) as u64, + d, b, inv_batch)?; + self.launch_bias_grad_reduce( + self.d_proj_k_buf.raw_ptr(), + grad_base + (b_k_off * f32_sz) as u64, + d, b, inv_batch)?; + self.launch_bias_grad_reduce( + self.d_proj_v_buf.raw_ptr(), + grad_base + (b_v_off * f32_sz) as u64, + d, b, inv_batch)?; + + // 8. d_input = W_Q @ d_Q + W_K @ d_K + W_V @ d_V + // First: d_input = W_Q @ d_Q (beta=0, overwrites) + self.lt_matmul(&self.gemm_bwd_dx_q, + params_base + (w_q_off * f32_sz) as u64, + self.d_proj_q_buf.raw_ptr(), + self.d_input_scratch.raw_ptr(), + 1.0, 0.0, "bwd_dx_q")?; + // Accumulate: d_input += W_K @ d_K (beta=1) + self.lt_matmul(&self.gemm_bwd_dx_k, + params_base + (w_k_off * f32_sz) as u64, + self.d_proj_k_buf.raw_ptr(), + self.d_input_scratch.raw_ptr(), + 1.0, 1.0, "bwd_dx_k")?; + // Accumulate: d_input += W_V @ d_V (beta=1) + self.lt_matmul(&self.gemm_bwd_dx_v, + params_base + (w_v_off * f32_sz) as u64, + self.d_proj_v_buf.raw_ptr(), + self.d_input_scratch.raw_ptr(), + 1.0, 1.0, "bwd_dx_v")?; + Ok(()) } @@ -406,7 +758,7 @@ impl GpuAttention { .launch(norm_cfg) .map_err(|e| MLError::ModelError(format!("attention grad_norm phase1: {e}")))?; } - // Phase 2: deterministic reduction → L2 norm + // Phase 2: deterministic reduction -> L2 norm let norm_blocks_i32 = norm_blocks as i32; unsafe { self.stream @@ -418,7 +770,7 @@ impl GpuAttention { .map_err(|e| MLError::ModelError(format!("attention grad_norm phase2: {e}")))?; } - // 3. Adam update kernel — adam_step on GPU (async HtoD) + // 3. Adam update kernel -- adam_step on GPU (async HtoD) self.attn_adam_step += 1; unsafe { cudarc::driver::sys::cuMemcpyHtoDAsync_v2( @@ -473,4 +825,200 @@ impl GpuAttention { pub fn output(&self) -> &CudaSlice { &self.output_buf } + + // ── Private helpers ───────────────────────────────────────────────────── + + /// Execute cublasLtMatmul with pre-cached GEMM descriptor. + fn lt_matmul( + &self, + desc: &AttnGemmDesc, + a_ptr: u64, + b_ptr: u64, + c_ptr: u64, + alpha: f32, + beta: f32, + label: &str, + ) -> Result<(), MLError> { + unsafe { + let cu_stream = self.stream.cu_stream() as cublaslt_sys::cudaStream_t; + let matmul_status = cublaslt_sys::cublasLtMatmul( + self.shared_handle.lt_handle.0, + desc.matmul_desc, + &alpha as *const f32 as *const std::ffi::c_void, + a_ptr as *const std::ffi::c_void, + desc.a_layout, + b_ptr as *const std::ffi::c_void, + desc.b_layout, + &beta as *const f32 as *const std::ffi::c_void, + c_ptr as *const std::ffi::c_void, + desc.c_layout, + c_ptr as *mut std::ffi::c_void, + desc.d_layout, + &desc.algo as *const cublaslt_sys::cublasLtMatmulAlgo_t, + self.shared_handle.lt_workspace_ptr as *mut std::ffi::c_void, + self.shared_handle.lt_workspace_size, + cu_stream, + ); + if matmul_status != cublaslt_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + let e = cublaslt_result::CublasError(matmul_status); + return Err(MLError::ModelError(format!( + "cublasLtMatmul attention {label}: {e:?}" + ))); + } + } + Ok(()) + } + + /// Launch bias_add kernel: x[idx] += bias[idx % out_dim] + fn launch_bias_add(&self, out_ptr: u64, bias_ptr: u64, out_dim: usize, batch: usize) -> Result<(), MLError> { + let total = (batch * out_dim) as i32; + let out_dim_i32 = out_dim as i32; + let blocks = ((batch * out_dim + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.shared_handle.add_bias_f32_kernel) + .arg(&out_ptr) + .arg(&bias_ptr) + .arg(&out_dim_i32) + .arg(&total) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("attention bias_add: {e}")))?; + } + Ok(()) + } + + /// Launch bias gradient reduction kernel: d_bias[j] = sum(d_pre[j + b*out_dim]) * inv_batch + fn launch_bias_grad_reduce(&self, d_pre_ptr: u64, d_bias_ptr: u64, out_dim: usize, batch: usize, inv_batch: f32) -> Result<(), MLError> { + let out_dim_i32 = out_dim as i32; + let batch_i32 = batch as i32; + let blocks = ((out_dim + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.bias_grad_reduce_kernel) + .arg(&d_pre_ptr) + .arg(&d_bias_ptr) + .arg(&out_dim_i32) + .arg(&batch_i32) + .arg(&inv_batch) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("attention bias_grad_reduce: {e}")))?; + } + Ok(()) + } +} + +// ── GEMM descriptor factory ───────────────────────────────────────────────── + +/// Create an `AttnGemmDesc` for a given shape and transpose flags. +/// +/// Uses CUBLAS_COMPUTE_32F_FAST_TF32 for tensor-core acceleration. +/// Runs the heuristic ONCE at init time for optimal algorithm selection. +#[allow(clippy::too_many_arguments)] +fn create_attn_gemm_desc( + lt_handle: cublaslt_sys::cublasLtHandle_t, + transa_i32: i32, + transb_i32: i32, + m: i32, n: i32, k: i32, + lda: i32, ldb: i32, ldc: i32, + ws_size: usize, + label: &str, +) -> Result { + let f32_type = cublaslt_sys::cudaDataType_t::CUDA_R_32F; + let compute_type = cublaslt_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F_FAST_TF32; + + unsafe { + let matmul_desc = cublaslt_result::create_matmul_desc(compute_type, f32_type) + .map_err(|e| MLError::ModelError(format!("attn MatmulDescCreate ({label} m={m},n={n},k={k}): {e:?}")))?; + + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSA, + &transa_i32 as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("attn set TRANSA ({label}): {e:?}")))?; + + cublaslt_result::set_matmul_desc_attribute( + matmul_desc, + cublaslt_sys::cublasLtMatmulDescAttributes_t::CUBLASLT_MATMUL_DESC_TRANSB, + &transb_i32 as *const i32 as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| MLError::ModelError(format!("attn set TRANSB ({label}): {e:?}")))?; + + // Matrix layouts -- physical storage dimensions (before transpose) + let transa_is_n = transa_i32 == 0; + let transb_is_n = transb_i32 == 0; + + let a_layout = if transa_is_n { + cublaslt_result::create_matrix_layout(f32_type, m as u64, k as u64, lda as i64) + } else { + cublaslt_result::create_matrix_layout(f32_type, k as u64, m as u64, lda as i64) + }.map_err(|e| MLError::ModelError(format!("attn A layout ({label}): {e:?}")))?; + + let b_layout = if transb_is_n { + cublaslt_result::create_matrix_layout(f32_type, k as u64, n as u64, ldb as i64) + } else { + cublaslt_result::create_matrix_layout(f32_type, n as u64, k as u64, ldb as i64) + }.map_err(|e| MLError::ModelError(format!("attn B layout ({label}): {e:?}")))?; + + let c_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64) + .map_err(|e| MLError::ModelError(format!("attn C layout ({label}): {e:?}")))?; + + let d_layout = cublaslt_result::create_matrix_layout(f32_type, m as u64, n as u64, ldc as i64) + .map_err(|e| MLError::ModelError(format!("attn D layout ({label}): {e:?}")))?; + + // Algorithm selection via heuristic (graph-safe) + let matmul_pref = cublaslt_result::create_matmul_pref() + .map_err(|e| MLError::ModelError(format!("attn MatmulPrefCreate ({label}): {e:?}")))?; + cublaslt_result::set_matmul_pref_attribute( + matmul_pref, + cublaslt_sys::cublasLtMatmulPreferenceAttributes_t::CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &ws_size as *const usize as *const std::ffi::c_void, + std::mem::size_of::(), + ).map_err(|e| { + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + MLError::ModelError(format!("attn set pref ws ({label}): {e:?}")) + })?; + + let heuristic = cublaslt_result::get_matmul_algo_heuristic( + lt_handle, matmul_desc, + a_layout, b_layout, c_layout, d_layout, + matmul_pref, + ); + let _ = cublaslt_result::destroy_matmul_pref(matmul_pref); + + let heuristic = match heuristic { + Ok(h) => h, + Err(e) => { + let _ = cublaslt_result::destroy_matrix_layout(d_layout); + let _ = cublaslt_result::destroy_matrix_layout(c_layout); + let _ = cublaslt_result::destroy_matrix_layout(b_layout); + let _ = cublaslt_result::destroy_matrix_layout(a_layout); + let _ = cublaslt_result::destroy_matmul_desc(matmul_desc); + return Err(MLError::ModelError(format!( + "attn heuristic ({label} transa={transa_i32},transb={transb_i32},m={m},n={n},k={k}): {e:?}" + ))); + } + }; + + let algo = heuristic.algo; + tracing::info!( + label, + transa = transa_i32, transb = transb_i32, + m, n, k, lda, ldb, ldc, + ws_needed = heuristic.workspaceSize, + "attn GEMM desc created (heuristic algorithm)" + ); + + Ok(AttnGemmDesc { + matmul_desc, a_layout, b_layout, c_layout, d_layout, algo, + }) + } } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index f22dc8180..c2c9dc086 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -541,7 +541,7 @@ impl FusedTrainingCtx { num_heads: 4, batch_size, }; - let gpu_attention = match GpuAttention::new(stream.clone(), attn_config) { + let gpu_attention = match GpuAttention::new(Arc::clone(trainer.shared_cublas()), attn_config) { Ok(attn) => { info!( hidden_dim = shared_h2,