perf(rl): replace cuBLAS with capture-safe matmul kernels — unblocks mega-graph
cuBLAS GEMM calls inside CUDA-graph capture regions cause CUDA_ERROR_STREAM_CAPTURE_INVALIDATED on first use of new (m,n,k) shapes. The mega-graph was silently failing → fast-path replay was a NO-OP → sps stayed at ~10. Phase 1 — DQN distributional Q head: - New crates/ml-alpha/cuda/dqn_q_head_fwd_bwd.cu (3 kernels) - Removed cuBLAS field + gemm_f32 helper from dqn.rs Phase 2 — IQN ensemble heads: - New crates/ml-alpha/cuda/rl_iqn_matmul.cu (3 kernels) - Removed cuBLAS field from iqn.rs Phase 3 — Mamba2 SKIPPED: workspace pre-warms during 65+ eager warmup steps before mega-graph capture. Local smoke (RTX 3050 Ti, b=128, 500 steps): - Mega-graph captures cleanly at step 67 - 8.9 sps avg, 11 sps peak post-capture (GPU-bound on mobile GPU) - l_q=0.024 (healthy), l_pi rising, V converging - Production L40S should see full mega-graph speedup now Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -32,6 +32,7 @@ const KERNELS: &[&str] = &[
|
||||
"aux_loss", // SDD-3 Layer B4: Huber loss + grad for aux trade-outcome regression targets (NaN-masked)
|
||||
"aux_vec_add", // SDD-3 Layer B5: element-wise dst += src for aux→encoder gradient accumulation (lifted stop-grad)
|
||||
"dqn_distributional_q", // RL Phase C: C51 distributional Q-head fwd + Bellman TD bwd for integrated RL trainer
|
||||
"dqn_q_head_fwd_bwd", // Mega-graph capture: hand-written matmul fwd + grad_h_t + grad_w replacing cuBLAS SGEMM in DqnHead::forward_gemm / backward_gemm (cuBLAS internal allocs broke graph capture)
|
||||
"rl_gamma_controller", // RL Phase C: ISV controller emitting γ to ISV[RL_GAMMA_INDEX=400]
|
||||
"rl_target_tau_controller", // RL Phase C: ISV controller emitting τ to ISV[RL_TARGET_TAU_INDEX=401]
|
||||
"ppo_clipped_surrogate", // RL Phase D: PPO clipped-surrogate + entropy bonus + value MSE fwd/bwd
|
||||
@@ -95,6 +96,7 @@ const KERNELS: &[&str] = &[
|
||||
"rl_asymmetric_trail_decay", // auto-tighten losers, auto-widen winners — structural P&L asymmetry
|
||||
"rl_session_risk_check", // session-level loss limit circuit breaker
|
||||
"rl_iqn_forward", // IQN distributional Q-head: quantile embedding + action-value projection; complementary to C51
|
||||
"rl_iqn_matmul", // Mega-graph capture: hand-written matmul fwd (embed + out) + bwd (grad_combined) replacing cuBLAS SGEMM in IqnHead::forward_inner / backward (cuBLAS internal allocs broke graph capture)
|
||||
"rl_iqn_loss", // IQN quantile Huber loss: ρ_τ(δ) = |τ - 1(δ<0)| × Huber(δ, κ=1.0); forward + backward
|
||||
"rl_iqn_backward", // IQN backward through forward pass: grad_output → grad_w_out/b_out/w_embed/b_embed per-batch scratch
|
||||
"rl_ensemble_action_value", // C51+IQN ensemble: E_ensemble = α×E_C51 + (1-α)×E_IQN; α from ISV[544]
|
||||
|
||||
248
crates/ml-alpha/cuda/dqn_q_head_fwd_bwd.cu
Normal file
248
crates/ml-alpha/cuda/dqn_q_head_fwd_bwd.cu
Normal file
@@ -0,0 +1,248 @@
|
||||
// dqn_q_head_fwd_bwd.cu — capture-safe matmul kernels for the C51
|
||||
// distributional Q head.
|
||||
//
|
||||
// Replaces the cuBLAS SGEMM path used by `DqnHead::forward_gemm` and
|
||||
// `DqnHead::backward_gemm`. cuBLAS calls perform internal allocations
|
||||
// and HtoD memcpys on the first launch with a new shape, which breaks
|
||||
// CUDA-graph stream capture with `CUDA_ERROR_STREAM_CAPTURE_INVALIDATED`.
|
||||
// The hand-written kernels in this file have NO host work and NO
|
||||
// internal allocations, so the entire training step can be captured.
|
||||
//
|
||||
// Shapes (constants — match `crates/ml-alpha/src/rl/common.rs` and
|
||||
// `crates/ml-alpha/src/heads/mod.rs`):
|
||||
// HIDDEN_DIM = 128
|
||||
// N_ACTIONS = 11
|
||||
// Q_N_ATOMS = 21
|
||||
// N_OUT = N_ACTIONS * Q_N_ATOMS = 231
|
||||
// B = batch size (runtime, currently 1024)
|
||||
//
|
||||
// Layout (matches existing `dqn_distributional_q.cu` and the cuBLAS
|
||||
// path's row-major view of W):
|
||||
// W [N_OUT, HIDDEN_DIM] row-major — w[k * HIDDEN_DIM + c]
|
||||
// b_bias [N_OUT]
|
||||
// h_t [B, HIDDEN_DIM] row-major
|
||||
// logits [B, N_OUT] row-major — fused softmax over atoms
|
||||
// happens in the loss kernel
|
||||
//
|
||||
// Kernels:
|
||||
// dqn_q_head_fwd : logits = h_t @ W^T + b_bias (B × N_OUT)
|
||||
// dqn_q_head_grad_h_t : grad_h_t = grad_logits @ W (B × HIDDEN_DIM)
|
||||
// dqn_q_head_grad_w : grad_w = grad_logits^T @ h_t (N_OUT × HIDDEN_DIM)
|
||||
//
|
||||
// grad_b is produced by the existing `reduce_sum_axis0_kernel` on
|
||||
// `grad_logits [B × N_OUT]` — that helper is already capture-safe.
|
||||
//
|
||||
// Constraints honoured:
|
||||
// * `feedback_no_atomicadd.md` — no atomicAdd
|
||||
// * `feedback_no_nvrtc.md` — pre-compiled cubin via build.rs
|
||||
// * `feedback_nvidia_grade_perf_for_kernels.md` — warp-uniform, no
|
||||
// divergent shuffles, no host branches; tree-reductions only
|
||||
// * `pearl_no_host_branches_in_captured_graph` — no host scalars
|
||||
// affect control flow (only B as a captured int).
|
||||
|
||||
#define HIDDEN_DIM 128
|
||||
#define N_ACTIONS 11
|
||||
#define Q_N_ATOMS 21
|
||||
#define N_OUT (N_ACTIONS * Q_N_ATOMS) // 231
|
||||
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// dqn_q_head_fwd: forward matmul + bias add.
|
||||
//
|
||||
// logits[b, k] = b_bias[k] + Σ_c W[k, c] * h_t[b, c]
|
||||
//
|
||||
// Launch:
|
||||
// grid = (B, 1, 1)
|
||||
// block = (BLOCK_FWD = 256, 1, 1)
|
||||
// shared_mem_bytes = HIDDEN_DIM * sizeof(float) (h_t row stage)
|
||||
//
|
||||
// Each block handles one batch sample. h_t[b, :] is staged once into
|
||||
// shared memory; each thread then computes one or more output slots k.
|
||||
// At B=1024 the grid covers every CTA on L40S (sm_89, 142 SMs × ~6
|
||||
// concurrent CTAs); the inner stride loop covers k = tid, tid+256,
|
||||
// ... up to N_OUT=231 (so most blocks issue exactly one iteration).
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
#define BLOCK_FWD 256
|
||||
|
||||
extern "C" __global__ void dqn_q_head_fwd(
|
||||
const float* __restrict__ w, // [N_OUT * HIDDEN_DIM]
|
||||
const float* __restrict__ b_bias, // [N_OUT]
|
||||
const float* __restrict__ h_t, // [B * HIDDEN_DIM]
|
||||
int B,
|
||||
float* __restrict__ logits // [B * N_OUT]
|
||||
) {
|
||||
extern __shared__ float s_h[]; // [HIDDEN_DIM]
|
||||
|
||||
const int b = blockIdx.x;
|
||||
const int tid = threadIdx.x;
|
||||
if (b >= B) return;
|
||||
|
||||
// Cooperative stage of h_t[b, :] into shared memory.
|
||||
for (int c = tid; c < HIDDEN_DIM; c += BLOCK_FWD) {
|
||||
s_h[c] = h_t[b * HIDDEN_DIM + c];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Each thread handles one or more output slots.
|
||||
for (int k = tid; k < N_OUT; k += BLOCK_FWD) {
|
||||
float acc = b_bias[k];
|
||||
const float* w_row = w + k * HIDDEN_DIM;
|
||||
#pragma unroll 8
|
||||
for (int c = 0; c < HIDDEN_DIM; ++c) {
|
||||
acc += w_row[c] * s_h[c];
|
||||
}
|
||||
logits[b * N_OUT + k] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// dqn_q_head_grad_h_t: backward into encoder hidden state.
|
||||
//
|
||||
// grad_h_t[b, c] = Σ_k grad_logits[b, k] * W[k, c]
|
||||
//
|
||||
// Launch:
|
||||
// grid = (B, 1, 1)
|
||||
// block = (HIDDEN_DIM = 128, 1, 1)
|
||||
// shared_mem_bytes = N_OUT * sizeof(float) (grad_logits row stage)
|
||||
//
|
||||
// Each block handles one batch sample. grad_logits[b, :] is staged
|
||||
// once into shared memory; thread c then reduces W[:, c] × s_gl
|
||||
// over k = 0..N_OUT.
|
||||
//
|
||||
// Mirrors the grad_h_t portion of the legacy `dqn_grad_w_b_h_t` kernel
|
||||
// in `dqn_distributional_q.cu` (which also stages grad_logits[b, :] in
|
||||
// shared); the W column traversal is identical.
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
extern "C" __global__ void dqn_q_head_grad_h_t(
|
||||
const float* __restrict__ w, // [N_OUT * HIDDEN_DIM]
|
||||
const float* __restrict__ grad_logits, // [B * N_OUT]
|
||||
int B,
|
||||
float* __restrict__ grad_h_t // [B * HIDDEN_DIM] (OVERWRITE)
|
||||
) {
|
||||
extern __shared__ float s_gl[]; // [N_OUT]
|
||||
|
||||
const int b = blockIdx.x;
|
||||
const int c = threadIdx.x;
|
||||
if (b >= B) return;
|
||||
if (c >= HIDDEN_DIM) return;
|
||||
|
||||
// Cooperative stage of grad_logits[b, :] into shared. Stride-loop
|
||||
// over N_OUT in steps of HIDDEN_DIM (block size). N_OUT=231 with
|
||||
// 128 threads → each thread loads ⌈231/128⌉ = 2 elements.
|
||||
for (int k = c; k < N_OUT; k += HIDDEN_DIM) {
|
||||
s_gl[k] = grad_logits[b * N_OUT + k];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// Thread c accumulates grad_h_t[b, c] = Σ_k s_gl[k] * w[k, c].
|
||||
float acc = 0.0f;
|
||||
#pragma unroll 8
|
||||
for (int k = 0; k < N_OUT; ++k) {
|
||||
acc += s_gl[k] * w[k * HIDDEN_DIM + c];
|
||||
}
|
||||
grad_h_t[b * HIDDEN_DIM + c] = acc;
|
||||
}
|
||||
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// dqn_q_head_grad_w: backward into weight matrix.
|
||||
//
|
||||
// grad_w[k, c] = Σ_b grad_logits[b, k] * h_t[b, c]
|
||||
//
|
||||
// Tiled matmul: each block computes a TM × TN output tile of grad_w
|
||||
// by sweeping over B in chunks of TB.
|
||||
//
|
||||
// Launch:
|
||||
// grid = (ceil(N_OUT / TM), ceil(HIDDEN_DIM / TN), 1)
|
||||
// block = (TN, TM, 1)
|
||||
// shared_mem_bytes = (TB * TM + TB * TN) * sizeof(float)
|
||||
// = (32*16 + 32*16) * 4 = 4 KiB
|
||||
//
|
||||
// At B=1024, N_OUT=231, HIDDEN_DIM=128, TM=TN=16, TB=32:
|
||||
// grid = (15, 8) = 120 blocks
|
||||
// block = (16, 16) = 256 threads
|
||||
// shared = 4 KiB / block
|
||||
// B-reduction loop iters = 1024 / 32 = 32
|
||||
//
|
||||
// No atomicAdd — each block is the sole writer of its TM × TN tile.
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
#define TM 16
|
||||
#define TN 16
|
||||
#define TB 32
|
||||
|
||||
extern "C" __global__ void dqn_q_head_grad_w(
|
||||
const float* __restrict__ grad_logits, // [B * N_OUT]
|
||||
const float* __restrict__ h_t, // [B * HIDDEN_DIM]
|
||||
int B,
|
||||
float* __restrict__ grad_w // [N_OUT * HIDDEN_DIM]
|
||||
) {
|
||||
// Shared-mem tiles. Layout chosen for bank-conflict-free access:
|
||||
// s_gl[bi][ki] → thread (ty, ?) reads s_gl[bi][ty] → bank = ty
|
||||
// s_ht[bi][ci] → thread (?, tx) reads s_ht[bi][tx] → bank = tx
|
||||
__shared__ float s_gl[TB][TM]; // grad_logits tile [TB rows × TM cols]
|
||||
__shared__ float s_ht[TB][TN]; // h_t tile [TB rows × TN cols]
|
||||
|
||||
const int k_base = blockIdx.x * TM;
|
||||
const int c_base = blockIdx.y * TN;
|
||||
const int tx = threadIdx.x; // c offset [0..TN)
|
||||
const int ty = threadIdx.y; // k offset [0..TM)
|
||||
const int k = k_base + ty;
|
||||
const int c = c_base + tx;
|
||||
|
||||
// Linear thread id for cooperative tile loads.
|
||||
const int tid = ty * TN + tx; // [0..TM*TN) = [0..256)
|
||||
const int BS = TM * TN; // 256
|
||||
|
||||
float acc = 0.0f;
|
||||
|
||||
for (int b_base = 0; b_base < B; b_base += TB) {
|
||||
// Load grad_logits tile [TB × TM] = 512 elements via 256 threads
|
||||
// (2 elements/thread). Boundary-guard with zero so the inner
|
||||
// multiply is safe for k >= N_OUT or b >= B padding lanes.
|
||||
#pragma unroll
|
||||
for (int idx = tid; idx < TB * TM; idx += BS) {
|
||||
const int bi = idx / TM;
|
||||
const int ki = idx % TM;
|
||||
const int b_g = b_base + bi;
|
||||
const int k_g = k_base + ki;
|
||||
float v = 0.0f;
|
||||
if (b_g < B && k_g < N_OUT) {
|
||||
v = grad_logits[b_g * N_OUT + k_g];
|
||||
}
|
||||
s_gl[bi][ki] = v;
|
||||
}
|
||||
|
||||
// Load h_t tile [TB × TN] = 512 elements via 256 threads.
|
||||
// c is always < HIDDEN_DIM at grid level (HIDDEN_DIM % TN == 0),
|
||||
// but the boundary guard keeps the kernel robust for future
|
||||
// shape changes.
|
||||
#pragma unroll
|
||||
for (int idx = tid; idx < TB * TN; idx += BS) {
|
||||
const int bi = idx / TN;
|
||||
const int ci = idx % TN;
|
||||
const int b_g = b_base + bi;
|
||||
const int c_g = c_base + ci;
|
||||
float v = 0.0f;
|
||||
if (b_g < B && c_g < HIDDEN_DIM) {
|
||||
v = h_t[b_g * HIDDEN_DIM + c_g];
|
||||
}
|
||||
s_ht[bi][ci] = v;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
// Each thread accumulates its grad_w[k, c] over the TB chunk.
|
||||
#pragma unroll
|
||||
for (int bi = 0; bi < TB; ++bi) {
|
||||
acc += s_gl[bi][ty] * s_ht[bi][tx];
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
// Boundary-guarded store: only threads owning a valid (k, c) write.
|
||||
if (k < N_OUT && c < HIDDEN_DIM) {
|
||||
grad_w[k * HIDDEN_DIM + c] = acc;
|
||||
}
|
||||
}
|
||||
273
crates/ml-alpha/cuda/rl_iqn_matmul.cu
Normal file
273
crates/ml-alpha/cuda/rl_iqn_matmul.cu
Normal file
@@ -0,0 +1,273 @@
|
||||
// rl_iqn_matmul.cu — capture-safe matmul kernels for the IQN head.
|
||||
//
|
||||
// Replaces the cuBLAS SGEMM calls in `IqnHead::forward_inner` and
|
||||
// `IqnHead::backward`. cuBLAS performs internal workspace allocations
|
||||
// and HtoD memcpys on the first launch with a new shape, which breaks
|
||||
// CUDA-graph stream capture with `CUDA_ERROR_STREAM_CAPTURE_INVALIDATED`.
|
||||
// The hand-written kernels below have NO host work and NO internal
|
||||
// allocations, so the entire mega-graph captures cleanly.
|
||||
//
|
||||
// Shapes (constants — match `crates/ml-alpha/src/rl/iqn.rs`):
|
||||
// HIDDEN_DIM = 128
|
||||
// N_ACTIONS = 11
|
||||
// EMBED_DIM = 64
|
||||
// M = B * N_TAU (e.g. 128 * 32 = 4096)
|
||||
//
|
||||
// Kernels:
|
||||
// iqn_embed_matmul_fwd : embed_out[M, HIDDEN_DIM] = cos_features[M, EMBED_DIM] @ W_embed[EMBED_DIM, HIDDEN_DIM]
|
||||
// iqn_out_matmul_fwd : q_raw [M, N_ACTIONS] = combined [M, HIDDEN_DIM] @ W_out [HIDDEN_DIM, N_ACTIONS]
|
||||
// iqn_grad_combined_bwd : grad_comb[M, HIDDEN_DIM] = grad_q [M, N_ACTIONS] @ W_out^T[N_ACTIONS, HIDDEN_DIM]
|
||||
//
|
||||
// All weight layouts are row-major and match what cuBLAS was reading
|
||||
// (see comments in `iqn.rs::forward_inner` / `backward`).
|
||||
//
|
||||
// Constraints honoured:
|
||||
// * `feedback_no_atomicadd.md` — no atomicAdd
|
||||
// * `feedback_no_nvrtc.md` — pre-compiled cubin via build.rs
|
||||
// * `feedback_nvidia_grade_perf_for_kernels.md` — tiled with shared
|
||||
// memory staging; cooperative loads; warp-uniform.
|
||||
|
||||
#define HIDDEN_DIM 128
|
||||
#define N_ACTIONS 11
|
||||
#define EMBED_DIM 64
|
||||
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// iqn_embed_matmul_fwd:
|
||||
// embed_out[m, c] = Σ_e cos_features[m, e] * W_embed[e, c]
|
||||
//
|
||||
// Inputs:
|
||||
// cos_features [M * EMBED_DIM] row-major
|
||||
// w_embed [EMBED_DIM * HIDDEN_DIM] row-major
|
||||
// Output:
|
||||
// embed_out [M * HIDDEN_DIM] row-major
|
||||
//
|
||||
// Tiling: each block computes a TM × TN tile of embed_out by sweeping
|
||||
// over EMBED_DIM in chunks of TK.
|
||||
//
|
||||
// At M=4096, HIDDEN_DIM=128, EMBED_DIM=64, TM=32, TN=32, TK=16:
|
||||
// grid = (ceil(M/TM), HIDDEN_DIM/TN) = (128, 4) = 512 blocks
|
||||
// block = (TN, TM) = (32, 32) = 1024 threads (max per block, ok)
|
||||
//
|
||||
// Shared mem: (TM*TK + TK*TN) * 4 B = (32*16 + 16*32) * 4 = 4 KiB
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
#define EMBED_FWD_TM 32
|
||||
#define EMBED_FWD_TN 32
|
||||
#define EMBED_FWD_TK 16
|
||||
|
||||
extern "C" __global__ void iqn_embed_matmul_fwd(
|
||||
const float* __restrict__ cos_features, // [M * EMBED_DIM]
|
||||
const float* __restrict__ w_embed, // [EMBED_DIM * HIDDEN_DIM]
|
||||
int M,
|
||||
float* __restrict__ embed_out // [M * HIDDEN_DIM]
|
||||
) {
|
||||
__shared__ float s_a[EMBED_FWD_TM][EMBED_FWD_TK]; // cos_features tile
|
||||
__shared__ float s_b[EMBED_FWD_TK][EMBED_FWD_TN]; // w_embed tile
|
||||
|
||||
const int m_base = blockIdx.x * EMBED_FWD_TM;
|
||||
const int c_base = blockIdx.y * EMBED_FWD_TN;
|
||||
const int tx = threadIdx.x; // c offset [0..TN)
|
||||
const int ty = threadIdx.y; // m offset [0..TM)
|
||||
const int m = m_base + ty;
|
||||
const int c = c_base + tx;
|
||||
|
||||
const int tid = ty * EMBED_FWD_TN + tx;
|
||||
const int BS = EMBED_FWD_TM * EMBED_FWD_TN; // 1024
|
||||
|
||||
float acc = 0.0f;
|
||||
|
||||
for (int e_base = 0; e_base < EMBED_DIM; e_base += EMBED_FWD_TK) {
|
||||
// Load cos_features tile [TM × TK] = 512 elements with 1024 threads.
|
||||
#pragma unroll
|
||||
for (int idx = tid; idx < EMBED_FWD_TM * EMBED_FWD_TK; idx += BS) {
|
||||
const int mi = idx / EMBED_FWD_TK;
|
||||
const int ei = idx % EMBED_FWD_TK;
|
||||
const int m_g = m_base + mi;
|
||||
const int e_g = e_base + ei;
|
||||
float v = 0.0f;
|
||||
if (m_g < M && e_g < EMBED_DIM) {
|
||||
v = cos_features[m_g * EMBED_DIM + e_g];
|
||||
}
|
||||
s_a[mi][ei] = v;
|
||||
}
|
||||
// Load w_embed tile [TK × TN] = 512 elements.
|
||||
#pragma unroll
|
||||
for (int idx = tid; idx < EMBED_FWD_TK * EMBED_FWD_TN; idx += BS) {
|
||||
const int ei = idx / EMBED_FWD_TN;
|
||||
const int ci = idx % EMBED_FWD_TN;
|
||||
const int e_g = e_base + ei;
|
||||
const int c_g = c_base + ci;
|
||||
float v = 0.0f;
|
||||
if (e_g < EMBED_DIM && c_g < HIDDEN_DIM) {
|
||||
v = w_embed[e_g * HIDDEN_DIM + c_g];
|
||||
}
|
||||
s_b[ei][ci] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
#pragma unroll
|
||||
for (int ei = 0; ei < EMBED_FWD_TK; ++ei) {
|
||||
acc += s_a[ty][ei] * s_b[ei][tx];
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
if (m < M && c < HIDDEN_DIM) {
|
||||
embed_out[m * HIDDEN_DIM + c] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// iqn_out_matmul_fwd:
|
||||
// q_raw[m, a] = Σ_c combined[m, c] * W_out[c, a]
|
||||
//
|
||||
// Inputs:
|
||||
// combined [M * HIDDEN_DIM] row-major
|
||||
// w_out [HIDDEN_DIM * N_ACTIONS] row-major
|
||||
// Output:
|
||||
// q_raw [M * N_ACTIONS] row-major
|
||||
//
|
||||
// N_ACTIONS=11 is tiny, so we tile only along M and use a thread-per-
|
||||
// (m, a) layout. Each block handles a single (M_TILE × N_ACTIONS) slab,
|
||||
// stages W_out once into shared memory, and sweeps over HIDDEN_DIM via
|
||||
// chunked staging of `combined`.
|
||||
//
|
||||
// Block layout:
|
||||
// grid = (ceil(M / OUT_TM), 1, 1)
|
||||
// block = (N_ACTIONS, OUT_TM, 1) = (11, 32) = 352 threads
|
||||
// shared = HIDDEN_DIM * N_ACTIONS + OUT_TM * HIDDEN_DIM
|
||||
// = 128*11 + 32*128 = 1408 + 4096 = 5504 floats = 22 KiB
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
#define OUT_FWD_TM 32
|
||||
|
||||
extern "C" __global__ void iqn_out_matmul_fwd(
|
||||
const float* __restrict__ combined, // [M * HIDDEN_DIM]
|
||||
const float* __restrict__ w_out, // [HIDDEN_DIM * N_ACTIONS]
|
||||
int M,
|
||||
float* __restrict__ q_raw // [M * N_ACTIONS]
|
||||
) {
|
||||
__shared__ float s_w[HIDDEN_DIM][N_ACTIONS]; // W_out staged once
|
||||
__shared__ float s_c[OUT_FWD_TM][HIDDEN_DIM]; // combined tile
|
||||
|
||||
const int m_base = blockIdx.x * OUT_FWD_TM;
|
||||
const int ax = threadIdx.x; // a offset [0..N_ACTIONS)
|
||||
const int my = threadIdx.y; // m offset [0..OUT_FWD_TM)
|
||||
const int m = m_base + my;
|
||||
|
||||
const int tid = my * N_ACTIONS + ax;
|
||||
const int BS = OUT_FWD_TM * N_ACTIONS; // 32 * 11 = 352
|
||||
|
||||
// Cooperative load of W_out [HIDDEN_DIM × N_ACTIONS] = 1408 floats
|
||||
// via 352 threads → 4 elements/thread.
|
||||
#pragma unroll
|
||||
for (int idx = tid; idx < HIDDEN_DIM * N_ACTIONS; idx += BS) {
|
||||
const int ci = idx / N_ACTIONS;
|
||||
const int ai = idx % N_ACTIONS;
|
||||
s_w[ci][ai] = w_out[ci * N_ACTIONS + ai];
|
||||
}
|
||||
|
||||
// Cooperative load of combined tile [OUT_FWD_TM × HIDDEN_DIM] = 4096
|
||||
// floats via 352 threads → ~12 elements/thread.
|
||||
#pragma unroll
|
||||
for (int idx = tid; idx < OUT_FWD_TM * HIDDEN_DIM; idx += BS) {
|
||||
const int mi = idx / HIDDEN_DIM;
|
||||
const int ci = idx % HIDDEN_DIM;
|
||||
const int m_g = m_base + mi;
|
||||
float v = 0.0f;
|
||||
if (m_g < M) {
|
||||
v = combined[m_g * HIDDEN_DIM + ci];
|
||||
}
|
||||
s_c[mi][ci] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (m < M && ax < N_ACTIONS) {
|
||||
float acc = 0.0f;
|
||||
#pragma unroll 8
|
||||
for (int c = 0; c < HIDDEN_DIM; ++c) {
|
||||
acc += s_c[my][c] * s_w[c][ax];
|
||||
}
|
||||
q_raw[m * N_ACTIONS + ax] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
// iqn_grad_combined_bwd:
|
||||
// grad_combined[m, c] = Σ_a grad_q[m, a] * W_out[c, a]
|
||||
//
|
||||
// ≡ Row-major matmul `grad_combined = grad_q @ W_out^T` where
|
||||
// W_out is stored row-major as [HIDDEN_DIM, N_ACTIONS].
|
||||
//
|
||||
// Inputs:
|
||||
// grad_q [M * N_ACTIONS] row-major
|
||||
// w_out [HIDDEN_DIM * N_ACTIONS] row-major
|
||||
// Output:
|
||||
// grad_combined [M * HIDDEN_DIM] row-major
|
||||
//
|
||||
// Block layout mirrors `iqn_out_matmul_fwd`:
|
||||
// grid = (ceil(M / GC_TM), HIDDEN_DIM / GC_TN, 1)
|
||||
// block = (GC_TN, GC_TM)
|
||||
// With N_ACTIONS=11 as the (small) reduction dim, the inner loop is
|
||||
// fully unrolled — no need for chunked TK staging.
|
||||
// ─────────────────────────────────────────────────────────────────────
|
||||
#define GC_TM 32
|
||||
#define GC_TN 32
|
||||
|
||||
extern "C" __global__ void iqn_grad_combined_bwd(
|
||||
const float* __restrict__ grad_q, // [M * N_ACTIONS]
|
||||
const float* __restrict__ w_out, // [HIDDEN_DIM * N_ACTIONS]
|
||||
int M,
|
||||
float* __restrict__ grad_combined // [M * HIDDEN_DIM]
|
||||
) {
|
||||
__shared__ float s_gq[GC_TM][N_ACTIONS]; // grad_q tile
|
||||
__shared__ float s_w [GC_TN][N_ACTIONS]; // W_out tile
|
||||
|
||||
const int m_base = blockIdx.x * GC_TM;
|
||||
const int c_base = blockIdx.y * GC_TN;
|
||||
const int tx = threadIdx.x;
|
||||
const int ty = threadIdx.y;
|
||||
const int m = m_base + ty;
|
||||
const int c = c_base + tx;
|
||||
|
||||
const int tid = ty * GC_TN + tx;
|
||||
const int BS = GC_TM * GC_TN; // 1024
|
||||
|
||||
// Load grad_q tile [GC_TM × N_ACTIONS] = 32*11 = 352 floats.
|
||||
#pragma unroll
|
||||
for (int idx = tid; idx < GC_TM * N_ACTIONS; idx += BS) {
|
||||
const int mi = idx / N_ACTIONS;
|
||||
const int ai = idx % N_ACTIONS;
|
||||
const int m_g = m_base + mi;
|
||||
float v = 0.0f;
|
||||
if (m_g < M) {
|
||||
v = grad_q[m_g * N_ACTIONS + ai];
|
||||
}
|
||||
s_gq[mi][ai] = v;
|
||||
}
|
||||
|
||||
// Load W_out tile [GC_TN × N_ACTIONS] = 32*11 = 352 floats —
|
||||
// row `c` of W_out is `[w_out[c, 0], w_out[c, 1], ..., w_out[c, N_ACTIONS-1]]`.
|
||||
#pragma unroll
|
||||
for (int idx = tid; idx < GC_TN * N_ACTIONS; idx += BS) {
|
||||
const int ci = idx / N_ACTIONS;
|
||||
const int ai = idx % N_ACTIONS;
|
||||
const int c_g = c_base + ci;
|
||||
float v = 0.0f;
|
||||
if (c_g < HIDDEN_DIM) {
|
||||
v = w_out[c_g * N_ACTIONS + ai];
|
||||
}
|
||||
s_w[ci][ai] = v;
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (m < M && c < HIDDEN_DIM) {
|
||||
float acc = 0.0f;
|
||||
#pragma unroll
|
||||
for (int a = 0; a < N_ACTIONS; ++a) {
|
||||
acc += s_gq[ty][a] * s_w[tx][a];
|
||||
}
|
||||
grad_combined[m * HIDDEN_DIM + c] = acc;
|
||||
}
|
||||
}
|
||||
@@ -41,9 +41,6 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
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};
|
||||
use cudarc::driver::sys::CUstream;
|
||||
use ml_core::cuda_autograd::init::scoped_init_seed;
|
||||
@@ -60,59 +57,23 @@ 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;
|
||||
|
||||
// ── cuBLAS SGEMM helper (mirrors ml-core::cuda_autograd::linear::gemm_ex_f32) ──
|
||||
|
||||
/// F32 x F32 -> F32 GEMM via `cublasGemmEx` with F32 internal accumulation.
|
||||
///
|
||||
/// # Safety
|
||||
/// All device pointers must be valid and dimensions must be correct.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
unsafe fn gemm_f32(
|
||||
cublas: &CudaBlas,
|
||||
transa: cublasOperation_t,
|
||||
transb: cublasOperation_t,
|
||||
m: i32,
|
||||
n: i32,
|
||||
k: i32,
|
||||
alpha: f32,
|
||||
a_ptr: u64,
|
||||
lda: i32,
|
||||
b_ptr: u64,
|
||||
ldb: i32,
|
||||
beta: f32,
|
||||
c_ptr: u64,
|
||||
ldc: i32,
|
||||
label: &str,
|
||||
) -> Result<()> {
|
||||
cudarc::cublas::result::gemm_ex(
|
||||
*cublas.handle(),
|
||||
transa,
|
||||
transb,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
(&alpha as *const f32).cast(),
|
||||
a_ptr as *const std::ffi::c_void,
|
||||
cublas_sys::cudaDataType_t::CUDA_R_32F,
|
||||
lda,
|
||||
b_ptr as *const std::ffi::c_void,
|
||||
cublas_sys::cudaDataType_t::CUDA_R_32F,
|
||||
ldb,
|
||||
(&beta as *const f32).cast(),
|
||||
c_ptr as *mut std::ffi::c_void,
|
||||
cublas_sys::cudaDataType_t::CUDA_R_32F,
|
||||
ldc,
|
||||
cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
|
||||
cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DFALT,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("cublasGemmEx {label}: {e:?}"))?;
|
||||
Ok(())
|
||||
}
|
||||
// ── Tile constants for `dqn_q_head_grad_w` (mirror the kernel file) ──
|
||||
//
|
||||
// Each block computes a TM × TN output tile of grad_w by sweeping over
|
||||
// the batch dimension in chunks of TB. Shared-mem usage per block:
|
||||
// (TB * TM + TB * TN) * sizeof(float) = 4 KiB at TM=TN=16, TB=32.
|
||||
const GRAD_W_TM: u32 = 16;
|
||||
const GRAD_W_TN: u32 = 16;
|
||||
const GRAD_W_TB: u32 = 32;
|
||||
|
||||
const DQN_HEAD_CUBIN: &[u8] = include_bytes!(concat!(
|
||||
env!("OUT_DIR"),
|
||||
"/dqn_distributional_q.cubin"
|
||||
));
|
||||
const DQN_Q_HEAD_FWD_BWD_CUBIN: &[u8] = include_bytes!(concat!(
|
||||
env!("OUT_DIR"),
|
||||
"/dqn_q_head_fwd_bwd.cubin"
|
||||
));
|
||||
const DQN_TARGET_SOFT_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(
|
||||
env!("OUT_DIR"),
|
||||
"/dqn_target_soft_update.cubin"
|
||||
@@ -213,16 +174,21 @@ pub struct DqnHead {
|
||||
/// Target-network biases, same shape as `b_d`.
|
||||
pub b_target_d: CudaSlice<f32>,
|
||||
|
||||
// ── cuBLAS infrastructure ────────────────────────────────────────
|
||||
/// cuBLAS handle for SGEMM-based forward / backward passes.
|
||||
/// Replaces the hand-written `dqn_distributional_q_fwd` and
|
||||
/// `dqn_grad_w_b_h_t` kernels: cuBLAS tensor-core SGEMM is 5-10x
|
||||
/// faster at b=256 for the [B, 128] x [128, 231] matrix shapes.
|
||||
pub cublas: CudaBlas,
|
||||
/// Pre-allocated cuBLAS workspace (8 MiB). Prevents per-call
|
||||
/// cudaMalloc inside cuBLAS that would break CUDA Graph capture.
|
||||
_cublas_workspace: CudaSlice<u8>,
|
||||
/// Bias-add and reduce-sum-axis0 kernels from ml-core.
|
||||
// ── Mega-graph capture-safe matmul kernels ──────────────────────
|
||||
// Replace cuBLAS SGEMM (which performs internal allocations / HtoD
|
||||
// on first launch with a new shape, breaking CUDA-graph stream
|
||||
// capture with `CUDA_ERROR_STREAM_CAPTURE_INVALIDATED`).
|
||||
/// `dqn_q_head_fwd` — logits = h_t @ W^T + b_bias.
|
||||
pub q_fwd_fn: CudaFunction,
|
||||
/// `dqn_q_head_grad_h_t` — grad_h_t = grad_logits @ W.
|
||||
pub q_grad_h_t_fn: CudaFunction,
|
||||
/// `dqn_q_head_grad_w` — grad_w = grad_logits^T @ h_t.
|
||||
pub q_grad_w_fn: CudaFunction,
|
||||
/// Owns the `dqn_q_head_fwd_bwd` cubin lifetime (all three matmul
|
||||
/// kernels live in the same translation unit).
|
||||
_q_head_fwd_bwd_module: Arc<CudaModule>,
|
||||
/// Reduce-sum-axis0 kernel handle from ml-core — used to fold
|
||||
/// `grad_logits [B × K_OUT]` into `grad_b [K_OUT]`. Capture-safe.
|
||||
bias_kernels: BiasKernels,
|
||||
}
|
||||
|
||||
@@ -293,33 +259,25 @@ impl DqnHead {
|
||||
let w_target_d = upload(&stream, &w_host)?;
|
||||
let b_target_d = upload(&stream, &b_host)?;
|
||||
|
||||
// cuBLAS handle + pre-allocated workspace (same pattern as
|
||||
// Mamba2Block::new). 8 MiB workspace prevents per-call
|
||||
// cudaMalloc that would break CUDA Graph stream capture.
|
||||
let cublas = CudaBlas::new(Arc::clone(&stream))
|
||||
.context("DqnHead: cuBLAS init")?;
|
||||
const CUBLAS_WORKSPACE_BYTES: usize = 8 * 1024 * 1024;
|
||||
let cublas_workspace = stream
|
||||
.alloc_zeros::<u8>(CUBLAS_WORKSPACE_BYTES)
|
||||
.context("DqnHead: cuBLAS workspace alloc")?;
|
||||
unsafe {
|
||||
let ws_ptr = cublas_workspace.raw_ptr();
|
||||
cudarc::cublas::sys::cublasSetWorkspace_v2(
|
||||
*cublas.handle(),
|
||||
ws_ptr as *mut std::ffi::c_void,
|
||||
CUBLAS_WORKSPACE_BYTES,
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow::anyhow!("DqnHead: cublasSetWorkspace_v2: {e:?}"))?;
|
||||
cudarc::cublas::sys::cublasSetMathMode(
|
||||
*cublas.handle(),
|
||||
cudarc::cublas::sys::cublasMath_t::CUBLAS_TF32_TENSOR_OP_MATH,
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow::anyhow!("DqnHead: cublasSetMathMode TF32: {e:?}"))?;
|
||||
}
|
||||
// Mega-graph capture-safe matmul cubin: replaces cuBLAS SGEMM
|
||||
// in `forward_gemm` / `backward_gemm`. cuBLAS performs internal
|
||||
// allocs / HtoD on first-use-of-new-shape which breaks CUDA
|
||||
// graph stream capture with CUDA_ERROR_STREAM_CAPTURE_INVALIDATED.
|
||||
let q_head_fwd_bwd_module = ctx
|
||||
.load_cubin(DQN_Q_HEAD_FWD_BWD_CUBIN.to_vec())
|
||||
.context("load dqn_q_head_fwd_bwd cubin")?;
|
||||
let q_fwd_fn = q_head_fwd_bwd_module
|
||||
.load_function("dqn_q_head_fwd")
|
||||
.context("load dqn_q_head_fwd")?;
|
||||
let q_grad_h_t_fn = q_head_fwd_bwd_module
|
||||
.load_function("dqn_q_head_grad_h_t")
|
||||
.context("load dqn_q_head_grad_h_t")?;
|
||||
let q_grad_w_fn = q_head_fwd_bwd_module
|
||||
.load_function("dqn_q_head_grad_w")
|
||||
.context("load dqn_q_head_grad_w")?;
|
||||
|
||||
// Bias-add and reduce-sum-axis0 kernel handles from ml-core.
|
||||
// reduce_sum_axis0 kernel handle from ml-core — folds
|
||||
// grad_logits[B × K_OUT] into grad_b[K_OUT].
|
||||
let bias_kernels = BiasKernels::shared(&stream)
|
||||
.map_err(|e| anyhow::anyhow!("DqnHead: BiasKernels init: {e}"))?;
|
||||
|
||||
@@ -343,8 +301,10 @@ impl DqnHead {
|
||||
b_d,
|
||||
w_target_d,
|
||||
b_target_d,
|
||||
cublas,
|
||||
_cublas_workspace: cublas_workspace,
|
||||
q_fwd_fn,
|
||||
q_grad_h_t_fn,
|
||||
q_grad_w_fn,
|
||||
_q_head_fwd_bwd_module: q_head_fwd_bwd_module,
|
||||
bias_kernels,
|
||||
})
|
||||
}
|
||||
@@ -401,12 +361,15 @@ impl DqnHead {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Phase E.2 forward via cuBLAS SGEMM:
|
||||
/// Forward via the capture-safe `dqn_q_head_fwd` kernel:
|
||||
/// `logits[B, K_OUT] = h_t[B, HIDDEN_DIM] @ W^T[HIDDEN_DIM, K_OUT] + bias[K_OUT]`
|
||||
///
|
||||
/// Replaces the hand-written `dqn_distributional_q_fwd` kernel whose
|
||||
/// per-thread serial loop over HIDDEN_DIM=128 consumed 18.1% of GPU
|
||||
/// time. cuBLAS tensor-core SGEMM is 5-10x faster at b=256.
|
||||
/// Replaces the previous cuBLAS SGEMM path. cuBLAS performs internal
|
||||
/// workspace allocs / HtoD memcpys on the first launch with a new
|
||||
/// shape, which breaks CUDA-graph stream capture with
|
||||
/// `CUDA_ERROR_STREAM_CAPTURE_INVALIDATED`. The hand-written kernel
|
||||
/// fuses bias-add into the matmul and has no host work, so the
|
||||
/// entire mega-graph can capture cleanly.
|
||||
///
|
||||
/// The trainer pre-allocates `logits_out` so the same memory is
|
||||
/// reused across step() calls.
|
||||
@@ -422,12 +385,18 @@ impl DqnHead {
|
||||
.context("dqn forward (online)")
|
||||
}
|
||||
|
||||
/// cuBLAS SGEMM forward shared by online and target network paths.
|
||||
/// Capture-safe matmul-plus-bias forward shared by online and target
|
||||
/// network paths.
|
||||
///
|
||||
/// Computes `logits[B, K_OUT] = h_t[B, HIDDEN_DIM] @ W^T + bias`:
|
||||
/// - cuBLAS column-major: C_col[K_OUT, B] = W_col^T[K_OUT, HIDDEN_DIM] @ X_col[HIDDEN_DIM, B]
|
||||
/// transA=T, transB=N, m=K_OUT, n=B, k=HIDDEN_DIM
|
||||
/// - Then add_bias_2d kernel broadcasts bias[K_OUT] into logits[B, K_OUT].
|
||||
/// Launches `dqn_q_head_fwd` from `cuda/dqn_q_head_fwd_bwd.cu`:
|
||||
/// * Grid = (B, 1, 1) — one block per batch sample.
|
||||
/// * Block = (256, 1, 1).
|
||||
/// * Shared mem = `HIDDEN_DIM * sizeof(float)` — staged `h_t[b, :]`.
|
||||
///
|
||||
/// Each thread handles one or more (action, atom) output slots,
|
||||
/// reducing W[k, :] · h_t[b, :] and adding the bias. Bias is fused
|
||||
/// into the same kernel so no separate `add_bias_2d` launch is
|
||||
/// needed.
|
||||
fn forward_gemm(
|
||||
&self,
|
||||
w: &CudaSlice<f32>,
|
||||
@@ -436,59 +405,22 @@ impl DqnHead {
|
||||
b_size: usize,
|
||||
logits_out: &mut CudaSlice<f32>,
|
||||
) -> Result<()> {
|
||||
let w_ptr = w.raw_ptr();
|
||||
let h_ptr = h_t.raw_ptr();
|
||||
let out_ptr = logits_out.raw_ptr();
|
||||
// SGEMM: logits = h_t @ W^T (alpha=1, beta=0 overwrites logits_out)
|
||||
unsafe {
|
||||
gemm_f32(
|
||||
&self.cublas,
|
||||
cublasOperation_t::CUBLAS_OP_T, // transA: W[K_OUT, HIDDEN_DIM] stored row-major = col[HIDDEN_DIM, K_OUT]
|
||||
cublasOperation_t::CUBLAS_OP_N, // transB: h_t[B, HIDDEN_DIM] stored row-major = col[HIDDEN_DIM, B]
|
||||
K_OUT as i32, // m
|
||||
b_size as i32, // n
|
||||
HIDDEN_DIM as i32, // k
|
||||
1.0, // alpha
|
||||
w_ptr, // A = W
|
||||
HIDDEN_DIM as i32, // lda (W row-major [K_OUT, HIDDEN_DIM] → col lda = HIDDEN_DIM)
|
||||
h_ptr, // B = h_t
|
||||
HIDDEN_DIM as i32, // ldb
|
||||
0.0, // beta
|
||||
out_ptr, // C = logits_out
|
||||
K_OUT as i32, // ldc
|
||||
"dqn_fwd_sgemm",
|
||||
)?;
|
||||
}
|
||||
// Broadcast bias: logits_out[b, j] += bias[j]
|
||||
self.add_bias(logits_out, b, b_size)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch the `add_bias_2d_kernel` from ml-core's BiasKernels.
|
||||
fn add_bias(
|
||||
&self,
|
||||
y: &mut CudaSlice<f32>,
|
||||
bias: &CudaSlice<f32>,
|
||||
rows: usize,
|
||||
) -> Result<()> {
|
||||
let total = rows * K_OUT;
|
||||
let threads = 256_u32;
|
||||
let blocks = ((total as u32) + threads - 1) / threads;
|
||||
let rows_i32 = rows as i32;
|
||||
let cols_i32 = K_OUT as i32;
|
||||
let b_i = b_size as i32;
|
||||
let smem = (HIDDEN_DIM * std::mem::size_of::<f32>()) as u32;
|
||||
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);
|
||||
args.push_ptr(w.raw_ptr());
|
||||
args.push_ptr(b.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.bias_kernels.add_fn.cu_function(),
|
||||
(blocks, 1, 1), (threads, 1, 1), 0,
|
||||
self.q_fwd_fn.cu_function(),
|
||||
(b_size as u32, 1, 1), (256, 1, 1), smem,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
).map_err(|e| anyhow::anyhow!("dqn add_bias_2d: {:?}", e))?;
|
||||
).map_err(|e| anyhow::anyhow!("dqn_q_head_fwd: {:?}", e))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -591,21 +523,24 @@ impl DqnHead {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// cuBLAS backward: compute grad_w, grad_b, and grad_h_t from
|
||||
/// `grad_logits` in three SGEMM + one reduce-sum launch. Replaces
|
||||
/// the hand-written `dqn_grad_w_b_h_t` kernel + `reduce_axis0`.
|
||||
/// Capture-safe backward: compute grad_w, grad_b, and grad_h_t from
|
||||
/// `grad_logits` in two hand-written matmul launches + one
|
||||
/// reduce-sum launch. Replaces the previous cuBLAS SGEMM path.
|
||||
///
|
||||
/// cuBLAS performs internal workspace allocs / HtoD memcpys on the
|
||||
/// first launch with a new shape, which breaks CUDA-graph stream
|
||||
/// capture with `CUDA_ERROR_STREAM_CAPTURE_INVALIDATED`. The
|
||||
/// `dqn_q_head_grad_h_t` and `dqn_q_head_grad_w` kernels in
|
||||
/// `cuda/dqn_q_head_fwd_bwd.cu` have no host work and no internal
|
||||
/// allocations, so the mega-graph captures cleanly.
|
||||
///
|
||||
/// Mathematics:
|
||||
/// grad_h_t [B, HIDDEN_DIM] = grad_logits [B, K_OUT] @ W [K_OUT, HIDDEN_DIM]
|
||||
/// grad_h_t [B, HIDDEN_DIM] = grad_logits [B, K_OUT] @ W [K_OUT, HIDDEN_DIM]
|
||||
/// grad_w [K_OUT, HIDDEN_DIM] = grad_logits^T [K_OUT, B] @ h_t [B, HIDDEN_DIM]
|
||||
/// grad_b [K_OUT] = sum(grad_logits [B, K_OUT], axis=0)
|
||||
/// grad_b [K_OUT] = sum(grad_logits [B, K_OUT], axis=0)
|
||||
///
|
||||
/// Benefits:
|
||||
/// - Eliminates the B*K_OUT*HIDDEN_DIM per-batch scratch (7.5 MiB at B=256).
|
||||
/// - Eliminates two `reduce_axis0` kernel launches.
|
||||
/// - cuBLAS SGEMM is 5-10x faster than the hand-written kernel for
|
||||
/// these shapes ([B=256, 128] x [128, 231]).
|
||||
/// - grad_w / grad_b are produced directly reduced, ready for Adam.
|
||||
/// grad_w / grad_b are produced directly reduced (no per-batch
|
||||
/// scratch + reduce_axis0 round-trip), ready for Adam.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn backward_gemm(
|
||||
&self,
|
||||
@@ -622,76 +557,78 @@ impl DqnHead {
|
||||
debug_assert_eq!(grad_b.len(), K_OUT);
|
||||
debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM);
|
||||
|
||||
let gl_ptr = grad_logits.raw_ptr();
|
||||
let ht_ptr = h_t.raw_ptr();
|
||||
let w_ptr = self.w_d.raw_ptr();
|
||||
let b_i = b_size as i32;
|
||||
|
||||
// ── grad_h_t [B, HIDDEN_DIM] = grad_logits [B, K_OUT] @ W [K_OUT, HIDDEN_DIM] ──
|
||||
// Column-major: grad_h_t_col[HIDDEN_DIM, B] = W_col[HIDDEN_DIM, K_OUT] @ gl_col[K_OUT, B]
|
||||
// transA=N, transB=N, m=HIDDEN_DIM, n=B, k=K_OUT
|
||||
let gh_ptr = grad_h_t.raw_ptr();
|
||||
unsafe {
|
||||
gemm_f32(
|
||||
&self.cublas,
|
||||
cublasOperation_t::CUBLAS_OP_N, // W[K_OUT, HIDDEN_DIM] row-major = col[HIDDEN_DIM, K_OUT]
|
||||
cublasOperation_t::CUBLAS_OP_N, // gl[B, K_OUT] row-major = col[K_OUT, B]
|
||||
HIDDEN_DIM as i32, // m
|
||||
b_size as i32, // n
|
||||
K_OUT as i32, // k
|
||||
1.0, // alpha
|
||||
w_ptr, // A = W
|
||||
HIDDEN_DIM as i32, // lda
|
||||
gl_ptr, // B = grad_logits
|
||||
K_OUT as i32, // ldb
|
||||
0.0, // beta
|
||||
gh_ptr, // C = grad_h_t
|
||||
HIDDEN_DIM as i32, // ldc
|
||||
"dqn_bwd_grad_h_t",
|
||||
)?;
|
||||
// Launch: grid=(B,), block=(HIDDEN_DIM=128,), smem = N_OUT * 4 B.
|
||||
// Each block stages grad_logits[b, :] in shared; thread c then
|
||||
// accumulates Σ_k s_gl[k] * W[k, c].
|
||||
{
|
||||
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(grad_logits.raw_ptr());
|
||||
args.push_i32(b_i);
|
||||
args.push_ptr(grad_h_t.raw_ptr());
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
raw_launch(
|
||||
self.q_grad_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_q_head_grad_h_t: {:?}", e))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── grad_w [K_OUT, HIDDEN_DIM] = grad_logits^T [K_OUT, B] @ h_t [B, HIDDEN_DIM] ──
|
||||
// Column-major: grad_w_col[HIDDEN_DIM, K_OUT] = h_t_col[HIDDEN_DIM, B] @ gl_col[K_OUT, B]^T
|
||||
// transA=N, transB=T, m=HIDDEN_DIM, n=K_OUT, k=B
|
||||
let gw_ptr = grad_w.raw_ptr();
|
||||
unsafe {
|
||||
gemm_f32(
|
||||
&self.cublas,
|
||||
cublasOperation_t::CUBLAS_OP_N, // h_t[B, HIDDEN_DIM] row-major = col[HIDDEN_DIM, B]
|
||||
cublasOperation_t::CUBLAS_OP_T, // gl[B, K_OUT] row-major = col[K_OUT, B]; transposed
|
||||
HIDDEN_DIM as i32, // m
|
||||
K_OUT as i32, // n
|
||||
b_size as i32, // k
|
||||
1.0, // alpha
|
||||
ht_ptr, // A = h_t
|
||||
HIDDEN_DIM as i32, // lda
|
||||
gl_ptr, // B = grad_logits
|
||||
K_OUT as i32, // ldb
|
||||
0.0, // beta
|
||||
gw_ptr, // C = grad_w
|
||||
HIDDEN_DIM as i32, // ldc
|
||||
"dqn_bwd_grad_w",
|
||||
)?;
|
||||
// Tiled matmul. Each block computes a TM × TN output tile by
|
||||
// sweeping over B in chunks of TB. Grid = (⌈K_OUT/TM⌉,
|
||||
// ⌈HIDDEN_DIM/TN⌉), Block = (TN, TM). Shared mem = (TB*TM +
|
||||
// TB*TN) * 4 B = 4 KiB at TM=TN=16, TB=32.
|
||||
{
|
||||
let grid_x = (K_OUT as u32).div_ceil(GRAD_W_TM);
|
||||
let grid_y = (HIDDEN_DIM as u32).div_ceil(GRAD_W_TN);
|
||||
let smem = ((GRAD_W_TB * GRAD_W_TM + GRAD_W_TB * GRAD_W_TN)
|
||||
* std::mem::size_of::<f32>() as u32) as u32;
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(grad_logits.raw_ptr());
|
||||
args.push_ptr(h_t.raw_ptr());
|
||||
args.push_i32(b_i);
|
||||
args.push_ptr(grad_w.raw_ptr());
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
raw_launch(
|
||||
self.q_grad_w_fn.cu_function(),
|
||||
(grid_x, grid_y, 1), (GRAD_W_TN, GRAD_W_TM, 1), smem,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
).map_err(|e| anyhow::anyhow!("dqn_q_head_grad_w: {:?}", e))?;
|
||||
}
|
||||
}
|
||||
|
||||
// ── grad_b [K_OUT] = sum(grad_logits [B, K_OUT], axis=0) ──────
|
||||
let threads = 256_u32;
|
||||
let blocks = ((K_OUT as u32) + threads - 1) / threads;
|
||||
let rows_i32 = b_size as i32;
|
||||
let cols_i32 = K_OUT as i32;
|
||||
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 {
|
||||
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))?;
|
||||
// Reuses the existing capture-safe ml-core reduce_sum_axis0
|
||||
// kernel — no cuBLAS involvement.
|
||||
{
|
||||
let threads = 256_u32;
|
||||
let blocks = (K_OUT as u32).div_ceil(threads);
|
||||
let rows_i32 = b_size as i32;
|
||||
let cols_i32 = K_OUT as i32;
|
||||
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 {
|
||||
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(())
|
||||
}
|
||||
|
||||
@@ -12,40 +12,48 @@
|
||||
//!
|
||||
//! where α is read from `ISV[RL_IQN_ENSEMBLE_ALPHA_INDEX=544]`.
|
||||
//!
|
||||
//! ## Architecture (cuBLAS-accelerated pipeline)
|
||||
//! ## Architecture (capture-safe custom-kernel pipeline)
|
||||
//!
|
||||
//! The forward pass is split into custom kernels + cuBLAS SGEMMs:
|
||||
//! The forward pass is a chain of hand-written CUDA kernels — no cuBLAS:
|
||||
//!
|
||||
//! ```text
|
||||
//! 1. rl_iqn_tau_cos_features (custom):
|
||||
//! 1. rl_iqn_tau_cos_features:
|
||||
//! Sample tau ~ U(0,1), compute cos_features[j] = cos((j+1)*pi*tau)
|
||||
//! → [B*N_TAU, EMBED_DIM]
|
||||
//!
|
||||
//! 2. cuBLAS SGEMM:
|
||||
//! 2. iqn_embed_matmul_fwd (cuda/rl_iqn_matmul.cu):
|
||||
//! embed_out = cos_features @ W_embed → [B*N_TAU, HIDDEN_DIM]
|
||||
//!
|
||||
//! 3. rl_iqn_relu_hadamard (custom):
|
||||
//! 3. rl_iqn_relu_hadamard:
|
||||
//! combined = h_t ⊙ ReLU(embed_out + b_embed) → [B*N_TAU, HIDDEN_DIM]
|
||||
//!
|
||||
//! 4. cuBLAS SGEMM:
|
||||
//! 4. iqn_out_matmul_fwd (cuda/rl_iqn_matmul.cu):
|
||||
//! q_raw = combined @ W_out → [B*N_TAU, N_ACTIONS]
|
||||
//!
|
||||
//! 5. rl_iqn_bias_add_q (custom):
|
||||
//! 5. rl_iqn_bias_add_q:
|
||||
//! q_values = q_raw + b_out → [B*N_TAU, N_ACTIONS]
|
||||
//! ```
|
||||
//!
|
||||
//! The backward pass uses one cuBLAS SGEMM for the largest matmul:
|
||||
//! The backward pass:
|
||||
//!
|
||||
//! ```text
|
||||
//! 1. cuBLAS SGEMM:
|
||||
//! 1. iqn_grad_combined_bwd (cuda/rl_iqn_matmul.cu):
|
||||
//! grad_combined = grad_q @ W_out^T → [B*N_TAU, HIDDEN_DIM]
|
||||
//!
|
||||
//! 2. rl_iqn_backward (custom):
|
||||
//! 2. rl_iqn_backward:
|
||||
//! Per-batch backward accumulation. Recomputes phi(tau) and combined
|
||||
//! inline from (tau, W_embed, b_embed, h_t). Reads grad_combined
|
||||
//! from cuBLAS. Produces per-batch grad_w_out/b_out/w_embed/b_embed.
|
||||
//! from the matmul kernel above. Produces per-batch
|
||||
//! grad_w_out/b_out/w_embed/b_embed.
|
||||
//! ```
|
||||
//!
|
||||
//! cuBLAS is intentionally avoided: it performs internal workspace
|
||||
//! allocations / HtoD memcpys on the first launch with a new shape,
|
||||
//! which breaks CUDA-graph stream capture with
|
||||
//! `CUDA_ERROR_STREAM_CAPTURE_INVALIDATED`. The hand-written matmul
|
||||
//! kernels in `cuda/rl_iqn_matmul.cu` have no host work and no internal
|
||||
//! allocations, so the mega-graph captures cleanly.
|
||||
//!
|
||||
//! ## Loss
|
||||
//!
|
||||
//! Quantile Huber loss (Dabney et al. 2018):
|
||||
@@ -71,14 +79,12 @@
|
||||
//! * `feedback_no_nvrtc.md` — pre-compiled cubins via `build.rs`.
|
||||
//! * `feedback_isv_for_adaptive_bounds.md` — N_TAU, ensemble α, and LR
|
||||
//! live in ISV slots 543-545.
|
||||
//! * `feedback_cpu_is_read_only.md` — forward/backward are GPU kernels
|
||||
//! + cuBLAS SGEMMs.
|
||||
//! * `feedback_cpu_is_read_only.md` — forward/backward are pure GPU
|
||||
//! kernels (no cuBLAS, no host work).
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use cudarc::cublas::CudaBlas;
|
||||
use cudarc::cublas::sys::{self as cublas_sys, cublasOperation_t};
|
||||
use cudarc::driver::{
|
||||
CudaFunction, CudaModule, CudaSlice, CudaStream,
|
||||
};
|
||||
@@ -101,14 +107,25 @@ pub const EMBED_DIM: usize = 64;
|
||||
/// via `RL_IQN_N_TAU_INDEX`. Matches the `N_TAU` default in the kernel.
|
||||
const DEFAULT_N_TAU: usize = 32;
|
||||
|
||||
/// cuBLAS workspace size — 8 MiB, prevents internal cudaMalloc during
|
||||
/// CUDA Graph stream capture.
|
||||
const CUBLAS_WORKSPACE_BYTES: usize = 8 * 1024 * 1024;
|
||||
// ── Tile constants for `iqn_embed_matmul_fwd` (mirror the kernel) ──
|
||||
const EMBED_FWD_TM: u32 = 32;
|
||||
const EMBED_FWD_TN: u32 = 32;
|
||||
|
||||
// ── Tile constants for `iqn_out_matmul_fwd` (mirror the kernel) ──
|
||||
const OUT_FWD_TM: u32 = 32;
|
||||
|
||||
// ── Tile constants for `iqn_grad_combined_bwd` (mirror the kernel) ──
|
||||
const GC_TM: u32 = 32;
|
||||
const GC_TN: u32 = 32;
|
||||
|
||||
const IQN_FWD_CUBIN: &[u8] = include_bytes!(concat!(
|
||||
env!("OUT_DIR"),
|
||||
"/rl_iqn_forward.cubin"
|
||||
));
|
||||
const IQN_MATMUL_CUBIN: &[u8] = include_bytes!(concat!(
|
||||
env!("OUT_DIR"),
|
||||
"/rl_iqn_matmul.cubin"
|
||||
));
|
||||
const IQN_LOSS_CUBIN: &[u8] = include_bytes!(concat!(
|
||||
env!("OUT_DIR"),
|
||||
"/rl_iqn_loss.cubin"
|
||||
@@ -146,57 +163,7 @@ impl Default for IqnHeadConfig {
|
||||
}
|
||||
}
|
||||
|
||||
// ── raw pointer helpers (same pattern as ml-core linear.rs) ──────────
|
||||
|
||||
/// F32 × F32 → F32 GEMM via `cublasGemmEx` with F32 internal accumulation.
|
||||
///
|
||||
/// Uses `CUBLAS_GEMM_DFALT` algo — deterministic, safe for CUDA Graph capture.
|
||||
///
|
||||
/// # Safety
|
||||
/// All device pointers must be valid and dimensions must be correct.
|
||||
unsafe fn gemm_ex_f32(
|
||||
cublas: &CudaBlas,
|
||||
transa: cublasOperation_t,
|
||||
transb: cublasOperation_t,
|
||||
m: i32,
|
||||
n: i32,
|
||||
k: i32,
|
||||
a_ptr: u64,
|
||||
lda: i32,
|
||||
b_ptr: u64,
|
||||
ldb: i32,
|
||||
c_ptr: u64,
|
||||
ldc: i32,
|
||||
label: &str,
|
||||
) -> Result<()> {
|
||||
let alpha = 1.0_f32;
|
||||
let beta = 0.0_f32;
|
||||
cudarc::cublas::result::gemm_ex(
|
||||
*cublas.handle(),
|
||||
transa,
|
||||
transb,
|
||||
m,
|
||||
n,
|
||||
k,
|
||||
(&alpha as *const f32).cast(),
|
||||
a_ptr as *const std::ffi::c_void,
|
||||
cublas_sys::cudaDataType_t::CUDA_R_32F,
|
||||
lda,
|
||||
b_ptr as *const std::ffi::c_void,
|
||||
cublas_sys::cudaDataType_t::CUDA_R_32F,
|
||||
ldb,
|
||||
(&beta as *const f32).cast(),
|
||||
c_ptr as *mut std::ffi::c_void,
|
||||
cublas_sys::cudaDataType_t::CUDA_R_32F,
|
||||
ldc,
|
||||
cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
|
||||
cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DFALT,
|
||||
)
|
||||
.map_err(|e| anyhow::anyhow!("cublasGemmEx IQN {label}: {e:?}"))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// IQN distributional Q-head with cuBLAS-accelerated matmuls.
|
||||
/// IQN distributional Q-head with capture-safe custom-kernel matmuls.
|
||||
///
|
||||
/// Owns device weights for the quantile embedding (`w_embed`, `b_embed`)
|
||||
/// and output projection (`w_out`, `b_out`), plus a target network copy
|
||||
@@ -206,9 +173,17 @@ pub struct IqnHead {
|
||||
stream: Arc<CudaStream>,
|
||||
raw_stream: CUstream,
|
||||
|
||||
// ── cuBLAS ───────────────────────────────────────────────────────
|
||||
cublas: CudaBlas,
|
||||
_cublas_workspace: CudaSlice<u8>,
|
||||
// ── Mega-graph capture-safe matmul kernels ──────────────────────
|
||||
// Replace cuBLAS SGEMM (which performs internal workspace allocs
|
||||
// / HtoD on first launch with a new shape, breaking CUDA-graph
|
||||
// stream capture with CUDA_ERROR_STREAM_CAPTURE_INVALIDATED).
|
||||
_matmul_module: Arc<CudaModule>,
|
||||
/// Forward stage 2: `embed_out = cos_features @ W_embed`.
|
||||
pub embed_matmul_fwd_fn: CudaFunction,
|
||||
/// Forward stage 4: `q_raw = combined @ W_out`.
|
||||
pub out_matmul_fwd_fn: CudaFunction,
|
||||
/// Backward stage 1: `grad_combined = grad_q @ W_out^T`.
|
||||
pub grad_combined_bwd_fn: CudaFunction,
|
||||
|
||||
// ── Forward kernels (split pipeline) ─────────────────────────────
|
||||
_fwd_module: Arc<CudaModule>,
|
||||
@@ -216,7 +191,7 @@ pub struct IqnHead {
|
||||
pub tau_cos_features_fn: CudaFunction,
|
||||
/// Stage 3: bias-add → ReLU → hadamard product with h_t.
|
||||
pub relu_hadamard_fn: CudaFunction,
|
||||
/// Stage 5: bias addition to cuBLAS output projection result.
|
||||
/// Stage 5: bias addition to output projection result.
|
||||
pub bias_add_q_fn: CudaFunction,
|
||||
/// Expected-Q reduction: mean over tau dimension.
|
||||
pub expected_q_fn: CudaFunction,
|
||||
@@ -257,22 +232,22 @@ impl IqnHead {
|
||||
let stream: Arc<CudaStream> = dev.cuda_stream().context("iqn_head stream")?.clone();
|
||||
let ctx = dev.cuda_context().context("iqn_head ctx")?;
|
||||
|
||||
// ── cuBLAS handle + pre-allocated workspace ──────────────────
|
||||
let cublas = CudaBlas::new(Arc::clone(&stream))
|
||||
.map_err(|e| anyhow::anyhow!("IqnHead: cuBLAS init: {e}"))?;
|
||||
let cublas_workspace = stream
|
||||
.alloc_zeros::<u8>(CUBLAS_WORKSPACE_BYTES)
|
||||
.map_err(|e| anyhow::anyhow!("IqnHead: cuBLAS workspace: {e}"))?;
|
||||
unsafe {
|
||||
let ws_ptr = cublas_workspace.raw_ptr();
|
||||
cudarc::cublas::sys::cublasSetWorkspace_v2(
|
||||
*cublas.handle(),
|
||||
ws_ptr as *mut std::ffi::c_void,
|
||||
CUBLAS_WORKSPACE_BYTES,
|
||||
)
|
||||
.result()
|
||||
.map_err(|e| anyhow::anyhow!("IqnHead: cublasSetWorkspace_v2: {e:?}"))?;
|
||||
}
|
||||
// ── Mega-graph capture-safe matmul cubin ─────────────────────
|
||||
// Replaces cuBLAS SGEMM. cuBLAS performs internal workspace
|
||||
// allocs / HtoD on first-use-of-new-shape which breaks CUDA
|
||||
// graph stream capture with CUDA_ERROR_STREAM_CAPTURE_INVALIDATED.
|
||||
let matmul_module = ctx
|
||||
.load_cubin(IQN_MATMUL_CUBIN.to_vec())
|
||||
.context("load rl_iqn_matmul cubin")?;
|
||||
let embed_matmul_fwd_fn = matmul_module
|
||||
.load_function("iqn_embed_matmul_fwd")
|
||||
.context("load iqn_embed_matmul_fwd")?;
|
||||
let out_matmul_fwd_fn = matmul_module
|
||||
.load_function("iqn_out_matmul_fwd")
|
||||
.context("load iqn_out_matmul_fwd")?;
|
||||
let grad_combined_bwd_fn = matmul_module
|
||||
.load_function("iqn_grad_combined_bwd")
|
||||
.context("load iqn_grad_combined_bwd")?;
|
||||
|
||||
// ── Forward kernel symbols ───────────────────────────────────
|
||||
let fwd_module = ctx
|
||||
@@ -356,8 +331,10 @@ impl IqnHead {
|
||||
cfg,
|
||||
stream,
|
||||
raw_stream,
|
||||
cublas,
|
||||
_cublas_workspace: cublas_workspace,
|
||||
_matmul_module: matmul_module,
|
||||
embed_matmul_fwd_fn,
|
||||
out_matmul_fwd_fn,
|
||||
grad_combined_bwd_fn,
|
||||
_fwd_module: fwd_module,
|
||||
tau_cos_features_fn,
|
||||
relu_hadamard_fn,
|
||||
@@ -479,20 +456,26 @@ impl IqnHead {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Stage 2: cuBLAS SGEMM — embed_out = cos_features @ W_embed ──
|
||||
// ── Stage 2: capture-safe matmul — embed_out = cos_features @ W_embed ──
|
||||
// Tiled matmul: grid = (⌈m/TM⌉, hd/TN); block = (TN, TM) = (32, 32).
|
||||
{
|
||||
let w_ptr = w_embed.raw_ptr();
|
||||
let grid_x = (m as u32).div_ceil(EMBED_FWD_TM);
|
||||
let grid_y = (hd as u32).div_ceil(EMBED_FWD_TN);
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(cos_ptr);
|
||||
args.push_ptr(w_embed.raw_ptr());
|
||||
args.push_i32(m as i32);
|
||||
args.push_ptr(embed_ptr);
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
gemm_ex_f32(
|
||||
&self.cublas,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
hd as i32, m as i32, EMBED_DIM as i32,
|
||||
w_ptr, hd as i32,
|
||||
cos_ptr, EMBED_DIM as i32,
|
||||
embed_ptr, hd as i32,
|
||||
"embed_fwd",
|
||||
)?;
|
||||
raw_launch(
|
||||
self.embed_matmul_fwd_fn.cu_function(),
|
||||
(grid_x, grid_y, 1),
|
||||
(EMBED_FWD_TN, EMBED_FWD_TM, 1),
|
||||
0,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
).map_err(|e| anyhow::anyhow!("iqn_embed_matmul_fwd: {:?}", e))?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -522,21 +505,26 @@ impl IqnHead {
|
||||
}
|
||||
}
|
||||
|
||||
// ── Stage 4: cuBLAS SGEMM — q_raw = combined @ W_out ────────
|
||||
// ── Stage 4: capture-safe matmul — q_raw = combined @ W_out ──
|
||||
// Grid = (⌈m/OUT_FWD_TM⌉, 1); block = (N_ACTIONS, OUT_FWD_TM) = (11, 32).
|
||||
// W_out is staged once into shared memory per block.
|
||||
{
|
||||
let wo_ptr = w_out.raw_ptr();
|
||||
let q_ptr = q_values_out.raw_ptr();
|
||||
let grid_x = (m as u32).div_ceil(OUT_FWD_TM);
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(comb_ptr);
|
||||
args.push_ptr(w_out.raw_ptr());
|
||||
args.push_i32(m as i32);
|
||||
args.push_ptr(q_values_out.raw_ptr());
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
gemm_ex_f32(
|
||||
&self.cublas,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
N_ACTIONS as i32, m as i32, hd as i32,
|
||||
wo_ptr, N_ACTIONS as i32,
|
||||
comb_ptr, hd as i32,
|
||||
q_ptr, N_ACTIONS as i32,
|
||||
"out_proj_fwd",
|
||||
)?;
|
||||
raw_launch(
|
||||
self.out_matmul_fwd_fn.cu_function(),
|
||||
(grid_x, 1, 1),
|
||||
(N_ACTIONS as u32, OUT_FWD_TM, 1),
|
||||
0,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
).map_err(|e| anyhow::anyhow!("iqn_out_matmul_fwd: {:?}", e))?;
|
||||
}
|
||||
}
|
||||
|
||||
@@ -670,29 +658,29 @@ impl IqnHead {
|
||||
let m = b_size * n_tau;
|
||||
let hd = self.cfg.hidden_dim;
|
||||
|
||||
// ── cuBLAS: grad_combined = grad_q @ W_out^T ─────────────────
|
||||
// Row-major: C[M, hd] = A[M, N_ACTIONS] @ B^T[N_ACTIONS, hd]
|
||||
// cuBLAS col-major: C_col[hd, M] = W_out_col^T[hd, N_ACTIONS] @ grad_q_col[N_ACTIONS, M]
|
||||
// W_out is row-major [hd, N_ACTIONS] = col-major [N_ACTIONS, hd].
|
||||
// transA=T on col-major [N_ACTIONS, hd] → [hd, N_ACTIONS].
|
||||
// transB=N on grad_q col-major [N_ACTIONS, M].
|
||||
// m=hd, n=M, k=N_ACTIONS, lda=N_ACTIONS, ldb=N_ACTIONS, ldc=hd
|
||||
// ── Capture-safe matmul: grad_combined = grad_q @ W_out^T ────
|
||||
// Tiled: grid = (⌈m/TM⌉, hd/TN); block = (GC_TN, GC_TM) = (32, 32).
|
||||
// Each block computes a TM × TN tile of grad_combined.
|
||||
debug_assert!(m * hd <= self.scratch_grad_combined.len());
|
||||
let gc_ptr = self.scratch_grad_combined.raw_ptr();
|
||||
{
|
||||
let wo_ptr = self.w_out_d.raw_ptr();
|
||||
let gq_ptr = grad_output.raw_ptr();
|
||||
let grid_x = (m as u32).div_ceil(GC_TM);
|
||||
let grid_y = (hd as u32).div_ceil(GC_TN);
|
||||
let mut args = RawArgs::new();
|
||||
args.push_ptr(grad_output.raw_ptr());
|
||||
args.push_ptr(self.w_out_d.raw_ptr());
|
||||
args.push_i32(m as i32);
|
||||
args.push_ptr(gc_ptr);
|
||||
let mut ptrs = args.build_arg_ptrs();
|
||||
unsafe {
|
||||
gemm_ex_f32(
|
||||
&self.cublas,
|
||||
cublasOperation_t::CUBLAS_OP_T,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
hd as i32, m as i32, N_ACTIONS as i32,
|
||||
wo_ptr, N_ACTIONS as i32,
|
||||
gq_ptr, N_ACTIONS as i32,
|
||||
gc_ptr, hd as i32,
|
||||
"grad_combined_bwd",
|
||||
)?;
|
||||
raw_launch(
|
||||
self.grad_combined_bwd_fn.cu_function(),
|
||||
(grid_x, grid_y, 1),
|
||||
(GC_TN, GC_TM, 1),
|
||||
0,
|
||||
self.raw_stream,
|
||||
&mut ptrs[..args.len()],
|
||||
).map_err(|e| anyhow::anyhow!("iqn_grad_combined_bwd: {:?}", e))?;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -211,7 +211,10 @@ fn frd_softmax_ce_grad_uniform_logits_match_log_n_atoms() -> Result<()> {
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
|
||||
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
|
||||
let loss = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
|
||||
let expected = (FRD_N_ATOMS as f32).ln();
|
||||
@@ -266,7 +269,10 @@ fn frd_softmax_ce_grad_sentinel_label_zeros_row() -> Result<()> {
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
|
||||
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
|
||||
let loss = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
|
||||
let grad = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?;
|
||||
@@ -306,7 +312,10 @@ fn frd_softmax_ce_grad_finite_diff_matches_analytical() -> Result<()> {
|
||||
let logits_d = upload_f32(&stream, &logits)?;
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
let grad_analytical = read_slice_d_pub(&stream, &grad_d, b_size * FRD_OUT_DIM)?;
|
||||
|
||||
// Finite-difference for slot (b=0, h=0, a=3). Note: gradient was
|
||||
@@ -320,14 +329,20 @@ fn frd_softmax_ce_grad_finite_diff_matches_analytical() -> Result<()> {
|
||||
// L(logits + ε · e_j) — perturb only the target slot upward.
|
||||
logits[probe_off] += eps;
|
||||
let logits_plus_d = upload_f32(&stream, &logits)?;
|
||||
head.softmax_ce_grad(&logits_plus_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_plus_d, &labels_d, &mut grad_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
let loss_plus = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
|
||||
let l_plus = loss_plus[probe_h]; // only h=0 affected — h=1,2 share the perturbation only if probe was in their horizon block
|
||||
|
||||
// L(logits - ε · e_j)
|
||||
logits[probe_off] -= 2.0 * eps;
|
||||
let logits_minus_d = upload_f32(&stream, &logits)?;
|
||||
head.softmax_ce_grad(&logits_minus_d, &labels_d, &mut grad_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_minus_d, &labels_d, &mut grad_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
let loss_minus = read_slice_d_pub(&stream, &loss_d, b_size * FRD_N_HORIZONS)?;
|
||||
let l_minus = loss_minus[probe_h];
|
||||
|
||||
@@ -365,7 +380,10 @@ fn ce_total_loss(
|
||||
let logits_d = upload_f32(stream, logits)?;
|
||||
let mut grad_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
|
||||
head.softmax_ce_grad(&logits_d, labels_d, &mut grad_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_d, labels_d, &mut grad_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
let loss = read_slice_d_pub(stream, &loss_d, b_size * FRD_N_HORIZONS)?;
|
||||
Ok(loss.iter().sum())
|
||||
}
|
||||
@@ -395,7 +413,10 @@ fn frd_layer2_bwd_finite_diff_w2() -> Result<()> {
|
||||
// Softmax+CE grad of logits.
|
||||
let mut grad_logits_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
|
||||
// Layer-2 backward: produce per-batch grad_W2 scratch.
|
||||
let mut grad_w2_pb_d =
|
||||
@@ -491,7 +512,10 @@ fn frd_layer2_bwd_db2_equals_grad_logits() -> Result<()> {
|
||||
|
||||
let mut grad_logits_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
|
||||
let mut grad_w2_pb_d =
|
||||
stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?;
|
||||
@@ -543,7 +567,10 @@ fn frd_layer1_bwd_finite_diff_w1() -> Result<()> {
|
||||
|
||||
let mut grad_logits_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
|
||||
let mut grad_w2_pb_d = stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?;
|
||||
let mut grad_b2_pb_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
@@ -669,7 +696,10 @@ fn frd_layer1_bwd_relu_mask_zeros_grad() -> Result<()> {
|
||||
|
||||
let mut grad_logits_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
let mut loss_d = stream.alloc_zeros::<f32>(b_size * FRD_N_HORIZONS)?;
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &mut loss_d, b_size)?;
|
||||
{
|
||||
let loss_ptr = loss_d.raw_ptr();
|
||||
head.softmax_ce_grad(&logits_d, &labels_d, &mut grad_logits_d, &loss_ptr, b_size)?;
|
||||
}
|
||||
|
||||
let mut grad_w2_pb_d = stream.alloc_zeros::<f32>(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?;
|
||||
let mut grad_b2_pb_d = stream.alloc_zeros::<f32>(b_size * FRD_OUT_DIM)?;
|
||||
|
||||
@@ -92,13 +92,14 @@ fn g1_isv_bootstrap_writes_canonical_values() {
|
||||
};
|
||||
let trainer = IntegratedTrainer::new(&dev, cfg).expect("IntegratedTrainer::new");
|
||||
|
||||
// Read full ISV slice to host. Uses the same pattern as the
|
||||
// trainer's own per-step ISV mirror refresh.
|
||||
// Read full ISV slice via the mapped-pinned `isv_mapped` buffer.
|
||||
// Mapped-pinned coherence: host pages reflect device writes after
|
||||
// a stream sync on the producing stream. The R1 bootstrap launches
|
||||
// run on the trainer's stream during `new()`, which finalises
|
||||
// synchronously inside `IntegratedTrainer::new`.
|
||||
let host_slice = trainer.isv_host_slice();
|
||||
let mut isv = vec![0.0_f32; RL_SLOTS_END];
|
||||
let stream = dev.cuda_stream().expect("cuda_stream");
|
||||
stream
|
||||
.memcpy_dtoh(&trainer.isv_d, isv.as_mut_slice())
|
||||
.expect("isv dtoh");
|
||||
isv.copy_from_slice(&host_slice[..RL_SLOTS_END]);
|
||||
|
||||
// Floating-point exact equality is the right oracle here — each
|
||||
// kernel's bootstrap path is `isv[slot] = K_BOOTSTRAP; return;`
|
||||
|
||||
@@ -62,16 +62,18 @@ fn upload(
|
||||
d
|
||||
}
|
||||
|
||||
/// Read the trainer's ISV state via the mapped-pinned `isv_mapped`
|
||||
/// buffer's host-side pages. Mapped-pinned coherence guarantees the
|
||||
/// host view is current after a stream synchronise on the producing
|
||||
/// stream — callers MUST sync before calling this.
|
||||
fn readback_isv(
|
||||
dev: &MlDevice,
|
||||
isv_d: &cudarc::driver::CudaSlice<f32>,
|
||||
_dev: &MlDevice,
|
||||
trainer: &IntegratedTrainer,
|
||||
) -> Vec<f32> {
|
||||
let mut isv = vec![0.0_f32; RL_SLOTS_END];
|
||||
let stream = dev.cuda_stream().expect("cuda_stream").clone();
|
||||
stream
|
||||
.memcpy_dtoh(isv_d, isv.as_mut_slice())
|
||||
.expect("isv dtoh");
|
||||
isv
|
||||
let host_slice = trainer.isv_host_slice();
|
||||
let mut out = vec![0.0_f32; RL_SLOTS_END];
|
||||
out.copy_from_slice(&host_slice[..RL_SLOTS_END]);
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -83,7 +85,7 @@ fn r3_ema_update_on_done_first_observation_bootstrap_replaces_directly() {
|
||||
// Pre-condition: the EMA-input slot is at sentinel zero (R1
|
||||
// bootstraps ISV[400..406] for controllers; ISV[417..423] EMA-input
|
||||
// slots stay at alloc_zeros per the R1 invariant).
|
||||
let isv_before = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv_before = readback_isv(&dev, &trainer);
|
||||
assert_eq!(
|
||||
isv_before[RL_MEAN_ABS_PNL_EMA_INDEX], 0.0,
|
||||
"pre-condition: EMA slot must be sentinel zero"
|
||||
@@ -100,7 +102,7 @@ fn r3_ema_update_on_done_first_observation_bootstrap_replaces_directly() {
|
||||
.expect("ema_update_on_done");
|
||||
stream.synchronize().expect("sync");
|
||||
|
||||
let isv_after = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv_after = readback_isv(&dev, &trainer);
|
||||
let val = isv_after[RL_MEAN_ABS_PNL_EMA_INDEX];
|
||||
// Exact equality — bootstrap path is `isv[slot] = mean_obs` with
|
||||
// no arithmetic. Any drift indicates a wrong code path.
|
||||
@@ -118,7 +120,7 @@ fn r3_ema_update_on_done_first_observation_bootstrap_replaces_directly() {
|
||||
.expect("ema_update_on_done hold");
|
||||
stream.synchronize().expect("sync");
|
||||
|
||||
let isv_hold = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv_hold = readback_isv(&dev, &trainer);
|
||||
assert_eq!(
|
||||
isv_hold[RL_MEAN_ABS_PNL_EMA_INDEX], 7.0,
|
||||
"hold step (no done) must preserve EMA, not blend toward 0"
|
||||
@@ -149,7 +151,7 @@ fn r3_ema_update_per_step_converges_to_constant_input() {
|
||||
}
|
||||
stream.synchronize().expect("sync");
|
||||
|
||||
let isv = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv = readback_isv(&dev, &trainer);
|
||||
let val = isv[RL_KL_PI_EMA_INDEX];
|
||||
assert!(
|
||||
(val - k).abs() < 1e-4,
|
||||
@@ -172,7 +174,7 @@ fn r3_compute_advantage_return_formula_holds() {
|
||||
// computes its expected values from whatever γ ISV holds, so the
|
||||
// pre-condition is just "γ is in the valid bounded range" rather
|
||||
// than a hardcoded canonical value.
|
||||
let isv = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv = readback_isv(&dev, &trainer);
|
||||
let gamma = isv[RL_GAMMA_INDEX];
|
||||
assert!(
|
||||
gamma >= 0.90 && gamma <= 0.999,
|
||||
|
||||
@@ -95,16 +95,18 @@ fn upload_f32(stream: &Arc<CudaStream>, host: &[f32]) -> cudarc::driver::CudaSli
|
||||
d
|
||||
}
|
||||
|
||||
/// Read the trainer's ISV state via the mapped-pinned `isv_mapped`
|
||||
/// buffer. Mapped-pinned coherence: host pages reflect device writes
|
||||
/// after the producing stream has been synchronised — callers MUST
|
||||
/// sync before calling.
|
||||
fn readback_isv(
|
||||
dev: &MlDevice,
|
||||
isv_d: &cudarc::driver::CudaSlice<f32>,
|
||||
_dev: &MlDevice,
|
||||
trainer: &IntegratedTrainer,
|
||||
) -> Vec<f32> {
|
||||
let mut isv = vec![0.0_f32; RL_SLOTS_END];
|
||||
let stream = dev.cuda_stream().expect("cuda_stream").clone();
|
||||
stream
|
||||
.memcpy_dtoh(isv_d, isv.as_mut_slice())
|
||||
.expect("isv dtoh");
|
||||
isv
|
||||
let host_slice = trainer.isv_host_slice();
|
||||
let mut out = vec![0.0_f32; RL_SLOTS_END];
|
||||
out.copy_from_slice(&host_slice[..RL_SLOTS_END]);
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -114,7 +116,7 @@ fn g3_per_step_controllers_move_isv_outputs_when_fed_real_emas() {
|
||||
let stream = dev.cuda_stream().expect("cuda_stream").clone();
|
||||
|
||||
// Pre-condition: R1 bootstrapped ISV[400..406].
|
||||
let isv_before = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv_before = readback_isv(&dev, &trainer);
|
||||
assert_eq!(isv_before[RL_GAMMA_INDEX], GAMMA_BOOTSTRAP);
|
||||
assert_eq!(isv_before[RL_TARGET_TAU_INDEX], TAU_BOOTSTRAP);
|
||||
assert_eq!(isv_before[RL_PPO_CLIP_INDEX], EPS_BOOTSTRAP);
|
||||
@@ -226,7 +228,7 @@ fn g3_per_step_controllers_move_isv_outputs_when_fed_real_emas() {
|
||||
|
||||
// Verify the EMA producers wrote what we expected (sanity check
|
||||
// before testing the controllers themselves).
|
||||
let isv_after_ema = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv_after_ema = readback_isv(&dev, &trainer);
|
||||
for (slot, expected) in inputs {
|
||||
let got = isv_after_ema[slot];
|
||||
assert!(
|
||||
@@ -235,14 +237,22 @@ fn g3_per_step_controllers_move_isv_outputs_when_fed_real_emas() {
|
||||
);
|
||||
}
|
||||
|
||||
// Fire all 7 RL controllers per-step. Each reads its EMA input
|
||||
// and Wiener-blends its output away from the bootstrap value.
|
||||
// Fire all 10 RL controllers in a single fused kernel launch.
|
||||
// Each reads its EMA input and Wiener-blends its output away
|
||||
// from the bootstrap value. The fused kernel replaces the prior
|
||||
// sequential `launch_rl_controllers_per_step` chain — see
|
||||
// `crates/ml-alpha/cuda/rl_fused_controllers.cu`.
|
||||
//
|
||||
// `build_trainer` pins `n_batch = 1` (see top of this file) — the
|
||||
// controllers don't actually consume the batch dim, but the
|
||||
// launcher requires a non-zero value.
|
||||
let b_size = 1_usize;
|
||||
trainer
|
||||
.launch_rl_controllers_per_step()
|
||||
.expect("launch_rl_controllers_per_step");
|
||||
.launch_rl_fused_controllers(b_size)
|
||||
.expect("launch_rl_fused_controllers");
|
||||
stream.synchronize().expect("sync after controllers");
|
||||
|
||||
let isv_after = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv_after = readback_isv(&dev, &trainer);
|
||||
|
||||
// Each output slot must have moved off the bootstrap value. If
|
||||
// the controller didn't fire (wrong slot wiring, missing launch,
|
||||
@@ -300,18 +310,20 @@ fn g4_dqn_target_soft_update_implements_polyak_formula() {
|
||||
.expect("dtoh w_target before");
|
||||
|
||||
// τ = ISV[401] bootstrap value = 0.005.
|
||||
let isv = readback_isv(&dev, &trainer.isv_d);
|
||||
let isv = readback_isv(&dev, &trainer);
|
||||
let tau = isv[RL_TARGET_TAU_INDEX];
|
||||
assert!(
|
||||
(tau - TAU_BOOTSTRAP).abs() < 1e-6,
|
||||
"pre-condition: τ should be R1-bootstrapped to {TAU_BOOTSTRAP}; got {tau}"
|
||||
);
|
||||
|
||||
// Fire the soft update.
|
||||
let isv_d_clone = trainer.isv_d.clone();
|
||||
// Fire the soft update. The new soft_update_target takes the
|
||||
// mapped-pinned ISV device pointer (stable, captured at trainer
|
||||
// construction).
|
||||
let isv_dev_ptr = trainer.isv_dev_ptr;
|
||||
trainer
|
||||
.dqn_head
|
||||
.soft_update_target(&isv_d_clone)
|
||||
.soft_update_target(&isv_dev_ptr)
|
||||
.expect("soft_update_target");
|
||||
stream.synchronize().expect("sync after soft_update");
|
||||
|
||||
|
||||
@@ -153,17 +153,18 @@ fn default_pyramid_ctx(stream: &Arc<CudaStream>, b_size: usize) -> Result<Pyrami
|
||||
})
|
||||
}
|
||||
|
||||
/// Overwrite a single ISV slot host-side then push the whole isv_d
|
||||
/// to the device. Cheap because the trainer's `isv_d` is small
|
||||
/// (~500 floats).
|
||||
/// Overwrite a single ISV slot via the mapped-pinned `isv_mapped`
|
||||
/// buffer's host pages. Mapped-pinned coherence: subsequent kernel
|
||||
/// reads via `isv_dev_ptr` see the write after a stream barrier on
|
||||
/// the consuming stream.
|
||||
fn set_isv_slot(
|
||||
trainer: &mut IntegratedTrainer,
|
||||
stream: &Arc<CudaStream>,
|
||||
_stream: &Arc<CudaStream>,
|
||||
slot: usize,
|
||||
value: f32,
|
||||
) -> Result<()> {
|
||||
trainer.isv_host[slot] = value;
|
||||
write_slice_f32_d_pub(stream, &trainer.isv_host, &mut trainer.isv_d)
|
||||
trainer.isv_mapped.write_record(slot, value);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user