diff --git a/crates/ml/src/cuda_pipeline/batched_forward.rs b/crates/ml/src/cuda_pipeline/batched_forward.rs index feb914190..9dc02dee3 100644 --- a/crates/ml/src/cuda_pipeline/batched_forward.rs +++ b/crates/ml/src/cuda_pipeline/batched_forward.rs @@ -111,6 +111,9 @@ pub struct CublasForward { // ── Network dimensions (baked at construction) ── batch_size: usize, state_dim: usize, + /// pad128(state_dim): CUTLASS K-tile alignment for the B-matrix (input states). + /// Used as `ldb` for the first-layer GemmEx to prevent 128-element OOB reads. + state_dim_padded: usize, shared_h1: usize, shared_h2: usize, value_h: usize, @@ -182,6 +185,7 @@ impl CublasForward { // Dimensions batch_size, state_dim, + state_dim_padded: (state_dim + 127) & !127, shared_h1, shared_h2, value_h, @@ -237,7 +241,9 @@ impl CublasForward { ) -> Result<(), MLError> { let b = self.batch_size; - self.gemmex_bf16(w_ptrs[0], states_ptr, h_s1_ptr, self.shared_h1, b, self.state_dim, "h_s1")?; + // First layer: ldb = state_dim_padded (CUTLASS K-tile alignment). + // States buffer is padded to [B, pad128(state_dim)] with zero columns. + self.gemmex_bf16_ldb(w_ptrs[0], states_ptr, h_s1_ptr, self.shared_h1, b, self.state_dim, self.state_dim_padded, "h_s1")?; self.launch_add_bias_relu_bf16_raw(stream, h_s1_ptr, w_ptrs[1], self.shared_h1, b)?; self.gemmex_bf16(w_ptrs[2], h_s1_ptr, h_s2_ptr, self.shared_h2, b, self.shared_h1, "h_s2")?; @@ -367,9 +373,9 @@ impl CublasForward { ) -> Result<(), MLError> { let b = self.batch_size; - // h_s1[B, SH1] - self.gemmex_bf16(tg_w_ptrs[0], states_ptr, h_s1_ptr, - self.shared_h1, b, self.state_dim, "tg_h_s1")?; + // h_s1[B, SH1] — first layer: ldb = state_dim_padded (CUTLASS K-tile alignment) + self.gemmex_bf16_ldb(tg_w_ptrs[0], states_ptr, h_s1_ptr, + self.shared_h1, b, self.state_dim, self.state_dim_padded, "tg_h_s1")?; self.launch_add_bias_relu_bf16_raw(stream, h_s1_ptr, tg_w_ptrs[1], self.shared_h1, b)?; // h_s2[B, SH2] @@ -503,6 +509,58 @@ impl CublasForward { Ok(()) } + /// BF16 x BF16 -> BF16 GEMM with an explicit `ldb` override. + /// + /// Identical to `gemmex_bf16` except the B-matrix (input) leading dimension + /// is `ldb` instead of `k`. Used for the first layer where the states + /// buffer is padded to `pad128(state_dim)` for CUTLASS K-tile alignment. + /// + /// `lda` is still `k` (weight matrix is contiguous with stride = in_dim). + #[allow(clippy::too_many_arguments)] + fn gemmex_bf16_ldb( + &self, + w_bf16_ptr: u64, // BF16 weights [out_dim, in_dim] + a_bf16_ptr: u64, // BF16 input [B, in_dim] (row stride = ldb) + c_bf16_ptr: u64, // BF16 output [B, out_dim] + n: usize, // out_dim + b: usize, // batch + k: usize, // in_dim (logical K dimension of the GEMM) + ldb: usize, // leading dim of B-matrix (>= k, padded stride) + _label: &str, + ) -> Result<(), MLError> { + let alpha = 1.0_f32; + let beta = 0.0_f32; + + unsafe { + let status = cublas_sys::cublasGemmEx( + self.handle.0, + cublas_sys::cublasOperation_t::CUBLAS_OP_T, // transa: transpose W + cublas_sys::cublasOperation_t::CUBLAS_OP_N, // transb: keep A as-is + n as i32, // m (rows of op(W) = N = out_dim) + b as i32, // n (cols of op(A) = B = batch) + k as i32, // k (inner dim) + &alpha as *const f32 as *const std::ffi::c_void, + w_bf16_ptr as *const std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_16BF, + k as i32, // lda: leading dim of W (before transpose) = K + a_bf16_ptr as *const std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_16BF, + ldb as i32, // ldb: padded leading dim of A + &beta as *const f32 as *const std::ffi::c_void, + c_bf16_ptr as *mut std::ffi::c_void, + cublas_sys::cudaDataType_t::CUDA_R_16BF, + n as i32, // ldc: leading dim of C = N + cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F, + cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DEFAULT_TENSOR_OP, + ); + if status != cublas_sys::cublasStatus_t::CUBLAS_STATUS_SUCCESS { + return Err(MLError::ModelError(format!("cublasGemmEx ldb {_label}: {status:?}"))); + } + } + + Ok(()) + } + // ══════════════════════════════════════════════════════════════════════════ // cublasGemmEx BF16-input → F32-output GEMM (output layer path) // ══════════════════════════════════════════════════════════════════════════ diff --git a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu index 8940b2b78..d3591a6d4 100644 --- a/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu +++ b/crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu @@ -417,3 +417,30 @@ extern "C" __global__ void spectral_norm_kernel( W[i] = W[i] * scale; } } + +/* ══════════════════════════════════════════════════════════════════════ + * PAD STATES KERNEL (CUTLASS K-tile alignment) + * + * Scatters contiguous [batch, sd] BF16 states into a padded + * [batch, padded_sd] layout with zero-filled columns [sd..padded_sd). + * + * CUTLASS uses 128-element K-tiles in GemmEx. When state_dim (K) is + * not a multiple of 128, the kernel reads past row boundaries. Padding + * each row to pad128(state_dim) with zeros eliminates the OOB reads. + * + * Launch config: grid=(ceil(batch * padded_sd / 256), 1, 1), block=(256, 1, 1). + * ══════════════════════════════════════════════════════════════════════ */ + +extern "C" __global__ void pad_states_kernel( + __nv_bfloat16* __restrict__ dst, /* [batch, padded_sd] — padded output */ + const __nv_bfloat16* __restrict__ src, /* [batch, sd] — contiguous input */ + int batch, + int sd, + int padded_sd +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= batch * padded_sd) return; + int b = idx / padded_sd; + int j = idx % padded_sd; + dst[idx] = (j < sd) ? src[b * sd + j] : bf16_zero(); +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d4c30d8c4..6588eb03d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -399,6 +399,7 @@ pub struct GpuDqnTrainer { spectral_norm_kernel: CudaFunction, clipped_saxpy_kernel: CudaFunction, clip_grad_kernel: CudaFunction, + pad_states_kernel: CudaFunction, /// Spectral norm left singular vectors: [out_dim] per weight matrix (W_s1, W_s2) spec_u_s1: CudaSlice, spec_v_s1: CudaSlice, @@ -1166,7 +1167,8 @@ impl GpuDqnTrainer { let td_ptr = self.td_errors_buf.raw_ptr(); let states_ptr = self.states_buf.raw_ptr(); - let state_dim_i32 = self.config.state_dim as i32; + // Use padded state_dim as stride — states_buf rows are pad128(state_dim) wide. + let state_dim_i32 = pad128(self.config.state_dim) as i32; unsafe { self.stream .launch_builder(&self.regime_scale_kernel) @@ -1185,6 +1187,40 @@ impl GpuDqnTrainer { Ok(()) } + /// Launch `pad_states_kernel`: scatter contiguous `[batch, sd]` BF16 states + /// into a padded `[batch, padded_sd]` destination with zero-filled columns. + /// + /// This eliminates CUTLASS K-tile OOB reads in the first-layer GemmEx. + fn launch_pad_states( + &self, + dst: u64, // [batch, padded_sd] — must be pre-allocated + src: u64, // [batch, sd] — contiguous source + batch: usize, + ) -> Result<(), MLError> { + let sd = self.config.state_dim as i32; + let padded_sd = pad128(self.config.state_dim) as i32; + let total = batch as i32 * padded_sd; + let blocks = ((total as u32) + 255) / 256; + + unsafe { + self.stream + .launch_builder(&self.pad_states_kernel) + .arg(&dst) + .arg(&src) + .arg(&(batch as i32)) + .arg(&sd) + .arg(&padded_sd) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("pad_states_kernel: {e}")))?; + } + + Ok(()) + } + /// Apply CQL (Conservative Q-Learning) gradient injection. /// /// Computes the CQL penalty gradient w.r.t. the C51 logits and runs a @@ -1723,7 +1759,7 @@ impl GpuDqnTrainer { // per array. Stack is set once in DQNTrainer::new() (64KB for all kernels). // ── Compile 4 utility kernels (grad_norm, adam_update, BF16 converters) ─ - let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel) = + let (grad_norm_kernel, grad_norm_finalize_kernel, adam_update_kernel, f32_to_bf16_kernel, bf16_to_f32_kernel, saxpy_kernel, zero_kernel, regime_scale_kernel, shrink_perturb, _relu_mask_in_module, spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel, pad_states_kernel) = compile_training_kernels(&stream, &config)?; // Separate grad_norm instance for non-graph launches (clip_grad_buf_inplace). @@ -1744,8 +1780,12 @@ impl GpuDqnTrainer { let per_update_kernel = compile_per_update_kernel(&stream)?; // ── Allocate batch input buffers ──────────────────────────── - let states_buf = alloc_bf16(&stream, b * config.state_dim, "states")?; - let next_states_buf = alloc_bf16(&stream, b * config.state_dim, "next_states")?; + // States buffers are padded to pad128(state_dim) per row so that + // CUTLASS 128-element K-tiles in the first-layer GemmEx never read + // past the row boundary. The padding columns are kept at zero. + let state_dim_padded = pad128(config.state_dim); + let states_buf = alloc_bf16(&stream, b * state_dim_padded, "states")?; + let next_states_buf = alloc_bf16(&stream, b * state_dim_padded, "next_states")?; let actions_buf = alloc_i32(&stream, b, "actions")?; let rewards_buf = alloc_bf16(&stream, b, "rewards")?; let dones_buf = alloc_bf16(&stream, b, "dones")?; @@ -2103,6 +2143,7 @@ impl GpuDqnTrainer { spectral_norm_kernel, clipped_saxpy_kernel, clip_grad_kernel, + pad_states_kernel, spec_u_s1, spec_v_s1, spec_u_s2, @@ -2747,22 +2788,21 @@ impl GpuDqnTrainer { ) -> Result { let _eg = EventTrackingGuard::new(self.stream.context()); - // Pad input to config.batch_size to prevent CUTLASS tile-boundary OOB reads. - // CUTLASS WMMA uses 32-element N-tiles; non-aligned batch sizes read past buffer. - let mut padded_buf; - let eff_states = if batch_size < self.config.batch_size { - let sd = self.config.state_dim; - padded_buf = self.stream.alloc_zeros::(self.config.batch_size * sd) - .map_err(|e| MLError::ModelError(format!("q_stats pad alloc: {e}")))?; - { - let mut dst = padded_buf.slice_mut(0..batch_size * sd); - self.stream.memcpy_dtod(states, &mut dst) - .map_err(|e| MLError::ModelError(format!("q_stats pad copy: {e}")))?; - } - &padded_buf - } else { - states - }; + // Pad input rows to pad128(state_dim) for CUTLASS K-tile alignment, + // and pad batch dimension to config.batch_size for N-tile alignment. + let sd = self.config.state_dim; + let padded_sd = pad128(sd); + let eff_batch = self.config.batch_size; + let mut padded_buf = self.stream.alloc_zeros::(eff_batch * padded_sd) + .map_err(|e| MLError::ModelError(format!("q_stats pad alloc: {e}")))?; + // Scatter contiguous [batch_size, sd] → padded [eff_batch, padded_sd] + // (only first `batch_size` rows are populated; rest stay zero) + self.launch_pad_states( + padded_buf.raw_ptr(), + states.raw_ptr(), + batch_size, + )?; + let eff_states = &padded_buf; self.compute_q_values(eff_states, batch_size)?; @@ -3169,15 +3209,23 @@ impl GpuDqnTrainer { let staging_base = self.upload_staging_buf.raw_ptr(); let mut byte_offset: u64 = 0; - // states: B * SD bf16 elements - let states_bytes = b * sd * bf16_size; - dtod_copy(self.states_buf.raw_ptr(), staging_base + byte_offset, states_bytes, &self.stream, 0, "upload_scatter")?; - byte_offset += states_bytes as u64; + // states: contiguous [B, SD] in staging → padded [B, pad128(SD)] in states_buf + let states_bytes = (b * sd * bf16_size) as u64; + self.launch_pad_states( + self.states_buf.raw_ptr(), + staging_base + byte_offset, + b, + )?; + byte_offset += states_bytes; - // next_states: B * SD bf16 elements - let next_states_bytes = b * sd * bf16_size; - dtod_copy(self.next_states_buf.raw_ptr(), staging_base + byte_offset, next_states_bytes, &self.stream, 1, "upload_scatter")?; - byte_offset += next_states_bytes as u64; + // next_states: contiguous [B, SD] → padded [B, pad128(SD)] + let next_states_bytes = (b * sd * bf16_size) as u64; + self.launch_pad_states( + self.next_states_buf.raw_ptr(), + staging_base + byte_offset, + b, + )?; + byte_offset += next_states_bytes; // rewards: B bf16 elements let rewards_bytes = b * bf16_size; @@ -3214,20 +3262,19 @@ impl GpuDqnTrainer { gpu_batch: &crate::dqn::replay_buffer_type::GpuBatch, ) -> Result<(), MLError> { let b = self.config.batch_size; - let sd = self.config.state_dim; let bf16_size = std::mem::size_of::(); - // States + next_states: bf16 DtoD - dtod_copy( + // States + next_states: contiguous [B, SD] in GpuBatch → padded [B, pad128(SD)] + self.launch_pad_states( self.states_buf.raw_ptr(), gpu_batch.states.data().raw_ptr(), - b * sd * bf16_size, &self.stream, 0, "states", + b, )?; - dtod_copy( + self.launch_pad_states( self.next_states_buf.raw_ptr(), gpu_batch.next_states.data().raw_ptr(), - b * sd * bf16_size, &self.stream, 1, "next_states", + b, )?; // Rewards, dones: bf16 DtoD @@ -4054,7 +4101,7 @@ impl GpuDqnTrainer { fn compile_training_kernels( stream: &Arc, config: &GpuDqnTrainConfig, -) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { +) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> { info!( state_dim = config.state_dim, total_params = compute_total_params(config), @@ -4092,9 +4139,11 @@ fn compile_training_kernels( .map_err(|e| MLError::ModelError(format!("dqn_clipped_saxpy_kernel load: {e}")))?; let clip_grad = module.load_function("dqn_clip_grad_kernel") .map_err(|e| MLError::ModelError(format!("dqn_clip_grad_kernel load: {e}")))?; + let pad_states = module.load_function("pad_states_kernel") + .map_err(|e| MLError::ModelError(format!("pad_states_kernel load: {e}")))?; - info!("GpuDqnTrainer: 12 utility kernels loaded from precompiled cubin"); - Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad)) + info!("GpuDqnTrainer: 13 utility kernels loaded from precompiled cubin"); + Ok((grad_norm, grad_norm_finalize, adam_update, f32_to_bf16, bf16_to_f32, saxpy, zero, regime_scale, shrink_perturb, _relu_mask_from_module, spectral_norm, clipped_saxpy, clip_grad, pad_states)) } /// Load the standalone Polyak EMA kernel from precompiled cubin. @@ -4320,6 +4369,18 @@ fn pad32(n: usize) -> usize { (n + 31) & !31 } +/// Round up to the next multiple of 128 (CUTLASS K-tile alignment). +/// +/// cuBLAS CUTLASS kernels use 128-element K-tiles for BF16 GemmEx. +/// When the K dimension (state_dim) is not a multiple of 128, CUTLASS +/// reads past the row boundary of the B-matrix. Padding each row of +/// the states buffer to `pad128(state_dim)` with zeros eliminates OOB +/// reads without changing the GEMM logical K dimension. +#[inline(always)] +fn pad128(n: usize) -> usize { + (n + 127) & !127 +} + fn alloc_bf16( stream: &Arc, n: usize,