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

perf: fuse 20 BF16 conversion launches into 1 flat-buffer conversion

Browse Source 
Replace 40 per-tensor f32_to_bf16_kernel launches per EMA step (20 online
+ 20 target) with 2 single-launch conversions over flat contiguous buffers.

- Add bf16_params_buf and bf16_target_params_buf (flat CudaSlice<u16>) that
  mirror the GOFF_* layout of the F32 params_buf/target_params_buf
- Precompute bf16_goff_byte_offsets[20] at construction for zero-cost pointer
  arithmetic into flat BF16 buffers during kernel launches
- sync_online_bf16: single f32_to_bf16_kernel(params_buf, bf16_params_buf, N)
- sync_target_bf16: single f32_to_bf16_kernel(target_params_buf, bf16_target_params_buf, N)
- Forward kernels pass raw u64 device pointers at GOFF offsets instead of
  individual CudaSlice<u16> references — zero additional allocation
- Remove DuelingWeightSetBf16/BranchingWeightSetBf16 dependency from trainer
- Add flat target_params_buf for fused single-kernel EMA update

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-21 12:00:13 +01:00
parent fa191983e6
commit e1b8b46255
1 changed files with 342 additions and 231 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -48,8 +48,6 @@ use cudarc::driver::{
use cudarc::nvrtc::Ptx;
use tracing::info;
use super::gpu_weights::{BranchingWeightSetBf16, DuelingWeightSetBf16};
use crate::MLError;
use super::gpu_weights::{DuelingWeightSet, BranchingWeightSet};
@@ -238,13 +236,17 @@ pub struct GpuDqnTrainer {
f32_to_bf16_kernel: CudaFunction,
bf16_to_f32_kernel: CudaFunction,
// ── BF16 weight mirrors (forward kernel reads BF16 for tensor core throughput) ──
// Allocated lazily on first train_step when weight sets are available.
// Synced after each Adam step (online) and after each EMA update (target).
online_dueling_bf16: Option<DuelingWeightSetBf16>,
online_branching_bf16: Option<BranchingWeightSetBf16>,
target_dueling_bf16: Option<DuelingWeightSetBf16>,
target_branching_bf16: Option<BranchingWeightSetBf16>,
// ── Flat BF16 weight mirrors (forward kernel reads BF16 for tensor core throughput) ──
// Single contiguous BF16 buffer per network (online + target), same GOFF_* layout
// as the F32 flat buffers. One f32_to_bf16_kernel launch converts the entire buffer
// instead of 20 per-tensor launches. Forward kernels read at precomputed byte offsets.
bf16_params_buf: CudaSlice<u16>, // [TOTAL_PARAMS] flat online BF16
bf16_target_params_buf: CudaSlice<u16>, // [TOTAL_PARAMS] flat target BF16
bf16_mirrors_initialized: bool,
/// Precomputed byte offsets into flat BF16 buffers for each of the 20 weight tensors.
/// Layout matches GOFF_* order: w_s1, b_s1, w_s2, b_s2, ..., w_bu2, b_bu2.
/// Each offset is in bytes (element offset * sizeof(u16)).
bf16_goff_byte_offsets: [u64; 20],
// ── Batch input buffers (uploaded per step) ─────────────────────
states_buf: CudaSlice<f32>, // [B, STATE_DIM]
@@ -275,6 +277,7 @@ pub struct GpuDqnTrainer {
// ── Backward / Adam buffers ─────────────────────────────────────
grad_buf: CudaSlice<f32>, // [TOTAL_PARAMS] gradient accumulator
params_buf: CudaSlice<f32>, // [TOTAL_PARAMS] flat online parameters
target_params_buf: CudaSlice<f32>, // [TOTAL_PARAMS] flat target parameters (EMA)
m_buf: CudaSlice<f32>, // [TOTAL_PARAMS] Adam first moment
v_buf: CudaSlice<f32>, // [TOTAL_PARAMS] Adam second moment
grad_norm_buf: CudaSlice<f32>, // [1] pre-clip gradient L2 norm
@@ -286,6 +289,7 @@ pub struct GpuDqnTrainer {
adam_step: i32,
total_params: usize,
params_initialized: bool,
target_params_initialized: bool,
// ── Shared memory size ──────────────────────────────────────────
shmem_bytes: usize,
@@ -468,11 +472,30 @@ impl GpuDqnTrainer {
// ── Allocate backward / Adam buffers ────────────────────────
let grad_buf = alloc_f32(&stream, total_params, "grad_buf")?;
let params_buf = alloc_f32(&stream, total_params, "params_buf")?;
let target_params_buf = alloc_f32(&stream, total_params, "target_params_buf")?;
let m_buf = alloc_f32(&stream, total_params, "adam_m")?;
let v_buf = alloc_f32(&stream, total_params, "adam_v")?;
let grad_norm_buf = alloc_f32(&stream, 1, "grad_norm")?;
let t_buf = alloc_i32(&stream, 1, "adam_t")?;
// ── Allocate flat BF16 weight mirror buffers ──────────────────
// One contiguous BF16 buffer per network; single f32_to_bf16_kernel
// launch converts the entire flat F32 buffer instead of 20 per-tensor
// launches. Forward kernels read at precomputed GOFF byte offsets.
let bf16_params_buf = alloc_u16(&stream, total_params, "bf16_params")?;
let bf16_target_params_buf = alloc_u16(&stream, total_params, "bf16_target_params")?;
// Precompute byte offsets into flat BF16 buffers (GOFF_* layout).
let param_sizes = compute_param_sizes(&config);
let mut bf16_goff_byte_offsets = [0_u64; 20];
{
let mut offset = 0_u64;
for i in 0..20 {
bf16_goff_byte_offsets[i] = offset;
offset += (param_sizes[i] as u64) * (std::mem::size_of::<u16>() as u64);
}
}
// ── Allocate consolidated transfer buffers ─────────────────
// Upload staging: states + next_states + actions(as f32) + rewards + dones + is_weights
let upload_staging_len = b * config.state_dim * 2 + b * 4; // 2*B*SD + 4*B
@@ -512,7 +535,7 @@ impl GpuDqnTrainer {
+ b * num_branches * config.num_atoms * 2
+ b * 2 + 1)
* std::mem::size_of::<f32>();
let optim_bytes = total_params * 4 * std::mem::size_of::<f32>(); // grad + params + m + v
let optim_bytes = total_params * 5 * std::mem::size_of::<f32>(); // grad + params + target_params + m + v
info!(
batch_size = b,
@@ -536,10 +559,10 @@ impl GpuDqnTrainer {
per_update_kernel,
f32_to_bf16_kernel,
bf16_to_f32_kernel,
online_dueling_bf16: None,
online_branching_bf16: None,
target_dueling_bf16: None,
target_branching_bf16: None,
bf16_params_buf,
bf16_target_params_buf,
bf16_mirrors_initialized: false,
bf16_goff_byte_offsets,
states_buf,
next_states_buf,
actions_buf,
@@ -560,6 +583,7 @@ impl GpuDqnTrainer {
q_out_buf,
grad_buf,
params_buf,
target_params_buf,
m_buf,
v_buf,
grad_norm_buf,
@@ -567,6 +591,7 @@ impl GpuDqnTrainer {
adam_step: 0,
total_params,
params_initialized: false,
target_params_initialized: false,
shmem_bytes,
training_graph: None,
upload_staging_buf,
@@ -600,78 +625,112 @@ impl GpuDqnTrainer {
// BF16 weight mirror management
// ═══════════════════════════════════════════════════════════════════
/// Ensure BF16 weight mirrors are allocated and synced from F32 originals.
/// Ensure flat BF16 weight mirrors are synced from F32 originals.
///
/// Called lazily on first `train_step()` or `forward_loss()`. Allocates
/// 4 BF16 mirror sets (online + target, dueling + branching) and syncs
/// initial F32 weights into them via the `f32_to_bf16_kernel`.
/// Called lazily on first `train_step()` or `forward_loss()`. Performs
/// initial F32 → BF16 conversion via 2 kernel launches (1 online + 1 target)
/// over the flat parameter buffers (same GOFF_* layout).
fn ensure_bf16_mirrors(
&mut self,
online_d: &DuelingWeightSet,
online_b: &BranchingWeightSet,
target_d: &DuelingWeightSet,
target_b: &BranchingWeightSet,
_online_d: &DuelingWeightSet,
_online_b: &BranchingWeightSet,
_target_d: &DuelingWeightSet,
_target_b: &BranchingWeightSet,
) -> Result<(), MLError> {
if self.online_dueling_bf16.is_some() {
if self.bf16_mirrors_initialized {
return Ok(());
}
// Allocate BF16 mirrors from F32 weight sets
let mut od_bf16 = DuelingWeightSetBf16::alloc_from(online_d, &self.stream)?;
let mut ob_bf16 = BranchingWeightSetBf16::alloc_from(online_b, &self.stream)?;
let mut td_bf16 = DuelingWeightSetBf16::alloc_from(target_d, &self.stream)?;
let mut tb_bf16 = BranchingWeightSetBf16::alloc_from(target_b, &self.stream)?;
// Single kernel launch: flat F32 params_buf → flat BF16 bf16_params_buf
self.launch_flat_bf16_convert_online()?;
// Single kernel launch: flat F32 target_params_buf → flat BF16 bf16_target_params_buf
self.launch_flat_bf16_convert_target()?;
// Initial F32 → BF16 sync
od_bf16.sync_from_f32(online_d, &self.f32_to_bf16_kernel, &self.stream)?;
ob_bf16.sync_from_f32(online_b, &self.f32_to_bf16_kernel, &self.stream)?;
td_bf16.sync_from_f32(target_d, &self.f32_to_bf16_kernel, &self.stream)?;
tb_bf16.sync_from_f32(target_b, &self.f32_to_bf16_kernel, &self.stream)?;
self.bf16_mirrors_initialized = true;
self.online_dueling_bf16 = Some(od_bf16);
self.online_branching_bf16 = Some(ob_bf16);
self.target_dueling_bf16 = Some(td_bf16);
self.target_branching_bf16 = Some(tb_bf16);
info!("GpuDqnTrainer: BF16 weight mirrors allocated and synced (4 sets, 40 tensors)");
info!(
total_params = self.total_params,
"GpuDqnTrainer: flat BF16 weight mirrors synced (2 kernel launches, 0 per-tensor)"
);
Ok(())
}
/// Sync online BF16 mirrors from F32 weight tensors.
/// Single fused F32 → BF16 conversion: params_buf → bf16_params_buf.
fn launch_flat_bf16_convert_online(&self) -> Result<(), MLError> {
let n = self.total_params;
if n == 0 { return Ok(()); }
let blocks = ((n + 255) / 256) as u32;
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
self.stream
.launch_builder(&self.f32_to_bf16_kernel)
.arg(&self.params_buf)
.arg(&self.bf16_params_buf)
.arg(&(n as i32))
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("flat f32_to_bf16 (online): {e}")))?;
}
Ok(())
}
/// Single fused F32 → BF16 conversion: target_params_buf → bf16_target_params_buf.
fn launch_flat_bf16_convert_target(&self) -> Result<(), MLError> {
let n = self.total_params;
if n == 0 { return Ok(()); }
let blocks = ((n + 255) / 256) as u32;
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
self.stream
.launch_builder(&self.f32_to_bf16_kernel)
.arg(&self.target_params_buf)
.arg(&self.bf16_target_params_buf)
.arg(&(n as i32))
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("flat f32_to_bf16 (target): {e}")))?;
}
Ok(())
}
/// Compute 20 raw device pointers into a flat BF16 buffer at GOFF_* offsets.
///
/// Returns `[ptr_w_s1, ptr_b_s1, ptr_w_s2, ..., ptr_w_bu2, ptr_b_bu2]`.
/// Each pointer is the base of the flat buffer + the precomputed byte offset.
fn bf16_weight_ptrs(&self, bf16_buf: &CudaSlice<u16>) -> [u64; 20] {
let base = raw_device_ptr_u16(bf16_buf, &self.stream);
let mut ptrs = [0_u64; 20];
for i in 0..20 {
ptrs[i] = base + self.bf16_goff_byte_offsets[i];
}
ptrs
}
/// Sync online BF16 mirrors from the flat F32 `params_buf`.
///
/// Called inside the CUDA Graph after `unflatten_online_weights()` so that
/// the next graph replay's forward kernel reads updated BF16 weights.
/// 12 conversion kernel launches (one per dueling tensor) + 8 (branching).
/// Single kernel launch over the entire flat buffer (was 20 per-tensor).
fn sync_online_bf16(
&mut self,
online_d: &DuelingWeightSet,
online_b: &BranchingWeightSet,
&self,
_online_d: &DuelingWeightSet,
_online_b: &BranchingWeightSet,
) -> Result<(), MLError> {
if let Some(ref mut bf16) = self.online_dueling_bf16 {
bf16.sync_from_f32(online_d, &self.f32_to_bf16_kernel, &self.stream)?;
}
if let Some(ref mut bf16) = self.online_branching_bf16 {
bf16.sync_from_f32(online_b, &self.f32_to_bf16_kernel, &self.stream)?;
}
Ok(())
self.launch_flat_bf16_convert_online()
}
/// Sync target BF16 mirrors from F32 weight tensors.
/// Sync target BF16 mirrors from the flat F32 `target_params_buf`.
///
/// Called after `target_ema_update()` so that the next forward kernel
/// reads updated BF16 target weights.
fn sync_target_bf16(
&mut self,
target_d: &DuelingWeightSet,
target_b: &BranchingWeightSet,
) -> Result<(), MLError> {
if let Some(ref mut bf16) = self.target_dueling_bf16 {
bf16.sync_from_f32(target_d, &self.f32_to_bf16_kernel, &self.stream)?;
}
if let Some(ref mut bf16) = self.target_branching_bf16 {
bf16.sync_from_f32(target_b, &self.f32_to_bf16_kernel, &self.stream)?;
}
Ok(())
/// reads updated BF16 target weights. Single kernel launch (was 20 per-tensor).
fn sync_target_bf16(&self) -> Result<(), MLError> {
self.launch_flat_bf16_convert_target()
}
// ═══════════════════════════════════════════════════════════════════
@@ -1065,8 +1124,8 @@ impl GpuDqnTrainer {
&mut self,
states: &CudaSlice<f32>,
batch_size: usize,
online_dueling: &DuelingWeightSet,
online_branching: &BranchingWeightSet,
_online_dueling: &DuelingWeightSet,
_online_branching: &BranchingWeightSet,
) -> Result<&CudaSlice<f32>, MLError> {
if batch_size > self.config.batch_size {
return Err(MLError::ModelError(format!(
@@ -1075,17 +1134,12 @@ impl GpuDqnTrainer {
)));
}
// Ensure BF16 mirrors are initialized (lazy on first call).
// forward_only_q only needs online weights but ensure_bf16_mirrors
// requires all 4 sets. Create dummy target refs if needed.
if self.online_dueling_bf16.is_none() {
// Allocate BF16 mirrors for online weights only
let mut od_bf16 = DuelingWeightSetBf16::alloc_from(online_dueling, &self.stream)?;
let mut ob_bf16 = BranchingWeightSetBf16::alloc_from(online_branching, &self.stream)?;
od_bf16.sync_from_f32(online_dueling, &self.f32_to_bf16_kernel, &self.stream)?;
ob_bf16.sync_from_f32(online_branching, &self.f32_to_bf16_kernel, &self.stream)?;
self.online_dueling_bf16 = Some(od_bf16);
self.online_branching_bf16 = Some(ob_bf16);
// Ensure flat BF16 online buffer is synced from params_buf (lazy on first call).
// forward_only_q only needs online weights — single kernel launch over
// the flat buffer if not yet initialized.
if !self.bf16_mirrors_initialized {
self.launch_flat_bf16_convert_online()?;
self.bf16_mirrors_initialized = true;
}
// Upload states to trainer's states_buf via DtoD copy
@@ -1104,10 +1158,8 @@ impl GpuDqnTrainer {
/// Argument order matches `dqn_forward_only_kernel` in `dqn_training_kernel.cu`:
/// 1 input + 20 online BF16 weights + 1 output + 1 config = 23 args.
fn launch_forward_only(&self, batch_size: usize) -> Result<(), MLError> {
let od = self.online_dueling_bf16.as_ref()
.ok_or_else(|| MLError::ModelError("BF16 online dueling mirrors not initialized".into()))?;
let ob = self.online_branching_bf16.as_ref()
.ok_or_else(|| MLError::ModelError("BF16 online branching mirrors not initialized".into()))?;
// Compute 20 raw device pointers into the flat BF16 online buffer
let op = self.bf16_weight_ptrs(&self.bf16_params_buf);
let batch_size_i32 = batch_size as i32;
@@ -1118,33 +1170,34 @@ impl GpuDqnTrainer {
};
// Safety: argument order matches the extern "C" kernel signature exactly.
// All CudaSlice lifetimes are valid (owned by self). Grid/block = 1 warp per sample.
// All pointers are into the pre-allocated bf16_params_buf (owned by self).
// Grid/block = 1 warp per sample.
unsafe {
self.stream
.launch_builder(&self.forward_only_kernel)
// ── Input states (1) ──────────────────────────────────
.arg(&self.states_buf)
// ── Online network BF16 weights (20) ─────────────────
.arg(&od.w_s1)
.arg(&od.b_s1)
.arg(&od.w_s2)
.arg(&od.b_s2)
.arg(&od.w_v1)
.arg(&od.b_v1)
.arg(&od.w_v2)
.arg(&od.b_v2)
.arg(&od.w_a1)
.arg(&od.b_a1)
.arg(&od.w_a2)
.arg(&od.b_a2)
.arg(&ob.w_bo1)
.arg(&ob.b_bo1)
.arg(&ob.w_bo2)
.arg(&ob.b_bo2)
.arg(&ob.w_bu1)
.arg(&ob.b_bu1)
.arg(&ob.w_bu2)
.arg(&ob.b_bu2)
// ── Online network BF16 weights (20) — views into flat bf16_params_buf ─
.arg(&op[0]) // w_s1
.arg(&op[1]) // b_s1
.arg(&op[2]) // w_s2
.arg(&op[3]) // b_s2
.arg(&op[4]) // w_v1
.arg(&op[5]) // b_v1
.arg(&op[6]) // w_v2
.arg(&op[7]) // b_v2
.arg(&op[8]) // w_a1 (branch 0 FC)
.arg(&op[9]) // b_a1
.arg(&op[10]) // w_a2 (branch 0 out)
.arg(&op[11]) // b_a2
.arg(&op[12]) // w_bo1 (branch 1 FC)
.arg(&op[13]) // b_bo1
.arg(&op[14]) // w_bo2 (branch 1 out)
.arg(&op[15]) // b_bo2
.arg(&op[16]) // w_bu1 (branch 2 FC)
.arg(&op[17]) // b_bu1
.arg(&op[18]) // w_bu2 (branch 2 out)
.arg(&op[19]) // b_bu2
// ── Q-value output (1) ─────────────────────────────────
.arg(&self.q_out_buf)
// ── Config (1) ──────────────────────────────────────────
@@ -1279,8 +1332,8 @@ impl GpuDqnTrainer {
// are stable — the graph replays writes to the same addresses.
self.unflatten_online_weights(online_d, online_b)?;
// ── 5. Sync online BF16 mirrors from updated F32 tensors ─────
// 20 conversion kernel launches (capturable in CUDA Graph).
// ── 5. Sync online BF16 mirror from updated flat params_buf ─────
// Single f32_to_bf16_kernel launch (capturable in CUDA Graph).
// Next graph replay's forward kernel reads updated BF16 weights.
self.sync_online_bf16(online_d, online_b)?;
@@ -1298,12 +1351,10 @@ impl GpuDqnTrainer {
pub fn invalidate_training_graph(&mut self) {
self.training_graph = None;
self.params_initialized = false;
// Drop BF16 mirrors — they'll be reallocated + synced on next train_step.
// This ensures the graph captures fresh BF16 pointers if weight sets change.
self.online_dueling_bf16 = None;
self.online_branching_bf16 = None;
self.target_dueling_bf16 = None;
self.target_branching_bf16 = None;
// Reset BF16 mirrors flag — they'll be re-synced on next train_step.
// Flat BF16 buffers are pre-allocated (no reallocation needed); only the
// content needs refreshing via a single f32_to_bf16_kernel launch each.
self.bf16_mirrors_initialized = false;
}
/// Invalidate cached state after external weight modifications.
@@ -1463,22 +1514,19 @@ impl GpuDqnTrainer {
// Kernel launch methods
// ═══════════════════════════════════════════════════════════════════
/// Launch the forward+loss kernel with BF16 weight mirrors.
/// Launch the forward+loss kernel with flat BF16 weight buffers.
///
/// Argument order matches `dqn_forward_loss_kernel` in `dqn_training_kernel.cu`:
/// 6 batch data + 20 online BF16 weights + 20 target BF16 weights + 8 activation saves
/// + 3 outputs + 2 config = 59 args.
///
/// BF16 mirrors must be populated before this call (via `ensure_bf16_mirrors`).
/// BF16 flat buffers must be synced before this call (via `ensure_bf16_mirrors`).
/// Weight pointers are raw u64 device addresses into the pre-allocated flat
/// `bf16_params_buf` / `bf16_target_params_buf` at precomputed GOFF_* offsets.
fn launch_forward_loss(&self) -> Result<(), MLError> {
let od = self.online_dueling_bf16.as_ref()
.ok_or_else(|| MLError::ModelError("BF16 online dueling mirrors not initialized".into()))?;
let ob = self.online_branching_bf16.as_ref()
.ok_or_else(|| MLError::ModelError("BF16 online branching mirrors not initialized".into()))?;
let td = self.target_dueling_bf16.as_ref()
.ok_or_else(|| MLError::ModelError("BF16 target dueling mirrors not initialized".into()))?;
let tb = self.target_branching_bf16.as_ref()
.ok_or_else(|| MLError::ModelError("BF16 target branching mirrors not initialized".into()))?;
// Compute 20 raw device pointers for online and target flat BF16 buffers
let op = self.bf16_weight_ptrs(&self.bf16_params_buf);
let tp = self.bf16_weight_ptrs(&self.bf16_target_params_buf);
let b = self.config.batch_size;
let batch_size_i32 = b as i32;
@@ -1491,8 +1539,7 @@ impl GpuDqnTrainer {
};
// Safety: argument order matches the extern "C" kernel signature exactly.
// All CudaSlice lifetimes are valid (owned by self).
// BF16 weight pointers (CudaSlice<u16>) match kernel's __nv_bfloat16* params.
// All pointers are into pre-allocated flat BF16 buffers (owned by self).
// Grid/block dimensions match kernel expectations (1 warp per sample).
unsafe {
self.stream
@@ -1504,51 +1551,48 @@ impl GpuDqnTrainer {
.arg(&self.rewards_buf)
.arg(&self.dones_buf)
.arg(&self.is_weights_buf)
// ── Online network BF16 weights (20) ─────────────────
.arg(&od.w_s1)
.arg(&od.b_s1)
.arg(&od.w_s2)
.arg(&od.b_s2)
.arg(&od.w_v1)
.arg(&od.b_v1)
.arg(&od.w_v2)
.arg(&od.b_v2)
// Branch 0 (exposure) — DuelingWeightSet advantage slot
.arg(&od.w_a1)
.arg(&od.b_a1)
.arg(&od.w_a2)
.arg(&od.b_a2)
// Branch 1 (order) — BranchingWeightSet
.arg(&ob.w_bo1)
.arg(&ob.b_bo1)
.arg(&ob.w_bo2)
.arg(&ob.b_bo2)
// Branch 2 (urgency) — BranchingWeightSet
.arg(&ob.w_bu1)
.arg(&ob.b_bu1)
.arg(&ob.w_bu2)
.arg(&ob.b_bu2)
// ── Target network BF16 weights (20) ─────────────────
.arg(&td.w_s1)
.arg(&td.b_s1)
.arg(&td.w_s2)
.arg(&td.b_s2)
.arg(&td.w_v1)
.arg(&td.b_v1)
.arg(&td.w_v2)
.arg(&td.b_v2)
.arg(&td.w_a1)
.arg(&td.b_a1)
.arg(&td.w_a2)
.arg(&td.b_a2)
.arg(&tb.w_bo1)
.arg(&tb.b_bo1)
.arg(&tb.w_bo2)
.arg(&tb.b_bo2)
.arg(&tb.w_bu1)
.arg(&tb.b_bu1)
.arg(&tb.w_bu2)
.arg(&tb.b_bu2)
// ── Online network BF16 weights (20) — views into flat bf16_params_buf ─
.arg(&op[0]) // w_s1
.arg(&op[1]) // b_s1
.arg(&op[2]) // w_s2
.arg(&op[3]) // b_s2
.arg(&op[4]) // w_v1
.arg(&op[5]) // b_v1
.arg(&op[6]) // w_v2
.arg(&op[7]) // b_v2
.arg(&op[8]) // w_a1 (branch 0 FC)
.arg(&op[9]) // b_a1
.arg(&op[10]) // w_a2 (branch 0 out)
.arg(&op[11]) // b_a2
.arg(&op[12]) // w_bo1 (branch 1 FC)
.arg(&op[13]) // b_bo1
.arg(&op[14]) // w_bo2 (branch 1 out)
.arg(&op[15]) // b_bo2
.arg(&op[16]) // w_bu1 (branch 2 FC)
.arg(&op[17]) // b_bu1
.arg(&op[18]) // w_bu2 (branch 2 out)
.arg(&op[19]) // b_bu2
// ── Target network BF16 weights (20) — views into flat bf16_target_params_buf ─
.arg(&tp[0]) // w_s1
.arg(&tp[1]) // b_s1
.arg(&tp[2]) // w_s2
.arg(&tp[3]) // b_s2
.arg(&tp[4]) // w_v1
.arg(&tp[5]) // b_v1
.arg(&tp[6]) // w_v2
.arg(&tp[7]) // b_v2
.arg(&tp[8]) // w_a1
.arg(&tp[9]) // b_a1
.arg(&tp[10]) // w_a2
.arg(&tp[11]) // b_a2
.arg(&tp[12]) // w_bo1
.arg(&tp[13]) // b_bo1
.arg(&tp[14]) // w_bo2
.arg(&tp[15]) // b_bo2
.arg(&tp[16]) // w_bu1
.arg(&tp[17]) // b_bu1
.arg(&tp[18]) // w_bu2
.arg(&tp[19]) // b_bu2
// ── Saved activations (8) ───────────────────────────
.arg(&self.save_h_s1)
.arg(&self.save_h_s2)
@@ -1833,86 +1877,146 @@ impl GpuDqnTrainer {
// GPU-native Polyak EMA target update
// ═══════════════════════════════════════════════════════════════════
/// Copy 20 individual target weight tensors into the flat `target_params_buf`.
///
/// Same GOFF_* layout as `flatten_online_weights()`.
/// Pure device-to-device copies -- zero host roundtrip.
fn flatten_target_weights(
&self,
target_d: &DuelingWeightSet,
target_b: &BranchingWeightSet,
) -> Result<(), MLError> {
let sizes = compute_param_sizes(&self.config);
let dst_base = raw_device_ptr(&self.target_params_buf, &self.stream);
let slices: [&CudaSlice<f32>; 20] = [
&target_d.w_s1, &target_d.b_s1,
&target_d.w_s2, &target_d.b_s2,
&target_d.w_v1, &target_d.b_v1,
&target_d.w_v2, &target_d.b_v2,
&target_d.w_a1, &target_d.b_a1,
&target_d.w_a2, &target_d.b_a2,
&target_b.w_bo1, &target_b.b_bo1,
&target_b.w_bo2, &target_b.b_bo2,
&target_b.w_bu1, &target_b.b_bu1,
&target_b.w_bu2, &target_b.b_bu2,
];
let mut byte_offset: u64 = 0;
for (i, slice) in slices.iter().enumerate() {
let num_bytes = sizes[i] * std::mem::size_of::<f32>();
let src = raw_device_ptr(slice, &self.stream);
dtod_copy(dst_base + byte_offset, src, num_bytes, &self.stream, i, "flatten_target")?;
byte_offset += num_bytes as u64;
}
Ok(())
}
/// Copy flat `target_params_buf` back to 20 individual target weight tensors.
///
/// Called after the fused EMA kernel to scatter updated flat target weights
/// back into the individual `DuelingWeightSet` + `BranchingWeightSet` tensors.
/// Pure device-to-device copies -- zero host roundtrip.
fn unflatten_target_weights(
&self,
target_d: &DuelingWeightSet,
target_b: &BranchingWeightSet,
) -> Result<(), MLError> {
let sizes = compute_param_sizes(&self.config);
let src_base = raw_device_ptr(&self.target_params_buf, &self.stream);
let slices: [&CudaSlice<f32>; 20] = [
&target_d.w_s1, &target_d.b_s1,
&target_d.w_s2, &target_d.b_s2,
&target_d.w_v1, &target_d.b_v1,
&target_d.w_v2, &target_d.b_v2,
&target_d.w_a1, &target_d.b_a1,
&target_d.w_a2, &target_d.b_a2,
&target_b.w_bo1, &target_b.b_bo1,
&target_b.w_bo2, &target_b.b_bo2,
&target_b.w_bu1, &target_b.b_bu1,
&target_b.w_bu2, &target_b.b_bu2,
];
let mut byte_offset: u64 = 0;
for (i, slice) in slices.iter().enumerate() {
let num_bytes = sizes[i] * std::mem::size_of::<f32>();
let dst = raw_device_ptr(slice, &self.stream);
dtod_copy(dst, src_base + byte_offset, num_bytes, &self.stream, i, "unflatten_target")?;
byte_offset += num_bytes as u64;
}
Ok(())
}
// ═══════════════════════════════════════════════════════════════════
// GPU-native Polyak EMA target update (fused single kernel)
// ═══════════════════════════════════════════════════════════════════
/// GPU-native Polyak EMA: `target[i] = (1-tau)*target[i] + tau*online[i]`
///
/// Updates all 20 target weight tensors in-place from the corresponding
/// online weight tensors using the EMA kernel. Runs OUTSIDE the captured
/// CUDA Graph — device pointers are stable so the graph stays valid.
/// Fused single-kernel update over flat parameter buffers. On first call,
/// flattens target weights into `target_params_buf` (same GOFF_* layout as
/// `params_buf`). Then launches ONE EMA kernel over the entire flat buffer
/// instead of 20 per-tensor launches. After the kernel, scatters the updated
/// flat target weights back to individual tensors via `unflatten_target_weights()`.
///
/// Eliminates the reverse-sync → CPU Polyak → forward-sync round-trip
/// that previously required 120 D2D copies + 120 Candle ops per step.
/// Runs OUTSIDE the captured CUDA Graph -- device pointers are stable so
/// the graph stays valid.
pub fn target_ema_update(
&mut self,
online_d: &DuelingWeightSet,
online_b: &BranchingWeightSet,
_online_d: &DuelingWeightSet,
_online_b: &BranchingWeightSet,
target_d: &mut DuelingWeightSet,
target_b: &mut BranchingWeightSet,
tau: f32,
) -> Result<(), MLError> {
let sizes = compute_param_sizes(&self.config);
// Sync stream and clear any stale errors from graph capture phase.
// cudarc stores errors from cuStreamWaitEvent on disabled events during
// graph capture. check_err() consumes them so bind_to_thread() succeeds.
unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
let _ = self.stream.context().check_err();
// Paired (target, online) slices in GOFF_* order (20 pairs)
let pairs: [(&CudaSlice<f32>, &CudaSlice<f32>); 20] = [
(&target_d.w_s1, &online_d.w_s1),
(&target_d.b_s1, &online_d.b_s1),
(&target_d.w_s2, &online_d.w_s2),
(&target_d.b_s2, &online_d.b_s2),
(&target_d.w_v1, &online_d.w_v1),
(&target_d.b_v1, &online_d.b_v1),
(&target_d.w_v2, &online_d.w_v2),
(&target_d.b_v2, &online_d.b_v2),
(&target_d.w_a1, &online_d.w_a1),
(&target_d.b_a1, &online_d.b_a1),
(&target_d.w_a2, &online_d.w_a2),
(&target_d.b_a2, &online_d.b_a2),
(&target_b.w_bo1, &online_b.w_bo1),
(&target_b.b_bo1, &online_b.b_bo1),
(&target_b.w_bo2, &online_b.w_bo2),
(&target_b.b_bo2, &online_b.b_bo2),
(&target_b.w_bu1, &online_b.w_bu1),
(&target_b.b_bu1, &online_b.b_bu1),
(&target_b.w_bu2, &online_b.w_bu2),
(&target_b.b_bu2, &online_b.b_bu2),
];
for (i, (target_slice, online_slice)) in pairs.iter().enumerate() {
let n = sizes[i] as i32;
if sizes[i] == 0 {
continue;
}
let blocks = ((sizes[i] + 255) / 256) as u32;
let launch_cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
// Use raw device pointers — graph capture left stale events on
// weight CudaSlices, making cudarc's launch_builder.arg() fail.
let t_ptr = raw_device_ptr(target_slice, &self.stream);
let o_ptr = raw_device_ptr(online_slice, &self.stream);
unsafe {
self.stream
.launch_builder(&self.ema_kernel)
.arg(&t_ptr)
.arg(&o_ptr)
.arg(&tau)
.arg(&n)
.launch(launch_cfg)
.map_err(|e| {
MLError::ModelError(format!("dqn_ema_kernel launch[{i}]: {e}"))
})?;
}
// First call: flatten target weights into flat buffer (20 DtoD copies, once only).
// Subsequent calls reuse the flat buffer which is kept in sync by unflatten at end.
if !self.target_params_initialized {
self.flatten_target_weights(target_d, target_b)?;
self.target_params_initialized = true;
}
// Single fused EMA kernel over all parameters at once.
// params_buf contains the latest online weights (kept in sync by
// unflatten_online_weights after each Adam step).
let n = self.total_params as i32;
let blocks = ((self.total_params + 255) / 256) as u32;
let launch_cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
let t_ptr = raw_device_ptr(&self.target_params_buf, &self.stream);
let o_ptr = raw_device_ptr(&self.params_buf, &self.stream);
unsafe {
self.stream
.launch_builder(&self.ema_kernel)
.arg(&t_ptr)
.arg(&o_ptr)
.arg(&tau)
.arg(&n)
.launch(launch_cfg)
.map_err(|e| {
MLError::ModelError(format!("dqn_ema_kernel flat launch: {e}"))
})?;
}
// Scatter flat target buffer back to individual weight tensors
self.unflatten_target_weights(target_d, target_b)?;
// Sync target BF16 mirrors from updated F32 target weights
self.sync_target_bf16(target_d, target_b)?;
// Single kernel launch over flat target_params_buf → bf16_target_params_buf
self.sync_target_bf16()?;
Ok(())
}
@@ -2262,6 +2366,13 @@ fn raw_device_ptr_u32(slice: &CudaSlice<u32>, stream: &CudaStream) -> u64 {
ptr
}
/// Extract raw CUDA device pointer from a `CudaSlice<u16>` (BF16 weight buffers).
fn raw_device_ptr_u16(slice: &CudaSlice<u16>, stream: &CudaStream) -> u64 {
let (ptr, guard) = slice.device_ptr(stream);
let _no_drop = std::mem::ManuallyDrop::new(guard);
ptr
}
/// Async device-to-device memcpy with error context.
pub(crate) fn dtod_copy(
dst: u64,