feat(bf16): BF16 cuBLAS forward infrastructure — gemmex_bf16 + BF16 bias kernels
- cublasGemmEx BF16×BF16→BF16 method (CUBLAS_COMPUTE_32F, tensor core path) - BF16 bias+relu and bias-only CUDA kernels (add_bias_relu_bf16_kernel) - 15 BF16 activation buffers allocated in CublasForward (online + target) - f32_to_bf16_kernel loaded for states conversion - bf16_weight_ptrs() helper for BF16 flat buffer offset computation - forward_online_bf16() method — complete BF16 forward pass (not yet wired) - cudarc f16 feature enabled, nvrtc removed Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -66,6 +66,8 @@ fn main() {
|
||||
"monitoring_kernel.cu",
|
||||
"nstep_kernel.cu",
|
||||
"ppo_experience_kernel.cu",
|
||||
// bias_kernels needs common header for f32_to_bf16_kernel (used by CublasForward)
|
||||
"bias_kernels.cu",
|
||||
];
|
||||
|
||||
// Standalone kernels: no common header needed.
|
||||
@@ -85,7 +87,6 @@ fn main() {
|
||||
"cql_grad_kernel.cu",
|
||||
// Inline kernels extracted from other modules
|
||||
"trade_stats_kernel.cu",
|
||||
"bias_kernels.cu",
|
||||
"backward_kernels.cu",
|
||||
"iqn_cvar_kernel.cu",
|
||||
];
|
||||
|
||||
@@ -7,6 +7,16 @@
|
||||
//! occupancy from ~1.56% to >60% on H100 by eliminating the per-sample block
|
||||
//! bottleneck caused by 134 KB shared memory usage.
|
||||
//!
|
||||
//! ## BF16 tensor core path (cublasGemmEx)
|
||||
//!
|
||||
//! Alongside the existing F32 `cublasSgemm` path, this module provides a
|
||||
//! `gemmex_bf16` method that calls `cublasGemmEx` with BF16 inputs and F32
|
||||
//! internal accumulation. On H100 this yields ~3x throughput vs F32 SGEMM.
|
||||
//!
|
||||
//! The BF16 path uses internal activation buffers (`h_s1_bf16`, etc.) allocated
|
||||
//! at construction time, plus BF16 bias+relu kernels (`add_bias_relu_bf16_kernel`,
|
||||
//! `add_bias_bf16_kernel`). The existing F32 `sgemm_layer` stays as fallback.
|
||||
//!
|
||||
//! ## Layout convention
|
||||
//!
|
||||
//! All tensors use row-major (C-style) layout:
|
||||
@@ -39,9 +49,10 @@
|
||||
//!
|
||||
//! ## cuBLAS CUDA Graph compatibility
|
||||
//!
|
||||
//! `cublasSgemm` is capturable in CUDA 12+ stream-capture mode.
|
||||
//! The cuBLAS handle must be set to the capture stream via `cublasSetStream` before
|
||||
//! capture begins. Handle creation itself is NOT capturable and must happen before.
|
||||
//! Both `cublasSgemm` and `cublasGemmEx` are capturable in CUDA 12+ stream-capture
|
||||
//! mode. The cuBLAS handle must be set to the capture stream via `cublasSetStream`
|
||||
//! before capture begins. Handle creation itself is NOT capturable and must happen
|
||||
//! before.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::mem::ManuallyDrop;
|
||||
@@ -86,14 +97,54 @@ pub struct CublasForward {
|
||||
_workspace_buf: CudaSlice<u8>,
|
||||
|
||||
// ── Compiled bias+activation kernels (loaded from the DQN module) ──
|
||||
/// Kernel: `add_bias_relu(output, bias, out_dim, total_elements)`
|
||||
/// Kernel: `add_bias_relu(output, bias, out_dim, total_elements)` — F32
|
||||
/// Grid: ceil(B * out_dim / 256), Block: 256.
|
||||
add_bias_relu_kernel: CudaFunction,
|
||||
|
||||
/// Kernel: `add_bias(output, bias, out_dim, total_elements)`
|
||||
/// Kernel: `add_bias(output, bias, out_dim, total_elements)` — F32
|
||||
/// Same shape as add_bias_relu but no clamping (used for final logit layers).
|
||||
add_bias_kernel: CudaFunction,
|
||||
|
||||
/// Kernel: `add_bias_relu_bf16(output, bias, out_dim, total_elements)` — BF16
|
||||
/// BF16 variant of add_bias_relu for tensor core forward path.
|
||||
add_bias_relu_bf16_kernel: CudaFunction,
|
||||
|
||||
/// Kernel: `add_bias_bf16(output, bias, out_dim, total_elements)` — BF16
|
||||
/// BF16 variant of add_bias for tensor core forward path.
|
||||
add_bias_bf16_kernel: CudaFunction,
|
||||
|
||||
/// Kernel: `f32_to_bf16_kernel(src_f32, dst_bf16, n)` — element-wise F32→BF16.
|
||||
/// Used by `forward_online_bf16` to convert F32 states from the experience
|
||||
/// collector into BF16 before the tensor core GEMM pipeline.
|
||||
/// Grid: ceil(n/256), Block: 256.
|
||||
f32_to_bf16_kernel: CudaFunction,
|
||||
|
||||
// ── BF16 activation buffers for tensor core forward path ─────────
|
||||
// Internal BF16 buffers needed because forward_online takes &CudaSlice<f32>
|
||||
// and we cannot change that signature without cascading to callers.
|
||||
// The gemmex_bf16 path reads/writes these instead.
|
||||
|
||||
/// Online network BF16 activation buffers
|
||||
states_bf16: CudaSlice<u16>, // [B, SD]
|
||||
h_s1_bf16: CudaSlice<u16>, // [B, SH1]
|
||||
h_s2_bf16: CudaSlice<u16>, // [B, SH2]
|
||||
h_v_bf16: CudaSlice<u16>, // [B, VH]
|
||||
h_b0_bf16: CudaSlice<u16>, // [B, AH]
|
||||
h_b1_bf16: CudaSlice<u16>, // [B, AH]
|
||||
h_b2_bf16: CudaSlice<u16>, // [B, AH]
|
||||
/// Online logit BF16 buffers
|
||||
v_logits_bf16: CudaSlice<u16>, // [B, NA]
|
||||
b_logits_bf16: CudaSlice<u16>, // [B, (B0+B1+B2)*NA]
|
||||
|
||||
/// Target network BF16 activation buffers (inference scratch — reused)
|
||||
tgt_h_s1_bf16: CudaSlice<u16>, // [B, SH1]
|
||||
tgt_h_s2_bf16: CudaSlice<u16>, // [B, SH2]
|
||||
tgt_h_v_bf16: CudaSlice<u16>, // [B, VH]
|
||||
tgt_h_b_bf16: CudaSlice<u16>, // [B, AH] (shared across branches)
|
||||
/// Target logit BF16 buffers
|
||||
tgt_v_logits_bf16: CudaSlice<u16>, // [B, NA]
|
||||
tgt_b_logits_bf16: CudaSlice<u16>, // [B, (B0+B1+B2)*NA]
|
||||
|
||||
// ── Network dimensions (baked at construction) ──
|
||||
batch_size: usize,
|
||||
state_dim: usize,
|
||||
@@ -155,15 +206,75 @@ impl CublasForward {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Compile bias kernels ────────────────────────────────────────
|
||||
let (add_bias_relu_kernel, add_bias_kernel) =
|
||||
// ── Compile bias kernels (F32 + BF16) + f32_to_bf16 converter ──
|
||||
let (add_bias_relu_kernel, add_bias_kernel,
|
||||
add_bias_relu_bf16_kernel, add_bias_bf16_kernel,
|
||||
f32_to_bf16_kernel) =
|
||||
compile_bias_kernels(stream)?;
|
||||
|
||||
// ── Allocate BF16 activation buffers ─────────────────────────
|
||||
let total_branch_logits = (branch_0_size + branch_1_size + branch_2_size) * num_atoms;
|
||||
|
||||
// Online network BF16 activation buffers
|
||||
let states_bf16 = stream.alloc_zeros::<u16>(batch_size * state_dim)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 states_bf16 alloc: {e}")))?;
|
||||
let h_s1_bf16 = stream.alloc_zeros::<u16>(batch_size * shared_h1)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 h_s1_bf16 alloc: {e}")))?;
|
||||
let h_s2_bf16 = stream.alloc_zeros::<u16>(batch_size * shared_h2)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 h_s2_bf16 alloc: {e}")))?;
|
||||
let h_v_bf16 = stream.alloc_zeros::<u16>(batch_size * value_h)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 h_v_bf16 alloc: {e}")))?;
|
||||
let h_b0_bf16 = stream.alloc_zeros::<u16>(batch_size * adv_h)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 h_b0_bf16 alloc: {e}")))?;
|
||||
let h_b1_bf16 = stream.alloc_zeros::<u16>(batch_size * adv_h)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 h_b1_bf16 alloc: {e}")))?;
|
||||
let h_b2_bf16 = stream.alloc_zeros::<u16>(batch_size * adv_h)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 h_b2_bf16 alloc: {e}")))?;
|
||||
let v_logits_bf16 = stream.alloc_zeros::<u16>(batch_size * num_atoms)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 v_logits_bf16 alloc: {e}")))?;
|
||||
let b_logits_bf16 = stream.alloc_zeros::<u16>(batch_size * total_branch_logits)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 b_logits_bf16 alloc: {e}")))?;
|
||||
|
||||
// Target network BF16 activation buffers (scratch — reused)
|
||||
let tgt_h_s1_bf16 = stream.alloc_zeros::<u16>(batch_size * shared_h1)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 tgt_h_s1_bf16 alloc: {e}")))?;
|
||||
let tgt_h_s2_bf16 = stream.alloc_zeros::<u16>(batch_size * shared_h2)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 tgt_h_s2_bf16 alloc: {e}")))?;
|
||||
let tgt_h_v_bf16 = stream.alloc_zeros::<u16>(batch_size * value_h)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 tgt_h_v_bf16 alloc: {e}")))?;
|
||||
let tgt_h_b_bf16 = stream.alloc_zeros::<u16>(batch_size * adv_h)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 tgt_h_b_bf16 alloc: {e}")))?;
|
||||
let tgt_v_logits_bf16 = stream.alloc_zeros::<u16>(batch_size * num_atoms)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 tgt_v_logits_bf16 alloc: {e}")))?;
|
||||
let tgt_b_logits_bf16 = stream.alloc_zeros::<u16>(batch_size * total_branch_logits)
|
||||
.map_err(|e| MLError::ModelError(format!("BF16 tgt_b_logits_bf16 alloc: {e}")))?;
|
||||
|
||||
Ok(Self {
|
||||
handle: SendSyncCublasHandle(raw_handle),
|
||||
_workspace_buf: workspace_buf,
|
||||
add_bias_relu_kernel,
|
||||
add_bias_kernel,
|
||||
add_bias_relu_bf16_kernel,
|
||||
add_bias_bf16_kernel,
|
||||
f32_to_bf16_kernel,
|
||||
// Online BF16 activation buffers
|
||||
states_bf16,
|
||||
h_s1_bf16,
|
||||
h_s2_bf16,
|
||||
h_v_bf16,
|
||||
h_b0_bf16,
|
||||
h_b1_bf16,
|
||||
h_b2_bf16,
|
||||
v_logits_bf16,
|
||||
b_logits_bf16,
|
||||
// Target BF16 activation buffers
|
||||
tgt_h_s1_bf16,
|
||||
tgt_h_s2_bf16,
|
||||
tgt_h_v_bf16,
|
||||
tgt_h_b_bf16,
|
||||
tgt_v_logits_bf16,
|
||||
tgt_b_logits_bf16,
|
||||
// Dimensions
|
||||
batch_size,
|
||||
state_dim,
|
||||
shared_h1,
|
||||
@@ -321,6 +432,180 @@ impl CublasForward {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
// Forward pass (online network — BF16 tensor core path)
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// Run the online network forward pass using cublasGemmEx BF16 tensor cores.
|
||||
///
|
||||
/// Same layer sequence as `forward_online` but all GEMMs use BF16 weights,
|
||||
/// BF16 activations, and F32 internal accumulation (H100: ~3x throughput).
|
||||
///
|
||||
/// ## Data flow
|
||||
///
|
||||
/// 1. **States F32→BF16**: `f32_to_bf16_kernel` converts the F32 experience
|
||||
/// collector states into `self.states_bf16`.
|
||||
/// 2. **Hidden layers**: `gemmex_bf16(W_bf16, input_bf16, output_bf16)` then
|
||||
/// `add_bias_relu_bf16(output_bf16, bias_bf16)`.
|
||||
/// 3. **Output layers**: `gemmex_bf16(W_bf16, input_bf16, output_bf16)` then
|
||||
/// `add_bias_bf16(output_bf16, bias_bf16)`.
|
||||
///
|
||||
/// All activations are written to `self.{h_s1_bf16, h_s2_bf16, ...}` internal
|
||||
/// buffers. Callers access them via `h_s1_bf16_ptr()` etc. for the backward
|
||||
/// pass (which will also operate in BF16).
|
||||
///
|
||||
/// ## Weight pointer layout
|
||||
///
|
||||
/// `bf16_w_ptrs` has the same 20-entry layout as the F32 `w_ptrs`:
|
||||
/// `[w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2,
|
||||
/// w_b0fc, b_b0fc, w_b0out, b_b0out, ...]`
|
||||
/// — but all pointers target BF16 (`u16`) device memory in `bf16_params_buf`.
|
||||
#[allow(dead_code, clippy::too_many_arguments)]
|
||||
pub fn forward_online_bf16(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
// ── Inputs ──────────────────────────────────────────────────────
|
||||
states_f32: &CudaSlice<f32>, // [B, SD] F32 from experience collector
|
||||
// ── BF16 weight pointers (raw u64 into flat bf16_params_buf) ──
|
||||
bf16_w_ptrs: &[u64; 20], // [w_s1, b_s1, w_s2, b_s2, ...]
|
||||
) -> Result<(), MLError> {
|
||||
let b = self.batch_size;
|
||||
let n_states = b * self.state_dim;
|
||||
|
||||
// ── Step 1: Convert states F32 → BF16 ────────────────────────────
|
||||
{
|
||||
let src_ptr = raw_f32_ptr(states_f32, stream);
|
||||
let dst_ptr = raw_u16_ptr(&self.states_bf16, stream);
|
||||
let n_i32 = n_states as i32;
|
||||
let blocks = ((n_states + 255) / 256) as u32;
|
||||
|
||||
unsafe {
|
||||
stream
|
||||
.launch_builder(&self.f32_to_bf16_kernel)
|
||||
.arg(&src_ptr)
|
||||
.arg(&dst_ptr)
|
||||
.arg(&n_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!("f32_to_bf16 states: {e}")))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── Step 2: Shared trunk layer 1 ─────────────────────────────────
|
||||
// h_s1[B, SH1] = ReLU(states_bf16 @ W_s1_bf16^T + b_s1_bf16)
|
||||
let states_bf16_ptr = raw_u16_ptr(&self.states_bf16, stream);
|
||||
let h_s1_bf16_ptr = raw_u16_ptr(&self.h_s1_bf16, stream);
|
||||
|
||||
self.gemmex_bf16(
|
||||
bf16_w_ptrs[0], // W_s1[SH1, SD]
|
||||
states_bf16_ptr,
|
||||
h_s1_bf16_ptr,
|
||||
self.shared_h1, b, self.state_dim,
|
||||
"bf16_h_s1",
|
||||
)?;
|
||||
self.launch_add_bias_relu_bf16_raw(
|
||||
stream, h_s1_bf16_ptr, bf16_w_ptrs[1], self.shared_h1, b,
|
||||
)?;
|
||||
|
||||
// ── Step 3: Shared trunk layer 2 ─────────────────────────────────
|
||||
// h_s2[B, SH2] = ReLU(h_s1_bf16 @ W_s2_bf16^T + b_s2_bf16)
|
||||
let h_s2_bf16_ptr = raw_u16_ptr(&self.h_s2_bf16, stream);
|
||||
|
||||
self.gemmex_bf16(
|
||||
bf16_w_ptrs[2], // W_s2[SH2, SH1]
|
||||
h_s1_bf16_ptr,
|
||||
h_s2_bf16_ptr,
|
||||
self.shared_h2, b, self.shared_h1,
|
||||
"bf16_h_s2",
|
||||
)?;
|
||||
self.launch_add_bias_relu_bf16_raw(
|
||||
stream, h_s2_bf16_ptr, bf16_w_ptrs[3], self.shared_h2, b,
|
||||
)?;
|
||||
|
||||
// ── Step 4: Value head layer 1 ───────────────────────────────────
|
||||
// h_v[B, VH] = ReLU(h_s2_bf16 @ W_v1_bf16^T + b_v1_bf16)
|
||||
let h_v_bf16_ptr = raw_u16_ptr(&self.h_v_bf16, stream);
|
||||
|
||||
self.gemmex_bf16(
|
||||
bf16_w_ptrs[4], // W_v1[VH, SH2]
|
||||
h_s2_bf16_ptr,
|
||||
h_v_bf16_ptr,
|
||||
self.value_h, b, self.shared_h2,
|
||||
"bf16_h_v",
|
||||
)?;
|
||||
self.launch_add_bias_relu_bf16_raw(
|
||||
stream, h_v_bf16_ptr, bf16_w_ptrs[5], self.value_h, b,
|
||||
)?;
|
||||
|
||||
// ── Step 5: Value head layer 2 (logits, no ReLU) ────────────────
|
||||
// v_logits[B, NA] = h_v_bf16 @ W_v2_bf16^T + b_v2_bf16
|
||||
let v_logits_bf16_ptr = raw_u16_ptr(&self.v_logits_bf16, stream);
|
||||
|
||||
self.gemmex_bf16(
|
||||
bf16_w_ptrs[6], // W_v2[NA, VH]
|
||||
h_v_bf16_ptr,
|
||||
v_logits_bf16_ptr,
|
||||
self.num_atoms, b, self.value_h,
|
||||
"bf16_v_logits",
|
||||
)?;
|
||||
self.launch_add_bias_bf16_raw(
|
||||
stream, v_logits_bf16_ptr, bf16_w_ptrs[7], self.num_atoms, b,
|
||||
)?;
|
||||
|
||||
// ── Step 6: Branch heads ─────────────────────────────────────────
|
||||
let branch_sizes = [self.branch_0_size, self.branch_1_size, self.branch_2_size];
|
||||
let branch_h_bf16_ptrs = [
|
||||
raw_u16_ptr(&self.h_b0_bf16, stream),
|
||||
raw_u16_ptr(&self.h_b1_bf16, stream),
|
||||
raw_u16_ptr(&self.h_b2_bf16, stream),
|
||||
];
|
||||
let branch_w_base = [8_usize, 12, 16];
|
||||
let b_logits_bf16_ptr = raw_u16_ptr(&self.b_logits_bf16, stream);
|
||||
let na = self.num_atoms;
|
||||
|
||||
let mut logit_byte_offset: u64 = 0;
|
||||
for d in 0..3 {
|
||||
let n_d = branch_sizes[d];
|
||||
let w_fc_idx = branch_w_base[d];
|
||||
let b_fc_idx = w_fc_idx + 1;
|
||||
let w_out_idx = w_fc_idx + 2;
|
||||
let b_out_idx = w_fc_idx + 3;
|
||||
|
||||
// h_bd[B, AH] = ReLU(h_s2_bf16 @ W_bdk_fc_bf16^T + b_bdk_fc_bf16)
|
||||
self.gemmex_bf16(
|
||||
bf16_w_ptrs[w_fc_idx], // W_bdk_fc[AH, SH2]
|
||||
h_s2_bf16_ptr,
|
||||
branch_h_bf16_ptrs[d],
|
||||
self.adv_h, b, self.shared_h2,
|
||||
"bf16_h_bd",
|
||||
)?;
|
||||
self.launch_add_bias_relu_bf16_raw(
|
||||
stream, branch_h_bf16_ptrs[d], bf16_w_ptrs[b_fc_idx], self.adv_h, b,
|
||||
)?;
|
||||
|
||||
// adv_logits_d[B, n_d*NA] = h_bd_bf16 @ W_bdk_out_bf16^T + b_bdk_out_bf16
|
||||
// Output pointer: offset into b_logits_bf16 by accumulated bytes
|
||||
let adv_out_ptr = b_logits_bf16_ptr + logit_byte_offset;
|
||||
self.gemmex_bf16(
|
||||
bf16_w_ptrs[w_out_idx], // W_bdk_out[n_d*NA, AH]
|
||||
branch_h_bf16_ptrs[d],
|
||||
adv_out_ptr,
|
||||
n_d * na, b, self.adv_h,
|
||||
"bf16_adv_logits",
|
||||
)?;
|
||||
self.launch_add_bias_bf16_raw(
|
||||
stream, adv_out_ptr, bf16_w_ptrs[b_out_idx], n_d * na, b,
|
||||
)?;
|
||||
|
||||
logit_byte_offset += (b * n_d * na * std::mem::size_of::<u16>()) as u64;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Run value head forward only: h_s2 → W_v1 → ReLU → W_v2 → v_logits.
|
||||
///
|
||||
/// Used by ensemble heads (1..K-1) to compute per-head value logits
|
||||
@@ -551,7 +836,69 @@ impl CublasForward {
|
||||
}
|
||||
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
// Bias + activation kernel launchers
|
||||
// cublasGemmEx BF16 helpers (tensor core path)
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// BF16 x BF16 -> BF16 GEMM via `cublasGemmEx` with F32 internal accumulation.
|
||||
///
|
||||
/// H100 tensor cores: ~3x throughput vs `cublasSgemm`.
|
||||
///
|
||||
/// Same row-major trick as `sgemm_layer`:
|
||||
/// `C[B, N] = A[B, K] @ W[N, K]^T` via
|
||||
/// `GemmEx(OP_T, OP_N, N, B, K, 1.0, W_bf16, K, A_bf16, K, 0.0, C_bf16, N)`
|
||||
///
|
||||
/// All three matrix pointers (W, A, C) must point to BF16 (`u16`) device memory.
|
||||
/// Alpha/beta are F32 scalars passed as `*const c_void`.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn gemmex_bf16(
|
||||
&self,
|
||||
w_bf16_ptr: u64, // BF16 weights [out_dim, in_dim]
|
||||
a_bf16_ptr: u64, // BF16 input [B, in_dim]
|
||||
c_bf16_ptr: u64, // BF16 output [B, out_dim]
|
||||
n: usize, // out_dim
|
||||
b: usize, // batch
|
||||
k: usize, // in_dim
|
||||
_label: &str,
|
||||
) -> Result<(), MLError> {
|
||||
let alpha = 1.0_f32;
|
||||
let beta = 0.0_f32;
|
||||
|
||||
// SAFETY: w_bf16_ptr, a_bf16_ptr, c_bf16_ptr are valid CUDA device pointers
|
||||
// in BF16 (u16) format. The cuBLAS handle is bound to the correct stream.
|
||||
// cublasGemmEx with CUDA_R_16BF + CUBLAS_COMPUTE_32F uses tensor cores
|
||||
// with F32 accumulation for numerical stability.
|
||||
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,
|
||||
k as i32, // ldb: leading dim of A = K
|
||||
&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 {_label}: {status:?}")));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
// Bias + activation kernel launchers (F32)
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// Launch `add_bias_relu` over an activation buffer.
|
||||
@@ -667,6 +1014,105 @@ impl CublasForward {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
// Bias + activation kernel launchers (BF16 — tensor core path)
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// Launch `add_bias_relu_bf16` over a BF16 activation buffer.
|
||||
///
|
||||
/// `output_bf16[i] = bf16(max(0, f32(output_bf16[i]) + f32(bias_bf16[i % out_dim])))`
|
||||
///
|
||||
/// Both `out_ptr` and `bias_ptr` must point to BF16 (u16) device memory.
|
||||
/// Grid: ceil(B * out_dim / 256), Block: 256.
|
||||
#[allow(dead_code)]
|
||||
fn launch_add_bias_relu_bf16_raw(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
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 {
|
||||
stream
|
||||
.launch_builder(&self.add_bias_relu_bf16_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!("add_bias_relu_bf16 kernel: {e}")))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `add_bias_bf16` (no activation) over a BF16 output buffer.
|
||||
///
|
||||
/// `output_bf16[i] = bf16(f32(output_bf16[i]) + f32(bias_bf16[i % out_dim]))`
|
||||
///
|
||||
/// Both `out_ptr` and `bias_ptr` must point to BF16 (u16) device memory.
|
||||
#[allow(dead_code)]
|
||||
fn launch_add_bias_bf16_raw(
|
||||
&self,
|
||||
stream: &Arc<CudaStream>,
|
||||
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 {
|
||||
stream
|
||||
.launch_builder(&self.add_bias_bf16_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!("add_bias_bf16 kernel: {e}")))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
// BF16 activation buffer accessors (for callers wiring up the full BF16 path)
|
||||
// ══════════════════════════════════════════════════════════════════════════
|
||||
|
||||
/// Raw device pointer to the online states BF16 buffer.
|
||||
#[allow(dead_code)]
|
||||
pub fn states_bf16_ptr(&self, stream: &Arc<CudaStream>) -> u64 {
|
||||
raw_u16_ptr(&self.states_bf16, stream)
|
||||
}
|
||||
|
||||
/// Raw device pointer to the online h_s1 BF16 buffer.
|
||||
#[allow(dead_code)]
|
||||
pub fn h_s1_bf16_ptr(&self, stream: &Arc<CudaStream>) -> u64 {
|
||||
raw_u16_ptr(&self.h_s1_bf16, stream)
|
||||
}
|
||||
|
||||
/// Raw device pointer to the online h_s2 BF16 buffer.
|
||||
#[allow(dead_code)]
|
||||
pub fn h_s2_bf16_ptr(&self, stream: &Arc<CudaStream>) -> u64 {
|
||||
raw_u16_ptr(&self.h_s2_bf16, stream)
|
||||
}
|
||||
}
|
||||
|
||||
// ── Raw device pointer helpers ────────────────────────────────────────────────
|
||||
@@ -685,15 +1131,27 @@ fn raw_f32_ptr_mut(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> u64 {
|
||||
ptr
|
||||
}
|
||||
|
||||
/// Extract raw u16 device pointer (BF16 buffers stored as u16).
|
||||
fn raw_u16_ptr(slice: &CudaSlice<u16>, stream: &Arc<CudaStream>) -> u64 {
|
||||
let (ptr, guard) = slice.device_ptr(stream);
|
||||
let _no_drop = ManuallyDrop::new(guard);
|
||||
ptr
|
||||
}
|
||||
|
||||
// ── Kernel compilation ────────────────────────────────────────────────────────
|
||||
|
||||
/// Precompiled bias kernels cubin (build.rs — ZERO runtime nvcc).
|
||||
static BIAS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bias_kernels.cubin"));
|
||||
|
||||
/// Load `add_bias_relu` and `add_bias` kernels from precompiled cubin.
|
||||
/// Load F32 and BF16 bias+activation kernels from precompiled cubin.
|
||||
///
|
||||
/// Returns (add_bias_relu, add_bias, add_bias_relu_bf16, add_bias_bf16, f32_to_bf16).
|
||||
///
|
||||
/// The `f32_to_bf16_kernel` is defined in `common_device_functions.cuh` which is
|
||||
/// prepended to `bias_kernels.cu` at build time (see `build.rs`).
|
||||
fn compile_bias_kernels(
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> Result<(CudaFunction, CudaFunction), MLError> {
|
||||
) -> Result<(CudaFunction, CudaFunction, CudaFunction, CudaFunction, CudaFunction), MLError> {
|
||||
let context = stream.context();
|
||||
let ptx = cudarc::nvrtc::Ptx::from_binary(BIAS_CUBIN.to_vec());
|
||||
let module = context
|
||||
@@ -706,8 +1164,17 @@ fn compile_bias_kernels(
|
||||
let add_bias = module
|
||||
.load_function("add_bias_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("add_bias_kernel load: {e}")))?;
|
||||
let add_bias_relu_bf16 = module
|
||||
.load_function("add_bias_relu_bf16_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("add_bias_relu_bf16_kernel load: {e}")))?;
|
||||
let add_bias_bf16 = module
|
||||
.load_function("add_bias_bf16_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("add_bias_bf16_kernel load: {e}")))?;
|
||||
let f32_to_bf16 = module
|
||||
.load_function("f32_to_bf16_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!("f32_to_bf16_kernel load: {e}")))?;
|
||||
|
||||
Ok((add_bias_relu, add_bias))
|
||||
Ok((add_bias_relu, add_bias, add_bias_relu_bf16, add_bias_bf16, f32_to_bf16))
|
||||
}
|
||||
|
||||
// ── Compute F32 weight pointers from flat params_buf ─────────────────────────
|
||||
@@ -740,3 +1207,29 @@ pub fn f32_weight_ptrs(
|
||||
}
|
||||
ptrs
|
||||
}
|
||||
|
||||
/// Compute 20 raw BF16 device pointers into a flat `CudaSlice<u16>` buffer.
|
||||
///
|
||||
/// Uses precomputed byte offsets (padded for 4-byte alignment — see
|
||||
/// `bf16_goff_byte_offsets` in `GpuDqnTrainer`). Each segment base is
|
||||
/// `buf_base + byte_offset[i]`.
|
||||
///
|
||||
/// Returns 20 raw u64 device pointers (BF16 stored as u16).
|
||||
#[allow(dead_code)]
|
||||
pub fn bf16_weight_ptrs(
|
||||
bf16_buf: &CudaSlice<u16>,
|
||||
bf16_byte_offsets: &[u64; 20],
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> [u64; 20] {
|
||||
let base = {
|
||||
let (ptr, guard) = bf16_buf.device_ptr(stream);
|
||||
let _no_drop = ManuallyDrop::new(guard);
|
||||
ptr
|
||||
};
|
||||
|
||||
let mut ptrs = [0_u64; 20];
|
||||
for i in 0..20 {
|
||||
ptrs[i] = base + bf16_byte_offsets[i];
|
||||
}
|
||||
ptrs
|
||||
}
|
||||
|
||||
@@ -30,3 +30,47 @@ extern "C" __global__ void add_bias_kernel(
|
||||
if (i >= total_elements) return;
|
||||
output[i] += bias[i % out_dim];
|
||||
}
|
||||
|
||||
/* ------------------------------------------------------------------ */
|
||||
/* BF16 bias+activation kernels for cublasGemmEx tensor core path */
|
||||
/* ------------------------------------------------------------------ */
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
/**
|
||||
* Fused bias-add + ReLU for BF16 hidden layers.
|
||||
* Both output[] and bias[] are __nv_bfloat16.
|
||||
* Intermediate arithmetic in F32 for precision.
|
||||
*
|
||||
* Launch config: grid=(ceil(total_elements/256), 1, 1), block=(256, 1, 1).
|
||||
*/
|
||||
extern "C" __global__ void add_bias_relu_bf16_kernel(
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
const __nv_bfloat16* __restrict__ bias,
|
||||
int out_dim,
|
||||
int total_elements)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= total_elements) return;
|
||||
float val = __bfloat162float(output[i]) + __bfloat162float(bias[i % out_dim]);
|
||||
output[i] = __float2bfloat16((val > 0.0f) ? val : 0.0f);
|
||||
}
|
||||
|
||||
/**
|
||||
* Bias-add only (no activation) for BF16 output layers.
|
||||
* Both output[] and bias[] are __nv_bfloat16.
|
||||
* Intermediate arithmetic in F32 for precision.
|
||||
*
|
||||
* Launch config: grid=(ceil(total_elements/256), 1, 1), block=(256, 1, 1).
|
||||
*/
|
||||
extern "C" __global__ void add_bias_bf16_kernel(
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
const __nv_bfloat16* __restrict__ bias,
|
||||
int out_dim,
|
||||
int total_elements)
|
||||
{
|
||||
int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= total_elements) return;
|
||||
float val = __bfloat162float(output[i]) + __bfloat162float(bias[i % out_dim]);
|
||||
output[i] = __float2bfloat16(val);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user