From 6b4e7f660b76e2cf066d7060ab16760be43bb098 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 28 May 2026 15:23:28 +0200 Subject: [PATCH] =?UTF-8?q?perf(rl):=20replace=20cuBLAS=20with=20capture-s?= =?UTF-8?q?afe=20matmul=20kernels=20=E2=80=94=20unblocks=20mega-graph?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- crates/ml-alpha/build.rs | 2 + crates/ml-alpha/cuda/dqn_q_head_fwd_bwd.cu | 248 ++++++++++++ crates/ml-alpha/cuda/rl_iqn_matmul.cu | 273 +++++++++++++ crates/ml-alpha/src/rl/dqn.rs | 373 ++++++++---------- crates/ml-alpha/src/rl/iqn.rs | 258 ++++++------ crates/ml-alpha/tests/frd_head.rs | 50 ++- crates/ml-alpha/tests/isv_bootstrap.rs | 13 +- crates/ml-alpha/tests/r3_ema_advantage.rs | 28 +- .../tests/r5_controllers_and_soft_update.rs | 50 ++- .../tests/trade_management_kernels.rs | 13 +- 10 files changed, 901 insertions(+), 407 deletions(-) create mode 100644 crates/ml-alpha/cuda/dqn_q_head_fwd_bwd.cu create mode 100644 crates/ml-alpha/cuda/rl_iqn_matmul.cu diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 61e286346..7cda8d993 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -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] diff --git a/crates/ml-alpha/cuda/dqn_q_head_fwd_bwd.cu b/crates/ml-alpha/cuda/dqn_q_head_fwd_bwd.cu new file mode 100644 index 000000000..9559edcd7 --- /dev/null +++ b/crates/ml-alpha/cuda/dqn_q_head_fwd_bwd.cu @@ -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; + } +} diff --git a/crates/ml-alpha/cuda/rl_iqn_matmul.cu b/crates/ml-alpha/cuda/rl_iqn_matmul.cu new file mode 100644 index 000000000..f9b29a19d --- /dev/null +++ b/crates/ml-alpha/cuda/rl_iqn_matmul.cu @@ -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; + } +} diff --git a/crates/ml-alpha/src/rl/dqn.rs b/crates/ml-alpha/src/rl/dqn.rs index c725b6a69..95dd0aa7a 100644 --- a/crates/ml-alpha/src/rl/dqn.rs +++ b/crates/ml-alpha/src/rl/dqn.rs @@ -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, - // ── 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, - /// 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, + /// 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::(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, @@ -436,59 +405,22 @@ impl DqnHead { b_size: usize, logits_out: &mut CudaSlice, ) -> 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, - bias: &CudaSlice, - 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::()) 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::()) 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::() 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(()) } diff --git a/crates/ml-alpha/src/rl/iqn.rs b/crates/ml-alpha/src/rl/iqn.rs index 1aeac3708..7d9e11e36 100644 --- a/crates/ml-alpha/src/rl/iqn.rs +++ b/crates/ml-alpha/src/rl/iqn.rs @@ -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, raw_stream: CUstream, - // ── cuBLAS ─────────────────────────────────────────────────────── - cublas: CudaBlas, - _cublas_workspace: CudaSlice, + // ── 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, + /// 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, @@ -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 = 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::(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))?; } } diff --git a/crates/ml-alpha/tests/frd_head.rs b/crates/ml-alpha/tests/frd_head.rs index e7a4bc1d9..f8a83de7d 100644 --- a/crates/ml-alpha/tests/frd_head.rs +++ b/crates/ml-alpha/tests/frd_head.rs @@ -211,7 +211,10 @@ fn frd_softmax_ce_grad_uniform_logits_match_log_n_atoms() -> Result<()> { let mut grad_d = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; let mut loss_d = stream.alloc_zeros::(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::(b_size * FRD_OUT_DIM)?; let mut loss_d = stream.alloc_zeros::(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::(b_size * FRD_OUT_DIM)?; let mut loss_d = stream.alloc_zeros::(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::(b_size * FRD_OUT_DIM)?; let mut loss_d = stream.alloc_zeros::(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::(b_size * FRD_OUT_DIM)?; let mut loss_d = stream.alloc_zeros::(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::(b_size * FRD_OUT_DIM)?; let mut loss_d = stream.alloc_zeros::(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::(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::(b_size * FRD_OUT_DIM)?; let mut loss_d = stream.alloc_zeros::(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::(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?; let mut grad_b2_pb_d = stream.alloc_zeros::(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::(b_size * FRD_OUT_DIM)?; let mut loss_d = stream.alloc_zeros::(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::(b_size * FRD_HIDDEN_DIM * FRD_OUT_DIM)?; let mut grad_b2_pb_d = stream.alloc_zeros::(b_size * FRD_OUT_DIM)?; diff --git a/crates/ml-alpha/tests/isv_bootstrap.rs b/crates/ml-alpha/tests/isv_bootstrap.rs index a12299ca5..295577ae3 100644 --- a/crates/ml-alpha/tests/isv_bootstrap.rs +++ b/crates/ml-alpha/tests/isv_bootstrap.rs @@ -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;` diff --git a/crates/ml-alpha/tests/r3_ema_advantage.rs b/crates/ml-alpha/tests/r3_ema_advantage.rs index c585c916a..846eca2d3 100644 --- a/crates/ml-alpha/tests/r3_ema_advantage.rs +++ b/crates/ml-alpha/tests/r3_ema_advantage.rs @@ -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, + _dev: &MlDevice, + trainer: &IntegratedTrainer, ) -> Vec { - 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, diff --git a/crates/ml-alpha/tests/r5_controllers_and_soft_update.rs b/crates/ml-alpha/tests/r5_controllers_and_soft_update.rs index 6ebec92e0..6ddef7bd5 100644 --- a/crates/ml-alpha/tests/r5_controllers_and_soft_update.rs +++ b/crates/ml-alpha/tests/r5_controllers_and_soft_update.rs @@ -95,16 +95,18 @@ fn upload_f32(stream: &Arc, 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, + _dev: &MlDevice, + trainer: &IntegratedTrainer, ) -> Vec { - 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"); diff --git a/crates/ml-alpha/tests/trade_management_kernels.rs b/crates/ml-alpha/tests/trade_management_kernels.rs index 7c08928a2..09f80f2d2 100644 --- a/crates/ml-alpha/tests/trade_management_kernels.rs +++ b/crates/ml-alpha/tests/trade_management_kernels.rs @@ -153,17 +153,18 @@ fn default_pyramid_ctx(stream: &Arc, b_size: usize) -> Result, + _stream: &Arc, 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]