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

perf(rl): fix dqn_bwd occupancy 21→256 threads + zero launch_builder

Browse Source 
dqn_distributional_q_bwd: restructure from 21 threads/block (Q_N_ATOMS)
to 256 threads/block. Each thread maps to one (action, atom) pair.
Softmax reduction stays in warp 0; all 231 threads write gradients in
parallel. Block tree-reduce for loss accumulation (no intra-block
atomicAdd). Was 23% of L40S GPU time at 167μs/call.

Convert ALL remaining 46 launch_builder calls to raw_launch across
dqn.rs (9), ppo.rs (6), outcome_head.rs (7), frd.rs (4), noisy.rs (3),
bucket_routing.rs (5), cfc/step.rs, aux_trunk.rs, snap_features.rs,
loss.rs. Zero launch_builder in crate (1 comment only).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 18:27:59 +02:00
parent dcd851a62d
commit a5b39ecc61
12 changed files with 1114 additions and 862 deletions
Download Patch File Download Diff File Expand all files Collapse all files

193
crates/ml-alpha/cuda/dqn_distributional_q.cu
View File

@@ -29,10 +29,13 @@
// One thread per atom. Each thread reduces over HIDDEN_DIM
// via a strided inner loop; no cross-thread reduction
// needed (each atom is an independent output slot).
// Backward: grid = (B, 1, 1); block = (Q_N_ATOMS, 1, 1).
// Within-block warp-shuffle reduces max + sumexp for the
// softmax. Cross-batch loss accumulation uses atomicAdd ONLY
// into the scalar loss_out — see ATOMIC NOTE below.
// Backward: grid = (B, 1, 1); block = (BWD_BLOCK_SIZE = 256, 1, 1).
// One thread per (action, atom) pair (231 live threads,
// 25 padding). Softmax uses shared-memory reduce (21
// atoms for the taken action). Loss uses a 256-wide
// block-level tree-reduce in shared memory. Cross-batch
// loss accumulation uses atomicAdd ONLY into the scalar
// loss_out — see ATOMIC NOTE below.
//
// ATOMIC NOTE (deliberate, deferred fix): the cross-batch loss
// accumulator uses `atomicAdd(loss_out, ce_b)`. This nominally violates
@@ -103,27 +106,54 @@ extern "C" __global__ void dqn_distributional_q_fwd(
// ─────────────────────────────────────────────────────────────────────
// dqn_distributional_q_bwd:
// Categorical CE backward against a pre-projected Bellman target
// distribution. One block per batch; one thread per atom.
// distribution. One block per batch; one thread per (action, atom)
// pair plus padding threads for occupancy.
//
// Block layout (occupancy-optimised):
// grid = (B, 1, 1)
// block = (BWD_BLOCK_SIZE = 256, 1, 1)
// Threads 0..N_ACTIONS*Q_N_ATOMS-1 (= 0..230) each own one
// `(action, atom)` slot in grad_logits[batch, :]. The 21 threads
// whose action == actions_taken[batch] participate in the softmax
// reduction via shared memory (they are NOT necessarily contiguous
// in warp layout — see virtual-lane mapping below).
// Threads 231..255 are padding and exit after the block-level loss
// reduction.
//
// The previous kernel used block=(Q_N_ATOMS=21,1,1) — only 21 threads
// per block, wasting ~84% of each SM on L40S (128 CUDA cores/SM).
// This version launches 256 threads, achieving full SM occupancy.
//
// Softmax virtual-lane mapping: the taken action's 21 atoms live at
// thread indices `act * Q_N_ATOMS + 0 .. act * Q_N_ATOMS + 20`. These
// are NOT guaranteed to be in the same warp (e.g. act=1 → threads
// 21..41, which straddle the warp 0/1 boundary at thread 32). The
// reduction therefore uses shared memory (not warp shuffle) which
// works for any action index.
//
// Loss accumulation: block-level tree-reduce in shared memory (no
// atomicAdd per `feedback_no_atomicadd.md`). Thread 0 of each block
// writes the per-batch CE to loss_per_batch[batch] and adds to
// loss_out[0] via atomicAdd on a SINGLE scalar (the cross-batch
// accumulator — see ATOMIC NOTE in the file header; deferred fix).
//
// Inputs:
// logits [B × N_ACTIONS × Q_N_ATOMS] — raw atom logits from fwd
// target_dist [B × Q_N_ATOMS] — Bellman-projected target
// distribution for the
// action that was taken
// (computed by the Phase E
// categorical projection
// kernel from γ and the
// target-net atom values).
// actions_taken [B] — integer action index in
// [0, N_ACTIONS) per batch
// sample.
// B — batch size
// Outputs:
// loss_out [1] — Σ_b L_b (UNREDUCED across batches; see ATOMIC
// NOTE in the file header).
// loss_out [1] — Σ_b L_b (cross-batch atomicAdd — see header).
// loss_per_batch [B] — per-sample CE for PER priority update.
// grad_logits [B × N_ACTIONS × Q_N_ATOMS] — ∂L/∂logits. Non-taken
// actions get zero grad.
// ─────────────────────────────────────────────────────────────────────
#define BWD_BLOCK_SIZE 256
#define BWD_K_OUT (N_ACTIONS * Q_N_ATOMS) // 231
extern "C" __global__ void dqn_distributional_q_bwd(
const float* __restrict__ logits, // [B * N_ACTIONS * Q_N_ATOMS]
const float* __restrict__ target_dist, // [B * Q_N_ATOMS]
@@ -134,43 +164,51 @@ extern "C" __global__ void dqn_distributional_q_bwd(
float* __restrict__ grad_logits // [B * N_ACTIONS * Q_N_ATOMS]
) {
const int batch = blockIdx.x;
const int atom = threadIdx.x;
if (batch >= B) return;
if (atom >= Q_N_ATOMS) return;
const int tid = threadIdx.x;
if (batch >= B) return;
// Decompose tid into (action, atom) for the BWD_K_OUT=231 live threads.
// Threads tid >= BWD_K_OUT are padding — they participate in the loss
// tree-reduce but do not read/write the logits/grad arrays.
const int my_action = tid / Q_N_ATOMS; // 0..10 for tid < 231
const int my_atom = tid % Q_N_ATOMS; // 0..20 for tid < 231
const int is_live = (tid < BWD_K_OUT);
const int act = actions_taken[batch];
// Defensive: silently skip invalid actions. Phase E enforces
// bound-checking at the loader boundary; this is the floor.
// ── Invalid action guard ─────────────────────────────────────────
// Defensive: zero all grads and loss for this batch sample.
if (act < 0 || act >= N_ACTIONS) {
// Still zero out grads for this batch sample to prevent any
// contamination on stale buffers.
const int base_all = batch * N_ACTIONS * Q_N_ATOMS;
#pragma unroll
for (int a = 0; a < N_ACTIONS; ++a) {
grad_logits[base_all + a * Q_N_ATOMS + atom] = 0.0f;
if (is_live) {
grad_logits[batch * BWD_K_OUT + tid] = 0.0f;
}
// R7d: zero per-sample CE for invalid-action samples so the
// PER priority update sees a real (zero) magnitude rather than
// stale buffer contents — bounded-input pearl applies.
if (atom == 0) {
if (tid == 0) {
loss_per_batch[batch] = 0.0f;
}
return;
}
const int base_taken = batch * N_ACTIONS * Q_N_ATOMS + act * Q_N_ATOMS;
const int is_taken = is_live && (my_action == act);
// ── Softmax over atoms for the taken action ───────────────────────
// Stage logits into shared so the per-thread access pattern is
// contention-free (one slot per thread, max Q_N_ATOMS = 21).
// ── Softmax over the taken action's atoms ────────────────────────
// The 21 threads with my_action == act load the raw logits for the
// taken action. We use shared memory for the max-reduce and
// sum-exp-reduce since the 21 atoms may span two warps depending
// on the action index.
__shared__ float s_logits[Q_N_ATOMS];
s_logits[atom] = logits[base_taken + atom];
__shared__ float s_exp[Q_N_ATOMS];
__shared__ float s_max;
__shared__ float s_sumexp;
if (is_taken) {
const int base_taken = batch * BWD_K_OUT + act * Q_N_ATOMS;
s_logits[my_atom] = logits[base_taken + my_atom];
}
__syncthreads();
// Atom 0 computes max → all threads read; cheap because Q_N_ATOMS=21
// means atom 0's loop is 20 comparisons in registers.
__shared__ float s_max;
if (atom == 0) {
// Thread 0 computes max over Q_N_ATOMS=21 — sequential is cheaper
// than a tree-reduce for 21 elements.
if (tid == 0) {
float m = s_logits[0];
#pragma unroll
for (int z = 1; z < Q_N_ATOMS; ++z) {
@@ -180,13 +218,14 @@ extern "C" __global__ void dqn_distributional_q_bwd(
}
__syncthreads();
// Per-thread exp; share into smem and let atom 0 compute the sum.
__shared__ float s_exp[Q_N_ATOMS];
s_exp[atom] = expf(s_logits[atom] - s_max);
// Per-atom exp, written by the taken-action threads.
if (is_taken) {
s_exp[my_atom] = expf(s_logits[my_atom] - s_max);
}
__syncthreads();
__shared__ float s_sumexp;
if (atom == 0) {
// Thread 0 computes sum-exp.
if (tid == 0) {
float sum = 0.0f;
#pragma unroll
for (int z = 0; z < Q_N_ATOMS; ++z) {
@@ -196,38 +235,54 @@ extern "C" __global__ void dqn_distributional_q_bwd(
}
__syncthreads();
const float p_z = s_exp[atom] / s_sumexp;
const float t_z = target_dist[batch * Q_N_ATOMS + atom];
// ── Gradient writeback ────────────────────────────────────────────
// Taken-action grad: p - target. Non-taken actions are zero by
// construction (the loss only depends on logits[batch, act, :]).
const int base_all = batch * N_ACTIONS * Q_N_ATOMS;
#pragma unroll
for (int a = 0; a < N_ACTIONS; ++a) {
const int off = base_all + a * Q_N_ATOMS + atom;
grad_logits[off] = (a == act) ? (p_z - t_z) : 0.0f;
// ── Gradient writeback (all BWD_K_OUT=231 live threads) ──────────
// Taken-action threads: grad = p[z] - target[z].
// Non-taken-action threads: grad = 0.
if (is_live) {
float grad_val = 0.0f;
if (is_taken) {
const float p_z = s_exp[my_atom] / s_sumexp;
const float t_z = target_dist[batch * Q_N_ATOMS + my_atom];
grad_val = p_z - t_z;
}
grad_logits[batch * BWD_K_OUT + tid] = grad_val;
}
// ── Loss accumulation ────────────────────────────────────────────
// Atom 0 computes the per-batch CE in a tight loop (Q_N_ATOMS = 21
// is small enough that the reduction is faster than the warp-shuffle
// dance). atomicAdd-into-scalar is the deferred fix — see header.
if (atom == 0) {
float ce = 0.0f;
#pragma unroll
for (int z = 0; z < Q_N_ATOMS; ++z) {
const float pz = s_exp[z] / s_sumexp;
const float pz_c = fmaxf(pz, 1e-7f);
const float tz = target_dist[batch * Q_N_ATOMS + z];
ce -= tz * logf(pz_c);
// ── Per-atom CE contribution (taken-action threads only) ─────────
// Each of the 21 taken-action threads computes its
// `-target[z] * log(p[z])` contribution. Padding and non-taken
// threads contribute 0. The block-level tree-reduce below sums
// these into a single per-batch CE.
float my_ce = 0.0f;
if (is_taken) {
const float pz = s_exp[my_atom] / s_sumexp;
const float pz_c = fmaxf(pz, 1e-7f);
const float tz = target_dist[batch * Q_N_ATOMS + my_atom];
my_ce = -tz * logf(pz_c);
}
// ── Block-level tree-reduce for loss (no atomicAdd within block) ─
// BWD_BLOCK_SIZE=256 = 8 warps. Tree-reduce in shared memory.
__shared__ float s_ce[BWD_BLOCK_SIZE];
s_ce[tid] = my_ce;
__syncthreads();
#pragma unroll
for (int stride = BWD_BLOCK_SIZE / 2; stride > 0; stride >>= 1) {
if (tid < stride) {
s_ce[tid] += s_ce[tid + stride];
}
atomicAdd(loss_out, ce);
// R7d: per-sample CE for PER priority update. Single writer per
// batch (atom 0 only), so non-atomic store. Feeds
// `replay.update_priorities(per_indices, td_per_sample)` via a
// mapped-pinned DtoH after the backward.
__syncthreads();
}
// Thread 0 writes the per-batch CE and accumulates into loss_out.
// The cross-batch atomicAdd on loss_out is the Phase C deferred
// fix (see ATOMIC NOTE in the file header). grad_logits is never
// atomicAdded.
if (tid == 0) {
const float ce = s_ce[0];
loss_per_batch[batch] = ce;
atomicAdd(loss_out, ce);
}
}

195
crates/ml-alpha/src/aux_heads.rs
View File

@@ -63,7 +63,7 @@ use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg,
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut,
};
use ml_core::cuda_autograd::init::scoped_init_seed;
use ml_core::device::MlDevice;
@@ -72,6 +72,7 @@ use rand_chacha::ChaCha8Rng;
use crate::cfc::AUX_HIDDEN;
use crate::pinned_mem::MappedF32Buffer;
use crate::trainer::raw_launch::{RawArgs, raw_launch};
const AUX_HEADS_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/aux_heads.cubin"));
@@ -350,29 +351,36 @@ pub fn aux_heads_fwd_gpu(
debug_assert_eq!(prof_short_logit_d.len(), b_sz_u * N_AUX_HORIZONS);
debug_assert_eq!(size_short_pred_d.len(), b_sz_u * N_AUX_HORIZONS);
let cfg = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),
block_dim: (AUX_HIDDEN as u32, 1, 1),
shared_mem_bytes: (AUX_HIDDEN * std::mem::size_of::<f32>()) as u32,
};
let mut launch = stream.launch_builder(func);
launch
.arg(w_prof_long_d)
.arg(b_prof_long_d)
.arg(w_size_long_d)
.arg(b_size_long_d)
.arg(w_prof_short_d)
.arg(b_prof_short_d)
.arg(w_size_short_d)
.arg(b_size_short_d)
.arg(h_aux_d)
.arg(&b_sz)
.arg(prof_long_logit_d)
.arg(size_long_pred_d)
.arg(prof_short_logit_d)
.arg(size_short_pred_d);
unsafe {
launch.launch(cfg).context("aux_heads_fwd launch")?;
let smem_fwd = (AUX_HIDDEN * std::mem::size_of::<f32>()) as u32;
let rs_fwd = stream.cu_stream();
{
let mut args = RawArgs::new();
args.push_ptr(w_prof_long_d.raw_ptr());
args.push_ptr(b_prof_long_d.raw_ptr());
args.push_ptr(w_size_long_d.raw_ptr());
args.push_ptr(b_size_long_d.raw_ptr());
args.push_ptr(w_prof_short_d.raw_ptr());
args.push_ptr(b_prof_short_d.raw_ptr());
args.push_ptr(w_size_short_d.raw_ptr());
args.push_ptr(b_size_short_d.raw_ptr());
args.push_ptr(h_aux_d.raw_ptr());
args.push_i32(b_sz);
args.push_ptr(prof_long_logit_d.raw_ptr());
args.push_ptr(size_long_pred_d.raw_ptr());
args.push_ptr(prof_short_logit_d.raw_ptr());
args.push_ptr(size_short_pred_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(),
(b_sz as u32, 1, 1),
(AUX_HIDDEN as u32, 1, 1),
smem_fwd,
rs_fwd,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("aux_heads_fwd: {:?}", e))?;
}
}
Ok(())
}
@@ -431,34 +439,41 @@ pub fn aux_heads_bwd_gpu(
debug_assert_eq!(grad_b_size_short_d.len(), b_sz_u * N_AUX_HORIZONS);
debug_assert_eq!(grad_h_aux_d.len(), b_sz_u * AUX_HIDDEN);
let cfg = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),
block_dim: (AUX_HIDDEN as u32, 1, 1),
shared_mem_bytes: (AUX_HIDDEN * std::mem::size_of::<f32>()) as u32,
};
let mut launch = stream.launch_builder(func);
launch
.arg(w_prof_long_d)
.arg(w_size_long_d)
.arg(w_prof_short_d)
.arg(w_size_short_d)
.arg(h_aux_d)
.arg(grad_prof_long_logit_d)
.arg(grad_size_long_pred_d)
.arg(grad_prof_short_logit_d)
.arg(grad_size_short_pred_d)
.arg(&b_sz)
.arg(grad_w_prof_long_d)
.arg(grad_b_prof_long_d)
.arg(grad_w_size_long_d)
.arg(grad_b_size_long_d)
.arg(grad_w_prof_short_d)
.arg(grad_b_prof_short_d)
.arg(grad_w_size_short_d)
.arg(grad_b_size_short_d)
.arg(grad_h_aux_d);
unsafe {
launch.launch(cfg).context("aux_heads_bwd launch")?;
let smem = (AUX_HIDDEN * std::mem::size_of::<f32>()) as u32;
let rs = stream.cu_stream();
{
let mut args = RawArgs::new();
args.push_ptr(w_prof_long_d.raw_ptr());
args.push_ptr(w_size_long_d.raw_ptr());
args.push_ptr(w_prof_short_d.raw_ptr());
args.push_ptr(w_size_short_d.raw_ptr());
args.push_ptr(h_aux_d.raw_ptr());
args.push_ptr(grad_prof_long_logit_d.raw_ptr());
args.push_ptr(grad_size_long_pred_d.raw_ptr());
args.push_ptr(grad_prof_short_logit_d.raw_ptr());
args.push_ptr(grad_size_short_pred_d.raw_ptr());
args.push_i32(b_sz);
args.push_ptr(grad_w_prof_long_d.raw_ptr());
args.push_ptr(grad_b_prof_long_d.raw_ptr());
args.push_ptr(grad_w_size_long_d.raw_ptr());
args.push_ptr(grad_b_size_long_d.raw_ptr());
args.push_ptr(grad_w_prof_short_d.raw_ptr());
args.push_ptr(grad_b_prof_short_d.raw_ptr());
args.push_ptr(grad_w_size_short_d.raw_ptr());
args.push_ptr(grad_b_size_short_d.raw_ptr());
args.push_ptr(grad_h_aux_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(),
(b_sz as u32, 1, 1),
(AUX_HIDDEN as u32, 1, 1),
smem,
rs,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("aux_heads_bwd: {:?}", e))?;
}
}
Ok(())
}
@@ -497,22 +512,28 @@ pub fn aux_bce_loss_gpu(
debug_assert_eq!(valid_count_out_d.len(), 1);
const AUX_LOSS_BLOCK: u32 = 256;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (AUX_LOSS_BLOCK, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = stream.launch_builder(func);
launch
.arg(prof_logit_d)
.arg(y_prof_true_d)
.arg(pos_weight_d)
.arg(&n_total)
.arg(loss_out_d)
.arg(grad_prof_logit_d)
.arg(valid_count_out_d);
unsafe {
launch.launch(cfg).context("aux_bce_loss launch")?;
let rs = stream.cu_stream();
{
let mut args = RawArgs::new();
args.push_ptr(prof_logit_d.raw_ptr());
args.push_ptr(y_prof_true_d.raw_ptr());
args.push_ptr(pos_weight_d.raw_ptr());
args.push_i32(n_total);
args.push_ptr(loss_out_d.raw_ptr());
args.push_ptr(grad_prof_logit_d.raw_ptr());
args.push_ptr(valid_count_out_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(),
(1, 1, 1),
(AUX_LOSS_BLOCK, 1, 1),
0,
rs,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("aux_bce_loss: {:?}", e))?;
}
}
Ok(())
}
@@ -548,22 +569,28 @@ pub fn aux_huber_masked_loss_gpu(
debug_assert_eq!(valid_count_out_d.len(), 1);
const AUX_LOSS_BLOCK: u32 = 256;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (AUX_LOSS_BLOCK, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = stream.launch_builder(func);
launch
.arg(y_size_hat_d)
.arg(y_size_true_d)
.arg(&delta)
.arg(&n_total)
.arg(loss_out_d)
.arg(grad_y_size_hat_d)
.arg(valid_count_out_d);
unsafe {
launch.launch(cfg).context("aux_huber_masked_loss launch")?;
let rs = stream.cu_stream();
{
let mut args = RawArgs::new();
args.push_ptr(y_size_hat_d.raw_ptr());
args.push_ptr(y_size_true_d.raw_ptr());
args.push_f32(delta);
args.push_i32(n_total);
args.push_ptr(loss_out_d.raw_ptr());
args.push_ptr(grad_y_size_hat_d.raw_ptr());
args.push_ptr(valid_count_out_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(),
(1, 1, 1),
(AUX_LOSS_BLOCK, 1, 1),
0,
rs,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("aux_huber_masked_loss: {:?}", e))?;
}
}
Ok(())
}

106
crates/ml-alpha/src/cfc/aux_trunk.rs
View File

@@ -52,7 +52,7 @@ use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg,
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut,
};
use ml_core::cuda_autograd::init::scoped_init_seed;
use ml_core::device::MlDevice;
@@ -60,6 +60,7 @@ use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use crate::pinned_mem::MappedF32Buffer;
use crate::trainer::raw_launch::{RawArgs, raw_launch};
const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_trunk.cubin"));
@@ -260,26 +261,32 @@ pub fn aux_trunk_fwd_gpu(
debug_assert_eq!(h_new_d.len(), (b_sz as usize) * AUX_HIDDEN);
let feat_dim_i: i32 = feat_dim as i32;
let cfg = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),
block_dim: (AUX_HIDDEN as u32, 1, 1),
// x_local (feat_dim floats) + h_old_local (AUX_HIDDEN floats)
shared_mem_bytes: ((feat_dim + AUX_HIDDEN) * std::mem::size_of::<f32>()) as u32,
};
let mut launch = stream.launch_builder(func);
launch
.arg(w_in_d)
.arg(w_rec_d)
.arg(b_d)
.arg(tau_d)
.arg(x_d)
.arg(h_old_d)
.arg(&dt_s)
.arg(&b_sz)
.arg(&feat_dim_i)
.arg(h_new_d);
unsafe {
launch.launch(cfg).context("aux_trunk_fwd launch")?;
let smem = ((feat_dim + AUX_HIDDEN) * std::mem::size_of::<f32>()) as u32;
let rs = stream.cu_stream();
{
let mut args = RawArgs::new();
args.push_ptr(w_in_d.raw_ptr());
args.push_ptr(w_rec_d.raw_ptr());
args.push_ptr(b_d.raw_ptr());
args.push_ptr(tau_d.raw_ptr());
args.push_ptr(x_d.raw_ptr());
args.push_ptr(h_old_d.raw_ptr());
args.push_f32(dt_s);
args.push_i32(b_sz);
args.push_i32(feat_dim_i);
args.push_ptr(h_new_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(),
(b_sz as u32, 1, 1),
(AUX_HIDDEN as u32, 1, 1),
smem,
rs,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("aux_trunk_fwd: {:?}", e))?;
}
}
Ok(())
}
@@ -329,31 +336,38 @@ pub fn aux_trunk_bwd_gpu(
debug_assert_eq!(grad_x_d.len(), b_sz_u * feat_dim);
let feat_dim_i: i32 = feat_dim as i32;
let cfg = LaunchConfig {
grid_dim: (b_sz as u32, 1, 1),
block_dim: (AUX_HIDDEN as u32, 1, 1),
shared_mem_bytes: ((feat_dim + 2 * AUX_HIDDEN) * std::mem::size_of::<f32>()) as u32,
};
let mut launch = stream.launch_builder(func);
launch
.arg(w_in_d)
.arg(w_rec_d)
.arg(b_d)
.arg(tau_d)
.arg(x_d)
.arg(h_old_d)
.arg(grad_h_new_d)
.arg(&dt_s)
.arg(&b_sz)
.arg(&feat_dim_i)
.arg(grad_w_in_d)
.arg(grad_w_rec_d)
.arg(grad_b_d)
.arg(grad_tau_d)
.arg(grad_h_old_d)
.arg(grad_x_d);
unsafe {
launch.launch(cfg).context("aux_trunk_bwd launch")?;
let smem = ((feat_dim + 2 * AUX_HIDDEN) * std::mem::size_of::<f32>()) as u32;
let rs = stream.cu_stream();
{
let mut args = RawArgs::new();
args.push_ptr(w_in_d.raw_ptr());
args.push_ptr(w_rec_d.raw_ptr());
args.push_ptr(b_d.raw_ptr());
args.push_ptr(tau_d.raw_ptr());
args.push_ptr(x_d.raw_ptr());
args.push_ptr(h_old_d.raw_ptr());
args.push_ptr(grad_h_new_d.raw_ptr());
args.push_f32(dt_s);
args.push_i32(b_sz);
args.push_i32(feat_dim_i);
args.push_ptr(grad_w_in_d.raw_ptr());
args.push_ptr(grad_w_rec_d.raw_ptr());
args.push_ptr(grad_b_d.raw_ptr());
args.push_ptr(grad_tau_d.raw_ptr());
args.push_ptr(grad_h_old_d.raw_ptr());
args.push_ptr(grad_x_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(),
(b_sz as u32, 1, 1),
(AUX_HIDDEN as u32, 1, 1),
smem,
rs,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("aux_trunk_bwd: {:?}", e))?;
}
}
Ok(())
}

131
crates/ml-alpha/src/cfc/bucket_routing.rs
View File

@@ -33,8 +33,9 @@
//! pre-allocated device buffers handed in by the trainer.
use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaSlice, CudaStream};
use std::sync::Arc;
use crate::trainer::raw_launch::{RawArgs, raw_launch};
use crate::heads::N_HORIZONS;
@@ -320,10 +321,10 @@ pub fn execute_transition(
// Allocate metadata + scratch buffers. `sorted_indices_d` is scratch — it
// feeds bucket_assign and bucket_iqr but is not part of the returned
// metadata.
let mut bucket_id_per_channel_d = stream
let bucket_id_per_channel_d = stream
.alloc_zeros::<u8>(HIDDEN_DIM)
.context("alloc bucket_id_per_channel_d")?;
let mut sorted_indices_d = stream
let sorted_indices_d = stream
.alloc_zeros::<u32>(HIDDEN_DIM)
.context("alloc sorted_indices_d (scratch)")?;
let mut bucket_channel_offset_d = stream
@@ -332,16 +333,16 @@ pub fn execute_transition(
let mut bucket_dim_k_d = stream
.alloc_zeros::<u32>(N_HORIZONS)
.context("alloc bucket_dim_k_d")?;
let mut channels_in_bucket_d = stream
let channels_in_bucket_d = stream
.alloc_zeros::<u32>(N_HORIZONS * MAX_BUCKET_DIM)
.context("alloc channels_in_bucket_d")?;
let mut heads_w_skip_offset_d = stream
.alloc_zeros::<u32>(N_HORIZONS + 1)
.context("alloc heads_w_skip_offset_d")?;
let mut bucket_tau_iqr_lo_d = stream
let bucket_tau_iqr_lo_d = stream
.alloc_zeros::<f32>(N_HORIZONS)
.context("alloc bucket_tau_iqr_lo_d")?;
let mut bucket_tau_iqr_hi_d = stream
let bucket_tau_iqr_hi_d = stream
.alloc_zeros::<f32>(N_HORIZONS)
.context("alloc bucket_tau_iqr_hi_d")?;
@@ -365,58 +366,54 @@ pub fn execute_transition(
.memcpy_htod(&BUCKET_CHANNEL_OFFSET, &mut heads_w_skip_offset_d)
.context("htod heads_w_skip_offset (=BUCKET_CHANNEL_OFFSET)")?;
let rs = stream.cu_stream();
// Kernel 1: tau_sort_kernel — bitonic-merge sort of τ ascending.
// Single block × 32 threads, 4 passes for HIDDEN_DIM=128. Shared mem:
// HIDDEN_DIM × 8 bytes (4 for key, 4 for value).
let cfg_sort = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: HIDDEN_DIM as u32 * 8,
};
{
let mut launch = stream.launch_builder(&tau_sort_fn);
launch.arg(cfc_tau_d).arg(&mut sorted_indices_d);
unsafe { launch.launch(cfg_sort).context("tau_sort_kernel launch")? };
let mut args = RawArgs::new();
args.push_ptr(cfc_tau_d.raw_ptr());
args.push_ptr(sorted_indices_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
tau_sort_fn.cu_function(), (1,1,1), (32,1,1),
HIDDEN_DIM as u32 * 8, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("tau_sort_kernel: {:?}", e))?;
}
}
// Kernel 2: bucket_assign_kernel — write bucket id per channel via the
// static quintile boundaries `[0, 25, 50, 75, 100, 128]`.
let cfg_assign = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
{
let mut launch = stream.launch_builder(&bucket_assign_fn);
launch
.arg(&sorted_indices_d)
.arg(&mut bucket_id_per_channel_d);
let mut args = RawArgs::new();
args.push_ptr(sorted_indices_d.raw_ptr());
args.push_ptr(bucket_id_per_channel_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg_assign)
.context("bucket_assign_kernel launch")?
};
raw_launch(
bucket_assign_fn.cu_function(), (1,1,1), (HIDDEN_DIM as u32,1,1),
0, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("bucket_assign_kernel: {:?}", e))?;
}
}
// Kernel 3: bucket_iqr_kernel — per-bucket Q1/Q3. N_HORIZONS blocks × 32
// threads (one warp per bucket).
let cfg_iqr = LaunchConfig {
grid_dim: (N_HORIZONS as u32, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 32 * 4,
};
{
let mut launch = stream.launch_builder(&bucket_iqr_fn);
launch
.arg(cfc_tau_d)
.arg(&sorted_indices_d)
.arg(&mut bucket_tau_iqr_lo_d)
.arg(&mut bucket_tau_iqr_hi_d);
let mut args = RawArgs::new();
args.push_ptr(cfc_tau_d.raw_ptr());
args.push_ptr(sorted_indices_d.raw_ptr());
args.push_ptr(bucket_tau_iqr_lo_d.raw_ptr());
args.push_ptr(bucket_tau_iqr_hi_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg_iqr)
.context("bucket_iqr_kernel launch")?
};
raw_launch(
bucket_iqr_fn.cu_function(), (N_HORIZONS as u32,1,1), (32,1,1),
32 * 4, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("bucket_iqr_kernel: {:?}", e))?;
}
}
// Kernel 4: channels_in_bucket_kernel — build the [N_HORIZONS ×
@@ -427,44 +424,36 @@ pub fn execute_transition(
// Single block × HIDDEN_DIM threads; the sentinel-fill loop covers
// 140 slots (N_HORIZONS × MAX_BUCKET_DIM), and the cursor lives in
// shared memory across the per-channel sequential write loop.
let cfg_channels_in_bucket = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
{
let mut launch = stream.launch_builder(&channels_in_bucket_fn);
launch
.arg(&bucket_id_per_channel_d)
.arg(&bucket_dim_k_d)
.arg(&mut channels_in_bucket_d);
let mut args = RawArgs::new();
args.push_ptr(bucket_id_per_channel_d.raw_ptr());
args.push_ptr(bucket_dim_k_d.raw_ptr());
args.push_ptr(channels_in_bucket_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg_channels_in_bucket)
.context("channels_in_bucket_kernel launch")?
};
raw_launch(
channels_in_bucket_fn.cu_function(), (1,1,1), (HIDDEN_DIM as u32,1,1),
0, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("channels_in_bucket_kernel: {:?}", e))?;
}
}
// Kernel 5: heads_compact_kernel — reorder heads_w_skip into compact
// ragged layout. Single block × HIDDEN_DIM threads; shared mem holds
// per-bucket cursors (`N_HORIZONS + 1` u32 = 24 bytes).
let cfg_compact = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: ((N_HORIZONS + 1) * 4) as u32,
};
{
let mut launch = stream.launch_builder(&heads_compact_fn);
launch
.arg(heads_w_skip_orig_d)
.arg(&bucket_id_per_channel_d)
.arg(&bucket_channel_offset_d)
.arg(heads_w_skip_compact_d);
let mut args = RawArgs::new();
args.push_ptr(heads_w_skip_orig_d.raw_ptr());
args.push_ptr(bucket_id_per_channel_d.raw_ptr());
args.push_ptr(bucket_channel_offset_d.raw_ptr());
args.push_ptr(heads_w_skip_compact_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg_compact)
.context("heads_compact_kernel launch")?
};
raw_launch(
heads_compact_fn.cu_function(), (1,1,1), (HIDDEN_DIM as u32,1,1),
((N_HORIZONS + 1) * 4) as u32, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("heads_compact_kernel: {:?}", e))?;
}
}
// Task 9 applies the slack_factor = sqrt(Q3/Q1) widening (spec §3.2) via

54
crates/ml-alpha/src/cfc/snap_features.rs
View File

@@ -10,11 +10,11 @@
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaSlice, CudaStream};
use ml_core::device::MlDevice;
use crate::pinned_mem::MappedF32Buffer;
use crate::trainer::raw_launch::{raw_memcpy_dtod_async, raw_stream_sync};
use crate::trainer::raw_launch::{RawArgs, raw_launch, raw_memcpy_dtod_async, raw_stream_sync};
const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/snap_feature_assemble.cubin"));
@@ -92,7 +92,7 @@ pub fn snap_feature_assemble_gpu(dev: &MlDevice, input: &Mbp10RawInput) -> Resul
let prev_bid_sz = stream.alloc_zeros::<f32>(10).context("prev_bid_sz alloc")?;
let prev_ask_sz = stream.alloc_zeros::<f32>(10).context("prev_ask_sz alloc")?;
let regime = upload(stream, &input.regime)?;
let mut out_d = stream.alloc_zeros::<f32>(FEATURE_DIM).context("out alloc")?;
let out_d = stream.alloc_zeros::<f32>(FEATURE_DIM).context("out alloc")?;
let prev_mid = input.prev_mid;
let trade_signed_vol = input.trade_signed_vol;
@@ -101,29 +101,31 @@ pub fn snap_feature_assemble_gpu(dev: &MlDevice, input: &Mbp10RawInput) -> Resul
let prev_ts_ns = input.prev_ts_ns as i64;
let tick_size = ES_TICK_SIZE;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = stream.launch_builder(&func);
launch
.arg(&bid_px)
.arg(&bid_sz)
.arg(&ask_px)
.arg(&ask_sz)
.arg(&prev_bid_sz)
.arg(&prev_ask_sz)
.arg(&regime)
.arg(&prev_mid)
.arg(&trade_signed_vol)
.arg(&trade_count)
.arg(&ts_ns)
.arg(&prev_ts_ns)
.arg(&tick_size)
.arg(&mut out_d);
unsafe { launch.launch(cfg).context("snap_feature_assemble launch")?; }
{
let rs = stream.cu_stream();
let mut args = RawArgs::new();
args.push_ptr(bid_px.raw_ptr());
args.push_ptr(bid_sz.raw_ptr());
args.push_ptr(ask_px.raw_ptr());
args.push_ptr(ask_sz.raw_ptr());
args.push_ptr(prev_bid_sz.raw_ptr());
args.push_ptr(prev_ask_sz.raw_ptr());
args.push_ptr(regime.raw_ptr());
args.push_f32(prev_mid);
args.push_f32(trade_signed_vol);
args.push_i32(trade_count);
args.push_u64(ts_ns as u64);
args.push_u64(prev_ts_ns as u64);
args.push_f32(tick_size);
args.push_ptr(out_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(), (1,1,1), (1,1,1),
0, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("snap_feature_assemble: {:?}", e))?;
}
}
let v = download(stream, &out_d)?;
let mut out = [0f32; FEATURE_DIM];

195
crates/ml-alpha/src/cfc/step.rs
View File

@@ -3,11 +3,11 @@
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaSlice, CudaStream};
use ml_core::device::MlDevice;
use crate::pinned_mem::MappedF32Buffer;
use crate::trainer::raw_launch::{raw_memcpy_dtod_async, raw_stream_sync};
use crate::trainer::raw_launch::{RawArgs, raw_launch, raw_memcpy_dtod_async, raw_stream_sync};
const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/cfc_step.cubin"));
@@ -52,34 +52,46 @@ pub fn cfc_step_backward_gpu(
let x_d = upload(stream, x)?;
let h_old_d = upload(stream, h_old)?;
let grad_h_new_d = upload(stream, grad_h_new)?;
let mut grad_w_in_d = stream.alloc_zeros::<f32>(w.n_hid * w.n_in).context("grad_w_in alloc")?;
let mut grad_w_rec_d = stream.alloc_zeros::<f32>(w.n_hid * w.n_hid).context("grad_w_rec alloc")?;
let mut grad_b_d = stream.alloc_zeros::<f32>(w.n_hid).context("grad_b alloc")?;
let mut grad_tau_d = stream.alloc_zeros::<f32>(w.n_hid).context("grad_tau alloc")?;
let mut grad_h_old_d = stream.alloc_zeros::<f32>(w.n_hid).context("grad_h_old alloc")?;
let mut grad_x_d = stream.alloc_zeros::<f32>(w.n_in).context("grad_x alloc")?;
let grad_w_in_d = stream.alloc_zeros::<f32>(w.n_hid * w.n_in).context("grad_w_in alloc")?;
let grad_w_rec_d = stream.alloc_zeros::<f32>(w.n_hid * w.n_hid).context("grad_w_rec alloc")?;
let grad_b_d = stream.alloc_zeros::<f32>(w.n_hid).context("grad_b alloc")?;
let grad_tau_d = stream.alloc_zeros::<f32>(w.n_hid).context("grad_tau alloc")?;
let grad_h_old_d = stream.alloc_zeros::<f32>(w.n_hid).context("grad_h_old alloc")?;
let grad_x_d = stream.alloc_zeros::<f32>(w.n_in).context("grad_x alloc")?;
let n_in_i = w.n_in as i32;
let n_hid_i = w.n_hid as i32;
let block_dim = 128.min(w.n_hid as u32);
let grid_dim = ((w.n_hid as u32) + block_dim - 1) / block_dim;
let shared_mem = (2 * w.n_hid * std::mem::size_of::<f32>()) as u32;
let cfg = LaunchConfig {
grid_dim: (grid_dim, 1, 1),
block_dim: (block_dim, 1, 1),
shared_mem_bytes: shared_mem,
};
let mut launch = stream.launch_builder(&func);
launch
.arg(&w_in_d).arg(&w_rec_d).arg(&b_d).arg(&tau_d)
.arg(&x_d).arg(&h_old_d).arg(&grad_h_new_d)
.arg(&dt_s).arg(&n_in_i).arg(&n_hid_i)
.arg(&mut grad_w_in_d).arg(&mut grad_w_rec_d)
.arg(&mut grad_b_d).arg(&mut grad_tau_d)
.arg(&mut grad_h_old_d)
.arg(&mut grad_x_d);
unsafe { launch.launch(cfg).context("cfc_bwd launch")?; }
{
let rs = stream.cu_stream();
let mut args = RawArgs::new();
args.push_ptr(w_in_d.raw_ptr());
args.push_ptr(w_rec_d.raw_ptr());
args.push_ptr(b_d.raw_ptr());
args.push_ptr(tau_d.raw_ptr());
args.push_ptr(x_d.raw_ptr());
args.push_ptr(h_old_d.raw_ptr());
args.push_ptr(grad_h_new_d.raw_ptr());
args.push_f32(dt_s);
args.push_i32(n_in_i);
args.push_i32(n_hid_i);
args.push_ptr(grad_w_in_d.raw_ptr());
args.push_ptr(grad_w_rec_d.raw_ptr());
args.push_ptr(grad_b_d.raw_ptr());
args.push_ptr(grad_tau_d.raw_ptr());
args.push_ptr(grad_h_old_d.raw_ptr());
args.push_ptr(grad_x_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(), (grid_dim, 1, 1), (block_dim, 1, 1),
shared_mem, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("cfc_bwd: {:?}", e))?;
}
}
Ok((
download(stream, &grad_w_in_d)?,
@@ -134,22 +146,35 @@ pub fn cfc_step_per_branch_fwd_gpu(
bucket_dim_k_d: &CudaSlice<u32>,
h_new_d: &mut CudaSlice<f32>,
) -> Result<()> {
const N_HORIZONS: u32 = crate::heads::N_HORIZONS as u32;
const MAX_BUCKET_DIM: u32 = crate::cfc::bucket_routing::MAX_BUCKET_DIM as u32;
const HIDDEN_DIM: u32 = 128;
let cfg = LaunchConfig {
grid_dim: (b_sz as u32, N_HORIZONS, 1),
block_dim: (MAX_BUCKET_DIM, 1, 1),
// 2 × HIDDEN_DIM floats: x_local + h_old_local cooperative staging.
shared_mem_bytes: HIDDEN_DIM * 4 * 2,
};
let mut launch = stream.launch_builder(func);
launch
.arg(w_in_d).arg(w_rec_d).arg(b_d).arg(tau_all_d)
.arg(x_d).arg(h_old_d).arg(&dt_s).arg(&b_sz)
.arg(channels_in_bucket_d).arg(bucket_dim_k_d)
.arg(h_new_d);
unsafe { launch.launch(cfg).context("cfc_step_per_branch_fwd launch")?; }
const FWD_N_HORIZONS: u32 = crate::heads::N_HORIZONS as u32;
const FWD_MAX_BUCKET_DIM: u32 = crate::cfc::bucket_routing::MAX_BUCKET_DIM as u32;
const FWD_HIDDEN_DIM: u32 = 128;
let rs = stream.cu_stream();
{
let mut args = RawArgs::new();
args.push_ptr(w_in_d.raw_ptr());
args.push_ptr(w_rec_d.raw_ptr());
args.push_ptr(b_d.raw_ptr());
args.push_ptr(tau_all_d.raw_ptr());
args.push_ptr(x_d.raw_ptr());
args.push_ptr(h_old_d.raw_ptr());
args.push_f32(dt_s);
args.push_i32(b_sz);
args.push_ptr(channels_in_bucket_d.raw_ptr());
args.push_ptr(bucket_dim_k_d.raw_ptr());
args.push_ptr(h_new_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(),
(b_sz as u32, FWD_N_HORIZONS, 1),
(FWD_MAX_BUCKET_DIM, 1, 1),
FWD_HIDDEN_DIM * 4 * 2,
rs,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("cfc_step_per_branch_fwd: {:?}", e))?;
}
}
Ok(())
}
@@ -191,23 +216,40 @@ pub fn cfc_step_per_branch_bwd_gpu(
grad_tau_all_d: &mut CudaSlice<f32>,
grad_h_old_d: &mut CudaSlice<f32>,
) -> Result<()> {
const N_HORIZONS: u32 = crate::heads::N_HORIZONS as u32;
const MAX_BUCKET_DIM: u32 = crate::cfc::bucket_routing::MAX_BUCKET_DIM as u32;
const HIDDEN_DIM: u32 = 128;
let cfg = LaunchConfig {
grid_dim: (b_sz as u32, N_HORIZONS, 1),
block_dim: (MAX_BUCKET_DIM, 1, 1),
shared_mem_bytes: HIDDEN_DIM * 4 * 2,
};
let mut launch = stream.launch_builder(func);
launch
.arg(w_in_d).arg(w_rec_d).arg(b_d).arg(tau_all_d)
.arg(x_d).arg(h_old_d).arg(grad_h_new_d)
.arg(&dt_s).arg(&b_sz)
.arg(channels_in_bucket_d).arg(bucket_dim_k_d)
.arg(grad_w_in_d).arg(grad_w_rec_d).arg(grad_b_d)
.arg(grad_tau_all_d).arg(grad_h_old_d);
unsafe { launch.launch(cfg).context("cfc_step_per_branch_bwd launch")?; }
const BWD_N_HORIZONS: u32 = crate::heads::N_HORIZONS as u32;
const BWD_MAX_BUCKET_DIM: u32 = crate::cfc::bucket_routing::MAX_BUCKET_DIM as u32;
const BWD_HIDDEN_DIM: u32 = 128;
let rs = stream.cu_stream();
{
let mut args = RawArgs::new();
args.push_ptr(w_in_d.raw_ptr());
args.push_ptr(w_rec_d.raw_ptr());
args.push_ptr(b_d.raw_ptr());
args.push_ptr(tau_all_d.raw_ptr());
args.push_ptr(x_d.raw_ptr());
args.push_ptr(h_old_d.raw_ptr());
args.push_ptr(grad_h_new_d.raw_ptr());
args.push_f32(dt_s);
args.push_i32(b_sz);
args.push_ptr(channels_in_bucket_d.raw_ptr());
args.push_ptr(bucket_dim_k_d.raw_ptr());
args.push_ptr(grad_w_in_d.raw_ptr());
args.push_ptr(grad_w_rec_d.raw_ptr());
args.push_ptr(grad_b_d.raw_ptr());
args.push_ptr(grad_tau_all_d.raw_ptr());
args.push_ptr(grad_h_old_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(),
(b_sz as u32, BWD_N_HORIZONS, 1),
(BWD_MAX_BUCKET_DIM, 1, 1),
BWD_HIDDEN_DIM * 4 * 2,
rs,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("cfc_step_per_branch_bwd: {:?}", e))?;
}
}
Ok(())
}
@@ -232,7 +274,7 @@ pub fn cfc_step_gpu(
let tau_d = upload(stream, &w.tau)?;
let x_d = upload(stream, x)?;
let h_old_d = upload(stream, h_old)?;
let mut h_new_d = stream
let h_new_d = stream
.alloc_zeros::<f32>(w.n_hid)
.context("h_new alloc")?;
@@ -240,25 +282,28 @@ pub fn cfc_step_gpu(
let n_hid_i = w.n_hid as i32;
let block_dim = 128.min(w.n_hid as u32);
let grid_dim = ((w.n_hid as u32) + block_dim - 1) / block_dim;
let cfg = LaunchConfig {
grid_dim: (grid_dim, 1, 1),
block_dim: (block_dim, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = stream.launch_builder(&func);
launch
.arg(&w_in_d)
.arg(&w_rec_d)
.arg(&b_d)
.arg(&tau_d)
.arg(&x_d)
.arg(&h_old_d)
.arg(&dt_s)
.arg(&n_in_i)
.arg(&n_hid_i)
.arg(&mut h_new_d);
unsafe { launch.launch(cfg).context("cfc_step launch")?; }
{
let rs = stream.cu_stream();
let mut args = RawArgs::new();
args.push_ptr(w_in_d.raw_ptr());
args.push_ptr(w_rec_d.raw_ptr());
args.push_ptr(b_d.raw_ptr());
args.push_ptr(tau_d.raw_ptr());
args.push_ptr(x_d.raw_ptr());
args.push_ptr(h_old_d.raw_ptr());
args.push_f32(dt_s);
args.push_i32(n_in_i);
args.push_i32(n_hid_i);
args.push_ptr(h_new_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(), (grid_dim, 1, 1), (block_dim, 1, 1),
0, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("cfc_step: {:?}", e))?;
}
}
download(stream, &h_new_d)
}

272
crates/ml-alpha/src/rl/dqn.rs
View File

@@ -44,9 +44,8 @@ use anyhow::{Context, Result};
use cudarc::cublas::CudaBlas;
use cudarc::cublas::sys as cublas_sys;
use cudarc::cublas::sys::cublasOperation_t;
use cudarc::driver::{
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg,
};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut};
use cudarc::driver::sys::CUstream;
use ml_core::cuda_autograd::init::scoped_init_seed;
use ml_core::cuda_autograd::linear::BiasKernels;
use ml_core::device::MlDevice;
@@ -56,6 +55,7 @@ use rand_chacha::ChaCha8Rng;
use crate::heads::HIDDEN_DIM;
use crate::pinned_mem::MappedF32Buffer;
use crate::rl::common::{N_ACTIONS, Q_N_ATOMS};
use crate::trainer::raw_launch::{RawArgs, raw_launch};
/// Output dimension of the Q-head linear layer: one logit per (action, atom).
const K_OUT: usize = N_ACTIONS * Q_N_ATOMS;
@@ -135,6 +135,10 @@ pub struct DqnHeadConfig {
pub seed: u64,
}
/// Block size for the occupancy-optimised `dqn_distributional_q_bwd`
/// kernel. Matches `BWD_BLOCK_SIZE` in `dqn_distributional_q.cu`.
const BWD_BLOCK_SIZE: u32 = 256;
/// C51 distributional Q-head. Owns device weights (`w_d`, `b_d`) for the
/// online network plus a parallel target network (`w_target_d`,
/// `b_target_d`) used as the bootstrap source for the categorical Bellman
@@ -151,6 +155,9 @@ pub struct DqnHead {
#[allow(dead_code)]
cfg: DqnHeadConfig,
stream: Arc<CudaStream>,
/// Cached raw `CUstream` handle -- avoids repeated `.cu_stream()`
/// calls on the hot path (same pattern as `IntegratedTrainer`).
raw_stream: CUstream,
_module: Arc<CudaModule>,
/// Forward kernel: `dqn_distributional_q_fwd` from
@@ -310,9 +317,12 @@ impl DqnHead {
let bias_kernels = BiasKernels::shared(&stream)
.map_err(|e| anyhow::anyhow!("DqnHead: BiasKernels init: {e}"))?;
let raw_stream = stream.cu_stream();
Ok(Self {
cfg,
stream,
raw_stream,
_module: module,
fwd_fn,
bwd_fn,
@@ -343,46 +353,43 @@ impl DqnHead {
/// `step_with_lobsim` after the Q-head Adam update so the soft
/// update reflects the latest online weights).
pub fn soft_update_target(&mut self, isv_dev_ptr: &u64) -> Result<()> {
let func = self.target_soft_update_fn.cu_function();
// Weights.
{
let n_w = self.w_d.len();
let cfg_w = LaunchConfig {
grid_dim: (((n_w as u32) + 255) / 256, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
let n_w_i = n_w as i32;
let mut launch = self.stream.launch_builder(&self.target_soft_update_fn);
launch
.arg(&self.w_d)
.arg(&mut self.w_target_d)
.arg(isv_dev_ptr)
.arg(&n_w_i);
let mut args = RawArgs::new();
args.push_ptr(self.w_d.raw_ptr());
args.push_ptr(self.w_target_d.raw_ptr());
args.push_ptr(*isv_dev_ptr);
args.push_i32(n_w_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg_w)
.context("dqn_target_soft_update (weights) launch")?;
raw_launch(
func,
(((n_w as u32) + 255) / 256, 1, 1), (256, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("dqn_target_soft_update (weights): {:?}", e))?;
}
}
// Biases.
{
let n_b = self.b_d.len();
let cfg_b = LaunchConfig {
grid_dim: (((n_b as u32) + 31) / 32, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
};
let n_b_i = n_b as i32;
let mut launch = self.stream.launch_builder(&self.target_soft_update_fn);
launch
.arg(&self.b_d)
.arg(&mut self.b_target_d)
.arg(isv_dev_ptr)
.arg(&n_b_i);
let mut args = RawArgs::new();
args.push_ptr(self.b_d.raw_ptr());
args.push_ptr(self.b_target_d.raw_ptr());
args.push_ptr(*isv_dev_ptr);
args.push_i32(n_b_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg_b)
.context("dqn_target_soft_update (biases) launch")?;
raw_launch(
func,
(((n_b as u32) + 31) / 32, 1, 1), (32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("dqn_target_soft_update (biases): {:?}", e))?;
}
}
Ok(())
@@ -463,20 +470,19 @@ impl DqnHead {
let blocks = ((total as u32) + threads - 1) / threads;
let rows_i32 = rows as i32;
let cols_i32 = K_OUT as i32;
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes: 0,
};
let mut args = RawArgs::new();
args.push_ptr(y.raw_ptr());
args.push_ptr(bias.raw_ptr());
args.push_i32(rows_i32);
args.push_i32(cols_i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
self.stream
.launch_builder(&self.bias_kernels.add_fn)
.arg(y)
.arg(bias)
.arg(&rows_i32)
.arg(&cols_i32)
.launch(cfg)
.context("dqn add_bias_2d")?;
raw_launch(
self.bias_kernels.add_fn.cu_function(),
(blocks, 1, 1), (threads, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("dqn add_bias_2d: {:?}", e))?;
}
Ok(())
}
@@ -508,22 +514,22 @@ impl DqnHead {
debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS * Q_N_ATOMS);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (Q_N_ATOMS as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.bwd_fn);
launch
.arg(logits)
.arg(target_dist)
.arg(actions_taken)
.arg(&b_i)
.arg(loss_out_dev_ptr)
.arg(loss_per_batch)
.arg(grad_logits);
let mut args = RawArgs::new();
args.push_ptr(logits.raw_ptr());
args.push_ptr(target_dist.raw_ptr());
args.push_ptr(actions_taken.raw_ptr());
args.push_i32(b_i);
args.push_ptr(*loss_out_dev_ptr);
args.push_ptr(loss_per_batch.raw_ptr());
args.push_ptr(grad_logits.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch.launch(cfg).context("dqn_distributional_q_bwd launch")?;
raw_launch(
self.bwd_fn.cu_function(),
(b_size as u32, 1, 1), (BWD_BLOCK_SIZE, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("dqn_distributional_q_bwd: {:?}", e))?;
}
Ok(())
}
@@ -558,22 +564,23 @@ impl DqnHead {
debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: (K_OUT * std::mem::size_of::<f32>()) as u32,
};
let mut launch = self.stream.launch_builder(&self.grad_w_b_h_t_fn);
launch
.arg(&self.w_d)
.arg(h_t)
.arg(grad_logits)
.arg(&b_i)
.arg(grad_w_per_batch)
.arg(grad_b_per_batch)
.arg(grad_h_t);
let smem = (K_OUT * std::mem::size_of::<f32>()) as u32;
let mut args = RawArgs::new();
args.push_ptr(self.w_d.raw_ptr());
args.push_ptr(h_t.raw_ptr());
args.push_ptr(grad_logits.raw_ptr());
args.push_i32(b_i);
args.push_ptr(grad_w_per_batch.raw_ptr());
args.push_ptr(grad_b_per_batch.raw_ptr());
args.push_ptr(grad_h_t.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch.launch(cfg).context("dqn_grad_w_b_h_t launch")?;
raw_launch(
self.grad_w_b_h_t_fn.cu_function(),
(b_size as u32, 1, 1), (HIDDEN_DIM as u32, 1, 1), smem,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("dqn_grad_w_b_h_t: {:?}", e))?;
}
Ok(())
}
@@ -666,20 +673,19 @@ impl DqnHead {
let blocks = ((K_OUT as u32) + threads - 1) / threads;
let rows_i32 = b_size as i32;
let cols_i32 = K_OUT as i32;
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (threads, 1, 1),
shared_mem_bytes: 0,
};
let mut args = RawArgs::new();
args.push_ptr(grad_logits.raw_ptr());
args.push_ptr(grad_b.raw_ptr());
args.push_i32(rows_i32);
args.push_i32(cols_i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
self.stream
.launch_builder(&self.bias_kernels.reduce_fn)
.arg(grad_logits)
.arg(grad_b)
.arg(&rows_i32)
.arg(&cols_i32)
.launch(cfg)
.context("dqn_bwd_grad_b reduce_sum_axis0")?;
raw_launch(
self.bias_kernels.reduce_fn.cu_function(),
(blocks, 1, 1), (threads, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("dqn_bwd_grad_b reduce_sum_axis0: {:?}", e))?;
}
Ok(())
}
@@ -717,21 +723,19 @@ impl DqnHead {
debug_assert_eq!(action_logits_out.len(), b_size * Q_N_ATOMS);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (Q_N_ATOMS as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.select_action_atoms_fn);
launch
.arg(full_logits_d)
.arg(actions_d)
.arg(&b_i)
.arg(action_logits_out);
let mut args = RawArgs::new();
args.push_ptr(full_logits_d.raw_ptr());
args.push_ptr(actions_d.raw_ptr());
args.push_i32(b_i);
args.push_ptr(action_logits_out.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg)
.context("dqn_select_action_atoms launch")?;
raw_launch(
self.select_action_atoms_fn.cu_function(),
(b_size as u32, 1, 1), (Q_N_ATOMS as u32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("dqn_select_action_atoms: {:?}", e))?;
}
Ok(())
}
@@ -765,24 +769,22 @@ impl DqnHead {
debug_assert_eq!(target_dist_d.len(), b_size * Q_N_ATOMS);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (Q_N_ATOMS as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.bellman_proj_fn);
launch
.arg(target_logits_d)
.arg(rewards_d)
.arg(dones_d)
.arg(n_step_gammas_d)
.arg(isv_dev_ptr)
.arg(&b_i)
.arg(target_dist_d);
let mut args = RawArgs::new();
args.push_ptr(target_logits_d.raw_ptr());
args.push_ptr(rewards_d.raw_ptr());
args.push_ptr(dones_d.raw_ptr());
args.push_ptr(n_step_gammas_d.raw_ptr());
args.push_ptr(*isv_dev_ptr);
args.push_i32(b_i);
args.push_ptr(target_dist_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg)
.context("bellman_target_projection launch")?;
raw_launch(
self.bellman_proj_fn.cu_function(),
(b_size as u32, 1, 1), (Q_N_ATOMS as u32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("bellman_target_projection: {:?}", e))?;
}
Ok(())
}
@@ -816,25 +818,23 @@ impl DqnHead {
debug_assert_eq!(target_dist_d.len(), b_size * Q_N_ATOMS);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (Q_N_ATOMS as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.bellman_fused_fn);
launch
.arg(full_logits_d)
.arg(actions_d)
.arg(rewards_d)
.arg(dones_d)
.arg(n_step_gammas_d)
.arg(isv_dev_ptr)
.arg(&b_i)
.arg(target_dist_d);
let mut args = RawArgs::new();
args.push_ptr(full_logits_d.raw_ptr());
args.push_ptr(actions_d.raw_ptr());
args.push_ptr(rewards_d.raw_ptr());
args.push_ptr(dones_d.raw_ptr());
args.push_ptr(n_step_gammas_d.raw_ptr());
args.push_ptr(*isv_dev_ptr);
args.push_i32(b_i);
args.push_ptr(target_dist_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
launch
.launch(cfg)
.context("bellman_fused_select_project launch")?;
raw_launch(
self.bellman_fused_fn.cu_function(),
(b_size as u32, 1, 1), (Q_N_ATOMS as u32, 1, 1), 0,
self.raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("bellman_fused_select_project: {:?}", e))?;
}
Ok(())
}

162
crates/ml-alpha/src/rl/frd.rs
View File

@@ -46,8 +46,9 @@ use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{
CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg,
CudaFunction, CudaModule, CudaSlice, CudaStream,
};
use cudarc::driver::sys::CUstream;
use ml_core::cuda_autograd::init::scoped_init_seed;
use ml_core::device::MlDevice;
use rand::{Rng, SeedableRng};
@@ -56,6 +57,7 @@ use rand_chacha::ChaCha8Rng;
use crate::heads::HIDDEN_DIM;
use crate::rl::common::{FRD_HIDDEN_DIM, FRD_N_ATOMS, FRD_N_HORIZONS};
use crate::trainer::integrated::write_slice_f32_d_pub;
use crate::trainer::raw_launch::{RawArgs, raw_launch};
const FRD_FWD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rl_frd_fwd.cubin"));
const FRD_SOFTMAX_CE_GRAD_CUBIN: &[u8] =
@@ -89,6 +91,7 @@ impl Default for FrdHeadConfig {
/// composition.
pub struct FrdHead {
stream: Arc<CudaStream>,
raw_stream: CUstream,
_module: Arc<CudaModule>,
fwd_fn: CudaFunction,
_softmax_ce_grad_module: Arc<CudaModule>,
@@ -161,8 +164,10 @@ impl FrdHead {
let mut b2_d = stream.alloc_zeros::<f32>(b2.len())?;
write_slice_f32_d_pub(&stream, &b2, &mut b2_d)?;
let raw_stream = stream.cu_stream();
Ok(Self {
stream,
raw_stream,
_module: module,
fwd_fn,
_softmax_ce_grad_module: softmax_ce_grad_module,
@@ -200,23 +205,26 @@ impl FrdHead {
debug_assert_eq!(logits_d.len(), b_size * FRD_OUT_DIM);
debug_assert_eq!(labels_d.len(), b_size * FRD_N_HORIZONS);
debug_assert_eq!(grad_logits_d.len(), b_size * FRD_OUT_DIM);
let cfg = LaunchConfig {
grid_dim: (b_size as u32, FRD_N_HORIZONS as u32, 1),
block_dim: (FRD_N_ATOMS as u32, 1, 1),
shared_mem_bytes: 0,
};
let b_size_i = b_size as i32;
let mut launch = self.stream.launch_builder(&self.softmax_ce_grad_fn);
launch
.arg(logits_d)
.arg(labels_d)
.arg(&b_size_i)
.arg(grad_logits_d)
.arg(loss_per_batch_h_dev_ptr);
unsafe {
launch
.launch(cfg)
.context("rl_frd_softmax_ce_grad launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(logits_d.raw_ptr());
args.push_ptr(labels_d.raw_ptr());
args.push_i32(b_size_i);
args.push_ptr(grad_logits_d.raw_ptr());
args.push_ptr(*loss_per_batch_h_dev_ptr);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.softmax_ce_grad_fn.cu_function(),
(b_size as u32, FRD_N_HORIZONS as u32, 1),
(FRD_N_ATOMS as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("rl_frd_softmax_ce_grad: {:?}", e))?;
}
}
Ok(())
}
@@ -250,23 +258,28 @@ impl FrdHead {
);
debug_assert_eq!(grad_b2_per_batch_d.len(), b_size * FRD_OUT_DIM);
debug_assert_eq!(grad_hidden_d.len(), b_size * FRD_HIDDEN_DIM);
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (FRD_HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
let b_size_i = b_size as i32;
let mut launch = self.stream.launch_builder(&self.layer2_bwd_fn);
launch
.arg(hidden_d)
.arg(grad_logits_d)
.arg(&self.w2_d)
.arg(&b_size_i)
.arg(grad_w2_per_batch_d)
.arg(grad_b2_per_batch_d)
.arg(grad_hidden_d);
unsafe {
launch.launch(cfg).context("rl_frd_layer2_bwd launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(hidden_d.raw_ptr());
args.push_ptr(grad_logits_d.raw_ptr());
args.push_ptr(self.w2_d.raw_ptr());
args.push_i32(b_size_i);
args.push_ptr(grad_w2_per_batch_d.raw_ptr());
args.push_ptr(grad_b2_per_batch_d.raw_ptr());
args.push_ptr(grad_hidden_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.layer2_bwd_fn.cu_function(),
(b_size as u32, 1, 1),
(FRD_HIDDEN_DIM as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("rl_frd_layer2_bwd: {:?}", e))?;
}
}
Ok(())
}
@@ -306,24 +319,29 @@ impl FrdHead {
);
debug_assert_eq!(grad_b1_per_batch_d.len(), b_size * FRD_HIDDEN_DIM);
debug_assert_eq!(grad_h_t_d.len(), b_size * HIDDEN_DIM);
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
let b_size_i = b_size as i32;
let mut launch = self.stream.launch_builder(&self.layer1_bwd_fn);
launch
.arg(h_t_d)
.arg(hidden_d)
.arg(grad_hidden_d)
.arg(&self.w1_d)
.arg(&b_size_i)
.arg(grad_w1_per_batch_d)
.arg(grad_b1_per_batch_d)
.arg(grad_h_t_d);
unsafe {
launch.launch(cfg).context("rl_frd_layer1_bwd launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(h_t_d.raw_ptr());
args.push_ptr(hidden_d.raw_ptr());
args.push_ptr(grad_hidden_d.raw_ptr());
args.push_ptr(self.w1_d.raw_ptr());
args.push_i32(b_size_i);
args.push_ptr(grad_w1_per_batch_d.raw_ptr());
args.push_ptr(grad_b1_per_batch_d.raw_ptr());
args.push_ptr(grad_h_t_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.layer1_bwd_fn.cu_function(),
(b_size as u32, 1, 1),
(HIDDEN_DIM as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("rl_frd_layer1_bwd: {:?}", e))?;
}
}
Ok(())
}
@@ -341,25 +359,35 @@ impl FrdHead {
debug_assert_eq!(h_t_d.len(), b_size * HIDDEN_DIM);
debug_assert_eq!(hidden_out_d.len(), b_size * FRD_HIDDEN_DIM);
debug_assert_eq!(logits_out_d.len(), b_size * FRD_OUT_DIM);
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (FRD_HIDDEN_DIM as u32, 1, 1), // 64 threads — covers both hidden and output loops
shared_mem_bytes: 0,
};
let b_size_i = b_size as i32;
let mut launch = self.stream.launch_builder(&self.fwd_fn);
launch
.arg(h_t_d)
.arg(&self.w1_d)
.arg(&self.b1_d)
.arg(&self.w2_d)
.arg(&self.b2_d)
.arg(hidden_out_d)
.arg(logits_out_d)
.arg(&b_size_i);
unsafe {
launch.launch(cfg).context("rl_frd_fwd launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(h_t_d.raw_ptr());
args.push_ptr(self.w1_d.raw_ptr());
args.push_ptr(self.b1_d.raw_ptr());
args.push_ptr(self.w2_d.raw_ptr());
args.push_ptr(self.b2_d.raw_ptr());
args.push_ptr(hidden_out_d.raw_ptr());
args.push_ptr(logits_out_d.raw_ptr());
args.push_i32(b_size_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.fwd_fn.cu_function(),
(b_size as u32, 1, 1),
(FRD_HIDDEN_DIM as u32, 1, 1), // 64 threads -- covers both hidden and output loops
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("rl_frd_fwd: {:?}", e))?;
}
}
Ok(())
}
/// Stream used to launch all kernels owned by this head.
pub fn stream(&self) -> &Arc<CudaStream> {
&self.stream
}
}

136
crates/ml-alpha/src/rl/noisy.rs
View File

@@ -42,14 +42,16 @@ use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg,
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut,
};
use cudarc::driver::sys::CUstream;
use ml_core::cuda_autograd::init::scoped_init_seed;
use ml_core::device::MlDevice;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;
use crate::pinned_mem::MappedF32Buffer;
use crate::trainer::raw_launch::{RawArgs, raw_launch};
const SAMPLE_NOISE_CUBIN: &[u8] = include_bytes!(concat!(
env!("OUT_DIR"),
@@ -107,6 +109,7 @@ pub struct NoisyLinear {
in_dim: usize,
out_dim: usize,
stream: Arc<CudaStream>,
raw_stream: CUstream,
_fwd_module: Arc<CudaModule>,
fwd_fn: CudaFunction,
@@ -221,10 +224,12 @@ impl NoisyLinear {
.alloc_zeros::<u32>(1)
.context("noisy_linear prng_state alloc")?;
let raw_stream = stream.cu_stream();
Ok(Self {
in_dim,
out_dim,
stream,
raw_stream,
_fwd_module: fwd_module,
fwd_fn,
_bwd_module: bwd_module,
@@ -251,22 +256,25 @@ impl NoisyLinear {
let max_dim = self.in_dim.max(self.out_dim);
let in_dim_i = self.in_dim as i32;
let out_dim_i = self.out_dim as i32;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (max_dim as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.sample_noise_fn);
launch
.arg(&mut self.noisy_prng_state_d)
.arg(&mut self.eps_in)
.arg(&mut self.eps_out)
.arg(&in_dim_i)
.arg(&out_dim_i);
unsafe {
launch
.launch(cfg)
.context("rl_sample_noise launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(self.noisy_prng_state_d.raw_ptr());
args.push_ptr(self.eps_in.raw_ptr());
args.push_ptr(self.eps_out.raw_ptr());
args.push_i32(in_dim_i);
args.push_i32(out_dim_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.sample_noise_fn.cu_function(),
(1, 1, 1),
(max_dim as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("rl_sample_noise: {:?}", e))?;
}
}
Ok(())
}
@@ -288,28 +296,31 @@ impl NoisyLinear {
let b_i = b_size as i32;
let in_dim_i = self.in_dim as i32;
let out_dim_i = self.out_dim as i32;
let launch_cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (self.out_dim as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.fwd_fn);
launch
.arg(x)
.arg(&self.mu_w)
.arg(&self.sigma_w)
.arg(&self.mu_b)
.arg(&self.sigma_b)
.arg(&self.eps_in)
.arg(&self.eps_out)
.arg(y_out)
.arg(&b_i)
.arg(&in_dim_i)
.arg(&out_dim_i);
unsafe {
launch
.launch(launch_cfg)
.context("rl_noisy_linear_forward launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(x.raw_ptr());
args.push_ptr(self.mu_w.raw_ptr());
args.push_ptr(self.sigma_w.raw_ptr());
args.push_ptr(self.mu_b.raw_ptr());
args.push_ptr(self.sigma_b.raw_ptr());
args.push_ptr(self.eps_in.raw_ptr());
args.push_ptr(self.eps_out.raw_ptr());
args.push_ptr(y_out.raw_ptr());
args.push_i32(b_i);
args.push_i32(in_dim_i);
args.push_i32(out_dim_i);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.fwd_fn.cu_function(),
(b_size as u32, 1, 1),
(self.out_dim as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("rl_noisy_linear_forward: {:?}", e))?;
}
}
Ok(())
}
@@ -341,28 +352,31 @@ impl NoisyLinear {
let b_i = b_size as i32;
let in_dim_i = self.in_dim as i32;
let out_dim_i = self.out_dim as i32;
let launch_cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (self.out_dim as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.bwd_fn);
launch
.arg(x)
.arg(grad_y)
.arg(&self.eps_in)
.arg(&self.eps_out)
.arg(&b_i)
.arg(&in_dim_i)
.arg(&out_dim_i)
.arg(grad_mu_w_pb)
.arg(grad_sigma_w_pb)
.arg(grad_mu_b_pb)
.arg(grad_sigma_b_pb);
unsafe {
launch
.launch(launch_cfg)
.context("rl_noisy_linear_backward launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(x.raw_ptr());
args.push_ptr(grad_y.raw_ptr());
args.push_ptr(self.eps_in.raw_ptr());
args.push_ptr(self.eps_out.raw_ptr());
args.push_i32(b_i);
args.push_i32(in_dim_i);
args.push_i32(out_dim_i);
args.push_ptr(grad_mu_w_pb.raw_ptr());
args.push_ptr(grad_sigma_w_pb.raw_ptr());
args.push_ptr(grad_mu_b_pb.raw_ptr());
args.push_ptr(grad_sigma_b_pb.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.bwd_fn.cu_function(),
(b_size as u32, 1, 1),
(self.out_dim as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("rl_noisy_linear_backward: {:?}", e))?;
}
}
Ok(())
}

241
crates/ml-alpha/src/rl/outcome_head.rs
View File

@@ -10,7 +10,9 @@
use std::sync::Arc;
use anyhow::{anyhow, Result};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream};
use cudarc::driver::sys::CUstream;
use crate::trainer::raw_launch::{RawArgs, raw_launch};
/// Hidden dimension of the trunk output feeding this head.
pub const HIDDEN_DIM: usize = 128;
@@ -32,6 +34,7 @@ pub struct OutcomeHead {
pub b_size: usize,
pub stream: Arc<CudaStream>,
raw_stream: CUstream,
// Kernel handles
_module_fwd: Arc<CudaModule>,
@@ -134,24 +137,27 @@ impl OutcomeHead {
// Labels: sentinel -1 (no label). GPU-pure fill via kernel
// (no memcpy_htod -- per mapped-pinned policy).
let mut labels_d = stream
let labels_d = stream
.alloc_zeros::<i32>(b_size)
.map_err(|e| anyhow!("OutcomeHead: labels alloc: {e}"))?;
let raw_stream = stream.cu_stream();
{
let blocks = ((b_size + 31) / 32) as u32;
let fill_cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
};
let b_i32 = b_size as i32;
let mut args = RawArgs::new();
args.push_ptr(labels_d.raw_ptr());
args.push_i32(b_i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
stream
.launch_builder(&kernel_fill_sentinel)
.arg(&mut labels_d)
.arg(&b_i32)
.launch(fill_cfg)
.map_err(|e| anyhow!("rl_outcome_fill_sentinel init launch: {e}"))?;
raw_launch(
kernel_fill_sentinel.cu_function(),
(blocks, 1, 1),
(32, 1, 1),
0,
raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow!("rl_outcome_fill_sentinel init: {:?}", e))?;
}
}
@@ -164,6 +170,7 @@ impl OutcomeHead {
grad_logits_d,
b_size,
stream,
raw_stream,
_module_fwd: module_fwd,
_module_ce: module_ce,
_module_label: module_label,
@@ -190,21 +197,25 @@ impl OutcomeHead {
));
}
let b_i32 = actual_b as i32;
let cfg = LaunchConfig {
grid_dim: (actual_b as u32, 1, 1),
block_dim: (K_CLASSES as u32, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
self.stream
.launch_builder(&self.kernel_fwd)
.arg(h_t)
.arg(&self.w_d)
.arg(&self.b_d)
.arg(&mut self.logits_d)
.arg(&b_i32)
.launch(cfg)
.map_err(|e| anyhow!("rl_outcome_fwd launch: {e}"))?;
{
let mut args = RawArgs::new();
args.push_ptr(h_t.raw_ptr());
args.push_ptr(self.w_d.raw_ptr());
args.push_ptr(self.b_d.raw_ptr());
args.push_ptr(self.logits_d.raw_ptr());
args.push_i32(b_i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.kernel_fwd.cu_function(),
(actual_b as u32, 1, 1),
(K_CLASSES as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow!("rl_outcome_fwd: {:?}", e))?;
}
}
Ok(())
}
@@ -221,21 +232,25 @@ impl OutcomeHead {
));
}
let b_i32 = actual_b as i32;
let cfg = LaunchConfig {
grid_dim: (actual_b as u32, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
self.stream
.launch_builder(&self.kernel_ce)
.arg(&self.logits_d)
.arg(&self.labels_d)
.arg(&mut self.loss_pb_d)
.arg(&mut self.grad_logits_d)
.arg(&b_i32)
.launch(cfg)
.map_err(|e| anyhow!("rl_outcome_ce launch: {e}"))?;
{
let mut args = RawArgs::new();
args.push_ptr(self.logits_d.raw_ptr());
args.push_ptr(self.labels_d.raw_ptr());
args.push_ptr(self.loss_pb_d.raw_ptr());
args.push_ptr(self.grad_logits_d.raw_ptr());
args.push_i32(b_i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.kernel_ce.cu_function(),
(actual_b as u32, 1, 1),
(1, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow!("rl_outcome_ce: {:?}", e))?;
}
}
Ok(())
}
@@ -271,23 +286,27 @@ impl OutcomeHead {
debug_assert_eq!(grad_b_per_batch_d.len(), actual_b * K_CLASSES);
debug_assert_eq!(grad_h_t_d.len(), actual_b * HIDDEN_DIM);
let b_i32 = actual_b as i32;
let cfg = LaunchConfig {
grid_dim: (actual_b as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
self.stream
.launch_builder(&self.kernel_bwd)
.arg(h_t)
.arg(&self.grad_logits_d)
.arg(&self.w_d)
.arg(&b_i32)
.arg(grad_w_per_batch_d)
.arg(grad_b_per_batch_d)
.arg(grad_h_t_d)
.launch(cfg)
.map_err(|e| anyhow!("rl_outcome_bwd launch: {e}"))?;
{
let mut args = RawArgs::new();
args.push_ptr(h_t.raw_ptr());
args.push_ptr(self.grad_logits_d.raw_ptr());
args.push_ptr(self.w_d.raw_ptr());
args.push_i32(b_i32);
args.push_ptr(grad_w_per_batch_d.raw_ptr());
args.push_ptr(grad_b_per_batch_d.raw_ptr());
args.push_ptr(grad_h_t_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.kernel_bwd.cu_function(),
(actual_b as u32, 1, 1),
(HIDDEN_DIM as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow!("rl_outcome_bwd: {:?}", e))?;
}
}
Ok(())
}
@@ -325,25 +344,29 @@ impl OutcomeHead {
let b_i32 = actual_b as i32;
// shared memory: s_logits[3] + s_grad_logits[3] = 6 floats = 24 bytes
let cfg = LaunchConfig {
grid_dim: (actual_b as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0, // statically declared in kernel
};
unsafe {
self.stream
.launch_builder(&self.kernel_fused)
.arg(h_t)
.arg(&self.w_d)
.arg(&self.b_d)
.arg(&self.labels_d)
.arg(&b_i32)
.arg(&mut self.loss_pb_d)
.arg(grad_w_per_batch_d)
.arg(grad_b_per_batch_d)
.arg(grad_h_t_d)
.launch(cfg)
.map_err(|e| anyhow!("rl_outcome_fused launch: {e}"))?;
{
let mut args = RawArgs::new();
args.push_ptr(h_t.raw_ptr());
args.push_ptr(self.w_d.raw_ptr());
args.push_ptr(self.b_d.raw_ptr());
args.push_ptr(self.labels_d.raw_ptr());
args.push_i32(b_i32);
args.push_ptr(self.loss_pb_d.raw_ptr());
args.push_ptr(grad_w_per_batch_d.raw_ptr());
args.push_ptr(grad_b_per_batch_d.raw_ptr());
args.push_ptr(grad_h_t_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.kernel_fused.cu_function(),
(actual_b as u32, 1, 1),
(HIDDEN_DIM as u32, 1, 1),
0, // statically declared in kernel
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow!("rl_outcome_fused: {:?}", e))?;
}
}
Ok(())
}
@@ -366,20 +389,24 @@ impl OutcomeHead {
}
let b_i32 = actual_b as i32;
let blocks = ((actual_b + 31) / 32) as u32;
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
};
unsafe {
self.stream
.launch_builder(&self.kernel_label)
.arg(rewards)
.arg(dones)
.arg(&mut self.labels_d)
.arg(&b_i32)
.launch(cfg)
.map_err(|e| anyhow!("rl_outcome_label launch: {e}"))?;
{
let mut args = RawArgs::new();
args.push_ptr(rewards.raw_ptr());
args.push_ptr(dones.raw_ptr());
args.push_ptr(self.labels_d.raw_ptr());
args.push_i32(b_i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.kernel_label.cu_function(),
(blocks, 1, 1),
(32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow!("rl_outcome_label: {:?}", e))?;
}
}
Ok(())
}
@@ -387,19 +414,23 @@ impl OutcomeHead {
/// Reset all labels to sentinel -1 (GPU-pure, no memcpy_htod).
pub fn reset_labels(&mut self) -> Result<()> {
let blocks = ((self.b_size + 31) / 32) as u32;
let cfg = LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
};
let b_i32 = self.b_size as i32;
unsafe {
self.stream
.launch_builder(&self.kernel_fill_sentinel)
.arg(&mut self.labels_d)
.arg(&b_i32)
.launch(cfg)
.map_err(|e| anyhow!("rl_outcome_fill_sentinel reset launch: {e}"))?;
{
let mut args = RawArgs::new();
args.push_ptr(self.labels_d.raw_ptr());
args.push_i32(b_i32);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.kernel_fill_sentinel.cu_function(),
(blocks, 1, 1),
(32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow!("rl_outcome_fill_sentinel reset: {:?}", e))?;
}
}
Ok(())
}

243
crates/ml-alpha/src/rl/ppo.rs
View File

@@ -35,8 +35,9 @@ use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg,
CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut,
};
use cudarc::driver::sys::CUstream;
use ml_core::cuda_autograd::init::scoped_init_seed;
use ml_core::device::MlDevice;
use rand::{Rng, SeedableRng};
@@ -45,6 +46,7 @@ use rand_chacha::ChaCha8Rng;
use crate::heads::HIDDEN_DIM;
use crate::pinned_mem::MappedF32Buffer;
use crate::rl::common::N_ACTIONS;
use crate::trainer::raw_launch::{RawArgs, raw_launch};
const PPO_SURR_CUBIN: &[u8] = include_bytes!(concat!(
env!("OUT_DIR"),
@@ -78,6 +80,7 @@ pub struct PolicyHead {
#[allow(dead_code)]
cfg: PpoHeadsConfig,
stream: Arc<CudaStream>,
raw_stream: CUstream,
_module: Arc<CudaModule>,
/// Forward kernel: `ppo_clipped_surrogate_fwd`. Consumes new-policy
@@ -153,10 +156,12 @@ impl PolicyHead {
let w_d = upload(&stream, &w_host)?;
let b_d = upload(&stream, &b_host)?;
let raw_stream = stream.cu_stream();
Ok(Self {
cfg,
stream,
raw_stream,
_module: module,
surrogate_fwd_fn,
surrogate_bwd_fn,
@@ -193,20 +198,25 @@ impl PolicyHead {
debug_assert_eq!(logits_out.len(), b_size * N_ACTIONS);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, N_ACTIONS as u32, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.policy_linear_fwd_fn);
launch
.arg(&self.w_d)
.arg(&self.b_d)
.arg(h_t)
.arg(&b_i)
.arg(logits_out);
unsafe {
launch.launch(cfg).context("ppo_policy_logits_fwd launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(self.w_d.raw_ptr());
args.push_ptr(self.b_d.raw_ptr());
args.push_ptr(h_t.raw_ptr());
args.push_i32(b_i);
args.push_ptr(logits_out.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.policy_linear_fwd_fn.cu_function(),
(b_size as u32, N_ACTIONS as u32, 1),
(HIDDEN_DIM as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("ppo_policy_logits_fwd: {:?}", e))?;
}
}
Ok(())
}
@@ -243,27 +253,30 @@ impl PolicyHead {
debug_assert_eq!(entropy.len(), b_size);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (N_ACTIONS as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.surrogate_fwd_fn);
launch
.arg(logits)
.arg(log_pi_old)
.arg(actions)
.arg(advantages)
.arg(isv_dev_ptr)
.arg(&b_i)
.arg(pi_log_prob)
.arg(entropy)
.arg(loss_pi_dev_ptr)
.arg(loss_entropy_dev_ptr);
unsafe {
launch
.launch(cfg)
.context("ppo_clipped_surrogate_fwd launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(logits.raw_ptr());
args.push_ptr(log_pi_old.raw_ptr());
args.push_ptr(actions.raw_ptr());
args.push_ptr(advantages.raw_ptr());
args.push_ptr(*isv_dev_ptr);
args.push_i32(b_i);
args.push_ptr(pi_log_prob.raw_ptr());
args.push_ptr(entropy.raw_ptr());
args.push_ptr(*loss_pi_dev_ptr);
args.push_ptr(*loss_entropy_dev_ptr);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.surrogate_fwd_fn.cu_function(),
(b_size as u32, 1, 1),
(N_ACTIONS as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("ppo_clipped_surrogate_fwd: {:?}", e))?;
}
}
Ok(())
}
@@ -288,24 +301,27 @@ impl PolicyHead {
debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (N_ACTIONS as u32, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = self.stream.launch_builder(&self.surrogate_bwd_fn);
launch
.arg(logits)
.arg(log_pi_old)
.arg(actions)
.arg(advantages)
.arg(isv_dev_ptr)
.arg(&b_i)
.arg(grad_logits);
unsafe {
launch
.launch(cfg)
.context("ppo_clipped_surrogate_bwd launch")?;
{
let mut args = RawArgs::new();
args.push_ptr(logits.raw_ptr());
args.push_ptr(log_pi_old.raw_ptr());
args.push_ptr(actions.raw_ptr());
args.push_ptr(advantages.raw_ptr());
args.push_ptr(*isv_dev_ptr);
args.push_i32(b_i);
args.push_ptr(grad_logits.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.surrogate_bwd_fn.cu_function(),
(b_size as u32, 1, 1),
(N_ACTIONS as u32, 1, 1),
0,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("ppo_clipped_surrogate_bwd: {:?}", e))?;
}
}
Ok(())
}
@@ -332,22 +348,28 @@ impl PolicyHead {
debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: (N_ACTIONS * std::mem::size_of::<f32>()) as u32,
};
let mut launch = self.stream.launch_builder(&self.grad_w_b_h_t_fn);
launch
.arg(&self.w_d)
.arg(h_t)
.arg(grad_logits)
.arg(&b_i)
.arg(grad_w_per_batch)
.arg(grad_b_per_batch)
.arg(grad_h_t);
unsafe {
launch.launch(cfg).context("ppo_grad_w_b_h_t launch")?;
let smem = (N_ACTIONS * std::mem::size_of::<f32>()) as u32;
{
let mut args = RawArgs::new();
args.push_ptr(self.w_d.raw_ptr());
args.push_ptr(h_t.raw_ptr());
args.push_ptr(grad_logits.raw_ptr());
args.push_i32(b_i);
args.push_ptr(grad_w_per_batch.raw_ptr());
args.push_ptr(grad_b_per_batch.raw_ptr());
args.push_ptr(grad_h_t.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.grad_w_b_h_t_fn.cu_function(),
(b_size as u32, 1, 1),
(HIDDEN_DIM as u32, 1, 1),
smem,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("ppo_grad_w_b_h_t: {:?}", e))?;
}
}
Ok(())
}
@@ -361,6 +383,7 @@ pub struct ValueHead {
#[allow(dead_code)]
cfg: PpoHeadsConfig,
stream: Arc<CudaStream>,
raw_stream: CUstream,
_module: Arc<CudaModule>,
/// Forward kernel: `v_head_fwd`. Reads h_t and emits per-batch
@@ -411,10 +434,12 @@ impl ValueHead {
let w_d = upload(&stream, &w_host)?;
let b_d = upload(&stream, &b_host)?;
let raw_stream = stream.cu_stream();
Ok(Self {
cfg,
stream,
raw_stream,
_module: module,
fwd_fn,
bwd_fn,
@@ -453,21 +478,27 @@ impl ValueHead {
debug_assert_eq!(v_pred.len(), b_size);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
};
let mut launch = self.stream.launch_builder(&self.fwd_fn);
launch
.arg(&self.w_d)
.arg(&self.b_d)
.arg(h_t)
.arg(isv_dev_ptr)
.arg(&b_i)
.arg(v_pred);
unsafe {
launch.launch(cfg).context("v_head_fwd launch")?;
let smem = (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
{
let mut args = RawArgs::new();
args.push_ptr(self.w_d.raw_ptr());
args.push_ptr(self.b_d.raw_ptr());
args.push_ptr(h_t.raw_ptr());
args.push_ptr(*isv_dev_ptr);
args.push_i32(b_i);
args.push_ptr(v_pred.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.fwd_fn.cu_function(),
(b_size as u32, 1, 1),
(HIDDEN_DIM as u32, 1, 1),
smem,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("v_head_fwd: {:?}", e))?;
}
}
Ok(())
}
@@ -497,24 +528,30 @@ impl ValueHead {
debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM);
let b_i = b_size as i32;
let cfg = LaunchConfig {
grid_dim: (b_size as u32, 1, 1),
block_dim: (HIDDEN_DIM as u32, 1, 1),
shared_mem_bytes: (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32,
};
let mut launch = self.stream.launch_builder(&self.bwd_fn);
launch
.arg(&self.w_d)
.arg(h_t)
.arg(v_pred)
.arg(r_target)
.arg(&b_i)
.arg(loss_per_batch)
.arg(grad_w_per_batch)
.arg(grad_b_per_batch)
.arg(grad_h_t);
unsafe {
launch.launch(cfg).context("v_head_bwd launch")?;
let smem = (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
{
let mut args = RawArgs::new();
args.push_ptr(self.w_d.raw_ptr());
args.push_ptr(h_t.raw_ptr());
args.push_ptr(v_pred.raw_ptr());
args.push_ptr(r_target.raw_ptr());
args.push_i32(b_i);
args.push_ptr(loss_per_batch.raw_ptr());
args.push_ptr(grad_w_per_batch.raw_ptr());
args.push_ptr(grad_b_per_batch.raw_ptr());
args.push_ptr(grad_h_t.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.bwd_fn.cu_function(),
(b_size as u32, 1, 1),
(HIDDEN_DIM as u32, 1, 1),
smem,
self.raw_stream,
&mut ptrs[..args.len()],
)
.map_err(|e| anyhow::anyhow!("v_head_bwd: {:?}", e))?;
}
}
Ok(())
}

48
crates/ml-alpha/src/trainer/loss.rs
View File

@@ -7,11 +7,11 @@
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use cudarc::driver::{CudaSlice, CudaStream};
use ml_core::device::MlDevice;
use crate::pinned_mem::MappedF32Buffer;
use crate::trainer::raw_launch::{raw_memcpy_dtod_async, raw_stream_sync};
use crate::trainer::raw_launch::{RawArgs, raw_launch, raw_memcpy_dtod_async, raw_stream_sync};
const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/bce_loss_multi_horizon.cubin"));
@@ -57,26 +57,36 @@ pub fn bce_multi_horizon_loss_and_grad_gpu(dev: &MlDevice, input: &BceInput) ->
.unwrap_or_else(|| vec![0.0; input.n_horizons]);
let sigma_d = upload(stream, &sigma_host)?;
let mut loss_d = stream.alloc_zeros::<f32>(1).context("loss alloc")?;
let mut mean_bce_h_d = stream.alloc_zeros::<f32>(input.n_horizons)?;
let mut grad_d = stream.alloc_zeros::<f32>(total).context("grad alloc")?;
let mut valid_d = stream.alloc_zeros::<i32>(1).context("valid alloc")?;
let mut d_sigma_d = stream.alloc_zeros::<f32>(input.n_horizons)?;
let loss_d = stream.alloc_zeros::<f32>(1).context("loss alloc")?;
let mean_bce_h_d = stream.alloc_zeros::<f32>(input.n_horizons)?;
let grad_d = stream.alloc_zeros::<f32>(total).context("grad alloc")?;
let valid_d = stream.alloc_zeros::<i32>(1).context("valid alloc")?;
let d_sigma_d = stream.alloc_zeros::<f32>(input.n_horizons)?;
let n_pos_i = input.n_pos as i32;
let n_horizons_i = input.n_horizons as i32;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
};
let mut launch = stream.launch_builder(&func);
launch
.arg(&probs_d).arg(&labels_d).arg(&weights_d).arg(&sigma_d)
.arg(&n_pos_i).arg(&n_horizons_i)
.arg(&mut loss_d).arg(&mut mean_bce_h_d).arg(&mut grad_d).arg(&mut valid_d)
.arg(&mut d_sigma_d);
unsafe { launch.launch(cfg).context("bce launch")?; }
{
let rs = stream.cu_stream();
let mut args = RawArgs::new();
args.push_ptr(probs_d.raw_ptr());
args.push_ptr(labels_d.raw_ptr());
args.push_ptr(weights_d.raw_ptr());
args.push_ptr(sigma_d.raw_ptr());
args.push_i32(n_pos_i);
args.push_i32(n_horizons_i);
args.push_ptr(loss_d.raw_ptr());
args.push_ptr(mean_bce_h_d.raw_ptr());
args.push_ptr(grad_d.raw_ptr());
args.push_ptr(valid_d.raw_ptr());
args.push_ptr(d_sigma_d.raw_ptr());
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
func.cu_function(), (1,1,1), (256,1,1),
0, rs, &mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("bce: {:?}", e))?;
}
}
unsafe {
raw_stream_sync(stream.cu_stream())