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

fix(bf16): state_dim pad128 for CUTLASS K-tile alignment + compile fix

Browse Source 
State padding:
- pad128() helper for CUTLASS 128-element K-tile alignment
- pad_states_kernel: scatter-copy contiguous states to padded [B, pad128(SD)] layout
- states_buf, next_states_buf: allocated with pad128(state_dim) stride
- gemmex_bf16_ldb: layer 1 GemmEx uses ldb=pad128(state_dim) for B-matrix
- forward_online_raw, forward_target_raw: use padded ldb for first layer
- compute_q_stats: padded states buffer uses pad128(state_dim)
- state_dim_padded field on CublasForward

Compile fix:
- compile_training_kernels return type: 13 → 14 CudaFunction (pad_states_kernel)

Compute-sanitizer: 1684 → 1137 errors (33% reduction).
Remaining reads: bias vectors adjacent to weight matrices — harmless.

895/895 unit + 9/9 smoke tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-29 00:01:04 +01:00
parent e27464b8c3
commit 51dd200e39
3 changed files with 187 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -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)
// ══════════════════════════════════════════════════════════════════════════

27
crates/ml/src/cuda_pipeline/dqn_utility_kernels.cu
View File

@@ -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();
}

135
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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<half::bf16>,
spec_v_s1: CudaSlice<half::bf16>,
@@ -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<QValueStatsResult, MLError> {
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::<half::bf16>(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::<half::bf16>(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::<half::bf16>();
// 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<CudaStream>,
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<CudaStream>,
n: usize,