feat(alpha): alpha_linear_q kernels + launchers for Task 12 DQN smoke
Phase E.1 Task 12a. Three new CUDA kernels for the H=600 DQN smoke
(Task 12 proper) that lands in a follow-up commit:
alpha_linear_q_forward_kernel Q = X · W^T + b
alpha_linear_q_grad_kernel dW, db sparse MSE-TD over taken actions
alpha_linear_q_sgd_step_kernel element-wise params -= lr · grad
Architecture: single linear layer, no hidden layer. The Phase E state
vector has meaningful direct features (alpha_logit, spread_bps, position,
ofi_sum_5, …) so linear Q can capture real relations like Q[Buy] ∝
alpha_logit. If linear can't pass the kill-criteria gate, no architecture
upgrade will save it — and the smoke proceeds with NoisyNet escalation
per the plan.
Sparse gradient: only the taken action contributes (standard DQN TD
loss). No atomicAdd needed — one thread per (i, j) loops over the batch
and adds only when actions[b] == i.
GPU contract:
- No host branches inside any kernel (graph-capture compatible)
- No atomicAdd (per feedback_no_atomicadd)
- All compute on GPU (forward, grad, weight update)
- Tiny launch overhead — fits per-step (batch=64 forward = 576 threads,
1 block; grad = 99 threads, 1 block)
Three pub(crate) Rust launchers in alpha_kernels.rs match the
launch_apply_pearls pattern. Cubin embedded via include_bytes!.
Smoke test `linear_q_forward_grad_sgd_round_trip_matches_hand_math`
exercises all three kernels end-to-end on a small (batch=2, state_dim=2,
n_actions=3) case with full hand-math:
Forward: Q = [[2.1, 3.2, 0.3], [4.1, 5.2, 0.3]] ✓
Grad: dW = [[-1.8, -2.7], [0.8, 1.0], [0, 0]]
db = [-0.9, 0.2, 0] ✓
SGD: W' = [[1.18, 0.27], [-0.08, 0.90], [0, 0]]
b' = [0.19, 0.18, 0.30] ✓
All within 1e-4 tolerance. `cargo test -p ml --lib alpha_kernels`:
5/5 pass on RTX 3050 Ti in 2.04s (compile witness + 4 GPU smokes).
Audit doc docs/isv-slots.md updated per Invariant 7.
This commit is contained in:
jgrusewski
@@ -1601,6 +1601,13 @@ fn main() {
|
||||
// plus a soft-V bootstrap V_soft(s') = τ · log Σ exp(Q/τ).
|
||||
// One thread per batch sample; log-sum-exp-stable softmax.
|
||||
"alpha_munchausen_target.cu",
|
||||
// Phase E.1 Task 12a (2026-05-15): linear Q-network kernels for
|
||||
// the H=600 DQN smoke. Forward (Q = X·W^T + b), MSE-TD gradient
|
||||
// (sparse over actions), and plain-SGD step. No hidden layer —
|
||||
// smoke validates whether ANY signal propagates; if linear Q
|
||||
// can't learn the kill-criteria gate, no architecture upgrade
|
||||
// will save it.
|
||||
"alpha_linear_q.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
@@ -92,6 +92,149 @@ pub(crate) unsafe fn launch_alpha_kill_criteria(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `alpha_linear_q_forward_kernel` on `stream`.
|
||||
///
|
||||
/// Computes `Q = X · W^T + b` for a batch of states. Grid: ceil(batch *
|
||||
/// n_actions / 256); Block: 256.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// All buffer pointers MUST be valid device pointers. `q_dev` MUST point
|
||||
/// at a writable `[batch * n_actions]`-float region.
|
||||
pub(crate) unsafe fn launch_alpha_linear_q_forward(
|
||||
stream: &cudarc::driver::CudaStream,
|
||||
kernel: &cudarc::driver::CudaFunction,
|
||||
w_dev: u64,
|
||||
b_dev: u64,
|
||||
x_dev: u64,
|
||||
q_dev: u64,
|
||||
batch: i32,
|
||||
state_dim: i32,
|
||||
n_actions: i32,
|
||||
) -> Result<(), MLError> {
|
||||
use cudarc::driver::{LaunchConfig, PushKernelArg};
|
||||
|
||||
debug_assert!(batch > 0, "batch must be positive");
|
||||
debug_assert!(state_dim > 0, "state_dim must be positive");
|
||||
debug_assert!(n_actions > 0, "n_actions must be positive");
|
||||
|
||||
const BLOCK: u32 = 256;
|
||||
let total = (batch as u32) * (n_actions as u32);
|
||||
let grid_x = (total + BLOCK - 1) / BLOCK;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (grid_x.max(1), 1, 1),
|
||||
block_dim: (BLOCK, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
stream
|
||||
.launch_builder(kernel)
|
||||
.arg(&w_dev)
|
||||
.arg(&b_dev)
|
||||
.arg(&x_dev)
|
||||
.arg(&q_dev)
|
||||
.arg(&batch)
|
||||
.arg(&state_dim)
|
||||
.arg(&n_actions)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("alpha_linear_q_forward launch: {e}")))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `alpha_linear_q_grad_kernel` on `stream`.
|
||||
///
|
||||
/// Computes sparse MSE-TD gradients `dW = scale · (Q − target) ⊗ X^T`
|
||||
/// and `db = scale · (Q − target)`, restricted to the taken action per
|
||||
/// sample.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// All buffer pointers MUST be valid device pointers. `dw_dev` MUST point
|
||||
/// at a writable `[n_actions * state_dim]`-float region; `db_dev` at
|
||||
/// `[n_actions]`.
|
||||
pub(crate) unsafe fn launch_alpha_linear_q_grad(
|
||||
stream: &cudarc::driver::CudaStream,
|
||||
kernel: &cudarc::driver::CudaFunction,
|
||||
q_dev: u64,
|
||||
target_dev: u64,
|
||||
actions_dev: u64,
|
||||
x_dev: u64,
|
||||
dw_dev: u64,
|
||||
db_dev: u64,
|
||||
batch: i32,
|
||||
state_dim: i32,
|
||||
n_actions: i32,
|
||||
scale: f32,
|
||||
) -> Result<(), MLError> {
|
||||
use cudarc::driver::{LaunchConfig, PushKernelArg};
|
||||
|
||||
debug_assert!(batch > 0, "batch must be positive");
|
||||
debug_assert!(state_dim > 0, "state_dim must be positive");
|
||||
debug_assert!(n_actions > 0, "n_actions must be positive");
|
||||
|
||||
const BLOCK: u32 = 256;
|
||||
let total = (n_actions as u32) * (state_dim as u32) + (n_actions as u32);
|
||||
let grid_x = (total + BLOCK - 1) / BLOCK;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (grid_x.max(1), 1, 1),
|
||||
block_dim: (BLOCK, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
stream
|
||||
.launch_builder(kernel)
|
||||
.arg(&q_dev)
|
||||
.arg(&target_dev)
|
||||
.arg(&actions_dev)
|
||||
.arg(&x_dev)
|
||||
.arg(&dw_dev)
|
||||
.arg(&db_dev)
|
||||
.arg(&batch)
|
||||
.arg(&state_dim)
|
||||
.arg(&n_actions)
|
||||
.arg(&scale)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("alpha_linear_q_grad launch: {e}")))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `alpha_linear_q_sgd_step_kernel` on `stream`.
|
||||
///
|
||||
/// Element-wise `params -= lr · grad`. Works for either the weight
|
||||
/// matrix or the bias vector — caller passes `n = numel`.
|
||||
///
|
||||
/// # Safety
|
||||
///
|
||||
/// `params_dev` and `grad_dev` MUST be valid device pointers to buffers
|
||||
/// of at least `n` floats. `params_dev` is modified in place.
|
||||
pub(crate) unsafe fn launch_alpha_linear_q_sgd_step(
|
||||
stream: &cudarc::driver::CudaStream,
|
||||
kernel: &cudarc::driver::CudaFunction,
|
||||
params_dev: u64,
|
||||
grad_dev: u64,
|
||||
lr: f32,
|
||||
n: i32,
|
||||
) -> Result<(), MLError> {
|
||||
use cudarc::driver::{LaunchConfig, PushKernelArg};
|
||||
|
||||
debug_assert!(n > 0, "n must be positive");
|
||||
|
||||
const BLOCK: u32 = 256;
|
||||
let grid_x = ((n as u32) + BLOCK - 1) / BLOCK;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (grid_x.max(1), 1, 1),
|
||||
block_dim: (BLOCK, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
stream
|
||||
.launch_builder(kernel)
|
||||
.arg(¶ms_dev)
|
||||
.arg(&grad_dev)
|
||||
.arg(&lr)
|
||||
.arg(&n)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("alpha_linear_q_sgd_step launch: {e}")))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Launch `alpha_munchausen_target_kernel` on `stream`.
|
||||
///
|
||||
/// Computes the Munchausen-augmented target for each batch sample:
|
||||
@@ -181,6 +324,11 @@ static ALPHA_KILL_CRITERIA_CUBIN: &[u8] =
|
||||
static APPLY_PEARLS_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/apply_pearls_kernel.cubin"));
|
||||
|
||||
/// Precompiled linear Q-network kernel cubin (forward / grad / SGD-step).
|
||||
/// Loaded by the Phase E.1 H=600 DQN smoke (Task 12).
|
||||
pub(crate) static ALPHA_LINEAR_Q_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/alpha_linear_q.cubin"));
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -654,4 +802,192 @@ mod tests {
|
||||
.expect("apply_pearls launch");
|
||||
}
|
||||
}
|
||||
|
||||
/// End-to-end GPU smoke for the linear Q-network kernels (Task 12a).
|
||||
///
|
||||
/// Validates forward → grad → SGD against hand-computed expected values.
|
||||
///
|
||||
/// Setup: W = [[1, 0], [0, 1], [0, 0]] (3 actions, 2 state dims),
|
||||
/// b = [0.1, 0.2, 0.3]
|
||||
/// X = [[2, 3], [4, 5]] (batch=2)
|
||||
///
|
||||
/// Forward (hand-math):
|
||||
/// Q[0, 0] = 1·2 + 0·3 + 0.1 = 2.1
|
||||
/// Q[0, 1] = 0·2 + 1·3 + 0.2 = 3.2
|
||||
/// Q[0, 2] = 0·2 + 0·3 + 0.3 = 0.3
|
||||
/// Q[1, 0] = 1·4 + 0·5 + 0.1 = 4.1
|
||||
/// Q[1, 1] = 0·4 + 1·5 + 0.2 = 5.2
|
||||
/// Q[1, 2] = 0·4 + 0·5 + 0.3 = 0.3
|
||||
///
|
||||
/// Grad (actions=[0, 1], target=[3.0, 5.0], scale=1.0):
|
||||
/// td_err[0, action 0] = 2.1 - 3.0 = -0.9
|
||||
/// td_err[1, action 1] = 5.2 - 5.0 = +0.2
|
||||
/// dW[0, 0] = -0.9·2 = -1.8 (only sample 0 taken action 0)
|
||||
/// dW[0, 1] = -0.9·3 = -2.7
|
||||
/// dW[1, 0] = +0.2·4 = +0.8 (only sample 1 taken action 1)
|
||||
/// dW[1, 1] = +0.2·5 = +1.0
|
||||
/// dW[2, *] = 0 (action 2 never taken)
|
||||
/// db[0] = -0.9, db[1] = +0.2, db[2] = 0
|
||||
///
|
||||
/// SGD (lr=0.1): W = W - 0.1·dW
|
||||
/// W[0, 0]: 1.0 - 0.1·(-1.8) = 1.18
|
||||
/// W[0, 1]: 0.0 - 0.1·(-2.7) = 0.27
|
||||
/// W[1, 0]: 0.0 - 0.1·(+0.8) = -0.08
|
||||
/// W[1, 1]: 1.0 - 0.1·(+1.0) = 0.90
|
||||
/// W[2, *]: unchanged
|
||||
#[test]
|
||||
fn linear_q_forward_grad_sgd_round_trip_matches_hand_math() {
|
||||
let ctx = match CudaContext::new(0) {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
eprintln!("Skipping GPU smoke — no CUDA device 0: {e}");
|
||||
return;
|
||||
}
|
||||
};
|
||||
let stream = ctx.default_stream();
|
||||
|
||||
let module = ctx
|
||||
.load_cubin(ALPHA_LINEAR_Q_CUBIN.to_vec())
|
||||
.expect("alpha_linear_q cubin load");
|
||||
let fwd_kernel = module
|
||||
.load_function("alpha_linear_q_forward_kernel")
|
||||
.expect("forward kernel load");
|
||||
let grad_kernel = module
|
||||
.load_function("alpha_linear_q_grad_kernel")
|
||||
.expect("grad kernel load");
|
||||
let sgd_kernel = module
|
||||
.load_function("alpha_linear_q_sgd_step_kernel")
|
||||
.expect("sgd kernel load");
|
||||
|
||||
let batch: i32 = 2;
|
||||
let state_dim: i32 = 2;
|
||||
let n_actions: i32 = 3;
|
||||
|
||||
// Row-major: W[i, j] at W[i * state_dim + j].
|
||||
let w_init: Vec<f32> = vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0];
|
||||
let b_init: Vec<f32> = vec![0.1, 0.2, 0.3];
|
||||
let x: Vec<f32> = vec![2.0, 3.0, 4.0, 5.0];
|
||||
let actions: Vec<i32> = vec![0, 1];
|
||||
let target: Vec<f32> = vec![3.0, 5.0];
|
||||
|
||||
let mut w_dev = stream.clone_htod(&w_init).expect("upload W");
|
||||
let mut b_dev = stream.clone_htod(&b_init).expect("upload b");
|
||||
let x_dev = stream.clone_htod(&x).expect("upload X");
|
||||
let actions_dev = stream.clone_htod(&actions).expect("upload actions");
|
||||
let target_dev = stream.clone_htod(&target).expect("upload target");
|
||||
let mut q_dev = stream
|
||||
.alloc_zeros::<f32>((batch * n_actions) as usize)
|
||||
.expect("alloc Q");
|
||||
let mut dw_dev = stream
|
||||
.alloc_zeros::<f32>((n_actions * state_dim) as usize)
|
||||
.expect("alloc dW");
|
||||
let mut db_dev = stream
|
||||
.alloc_zeros::<f32>(n_actions as usize)
|
||||
.expect("alloc db");
|
||||
|
||||
// --- Forward ---
|
||||
{
|
||||
let (w_ptr, _g0) = w_dev.device_ptr(&stream);
|
||||
let (b_ptr, _g1) = b_dev.device_ptr(&stream);
|
||||
let (x_ptr, _g2) = x_dev.device_ptr(&stream);
|
||||
let (q_ptr, _g3) = q_dev.device_ptr_mut(&stream);
|
||||
unsafe {
|
||||
launch_alpha_linear_q_forward(
|
||||
&stream, &fwd_kernel,
|
||||
w_ptr, b_ptr, x_ptr, q_ptr,
|
||||
batch, state_dim, n_actions,
|
||||
)
|
||||
.expect("forward launch");
|
||||
}
|
||||
}
|
||||
stream.synchronize().expect("sync after forward");
|
||||
let q_after = stream.clone_dtoh(&q_dev).expect("download Q");
|
||||
let q_expected = [2.1_f32, 3.2, 0.3, 4.1, 5.2, 0.3];
|
||||
for (i, (g, e)) in q_after.iter().zip(q_expected.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - e).abs() < 1e-4,
|
||||
"forward Q[{i}] = {g} expected {e}"
|
||||
);
|
||||
}
|
||||
|
||||
// --- Grad ---
|
||||
{
|
||||
let (q_ptr, _g0) = q_dev.device_ptr(&stream);
|
||||
let (target_ptr, _g1) = target_dev.device_ptr(&stream);
|
||||
let (actions_ptr, _g2) = actions_dev.device_ptr(&stream);
|
||||
let (x_ptr, _g3) = x_dev.device_ptr(&stream);
|
||||
let (dw_ptr, _g4) = dw_dev.device_ptr_mut(&stream);
|
||||
let (db_ptr, _g5) = db_dev.device_ptr_mut(&stream);
|
||||
unsafe {
|
||||
launch_alpha_linear_q_grad(
|
||||
&stream, &grad_kernel,
|
||||
q_ptr, target_ptr, actions_ptr, x_ptr,
|
||||
dw_ptr, db_ptr,
|
||||
batch, state_dim, n_actions,
|
||||
1.0, // scale
|
||||
)
|
||||
.expect("grad launch");
|
||||
}
|
||||
}
|
||||
stream.synchronize().expect("sync after grad");
|
||||
let dw_after = stream.clone_dtoh(&dw_dev).expect("download dW");
|
||||
let db_after = stream.clone_dtoh(&db_dev).expect("download db");
|
||||
let dw_expected = [-1.8_f32, -2.7, 0.8, 1.0, 0.0, 0.0];
|
||||
let db_expected = [-0.9_f32, 0.2, 0.0];
|
||||
for (i, (g, e)) in dw_after.iter().zip(dw_expected.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - e).abs() < 1e-4,
|
||||
"grad dW[{i}] = {g} expected {e}"
|
||||
);
|
||||
}
|
||||
for (i, (g, e)) in db_after.iter().zip(db_expected.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - e).abs() < 1e-4,
|
||||
"grad db[{i}] = {g} expected {e}"
|
||||
);
|
||||
}
|
||||
|
||||
// --- SGD step: W -= 0.1 · dW; b -= 0.1 · db ---
|
||||
{
|
||||
let (w_ptr, _g0) = w_dev.device_ptr_mut(&stream);
|
||||
let (dw_ptr, _g1) = dw_dev.device_ptr(&stream);
|
||||
unsafe {
|
||||
launch_alpha_linear_q_sgd_step(
|
||||
&stream, &sgd_kernel,
|
||||
w_ptr, dw_ptr, 0.1,
|
||||
n_actions * state_dim,
|
||||
)
|
||||
.expect("sgd W launch");
|
||||
}
|
||||
}
|
||||
{
|
||||
let (b_ptr, _g0) = b_dev.device_ptr_mut(&stream);
|
||||
let (db_ptr, _g1) = db_dev.device_ptr(&stream);
|
||||
unsafe {
|
||||
launch_alpha_linear_q_sgd_step(
|
||||
&stream, &sgd_kernel,
|
||||
b_ptr, db_ptr, 0.1,
|
||||
n_actions,
|
||||
)
|
||||
.expect("sgd b launch");
|
||||
}
|
||||
}
|
||||
stream.synchronize().expect("sync after sgd");
|
||||
let w_after = stream.clone_dtoh(&w_dev).expect("download W");
|
||||
let b_after = stream.clone_dtoh(&b_dev).expect("download b");
|
||||
let w_expected = [1.18_f32, 0.27, -0.08, 0.90, 0.0, 0.0];
|
||||
let b_expected = [0.19_f32, 0.18, 0.30];
|
||||
for (i, (g, e)) in w_after.iter().zip(w_expected.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - e).abs() < 1e-4,
|
||||
"sgd W[{i}] = {g} expected {e}"
|
||||
);
|
||||
}
|
||||
for (i, (g, e)) in b_after.iter().zip(b_expected.iter()).enumerate() {
|
||||
assert!(
|
||||
(g - e).abs() < 1e-4,
|
||||
"sgd b[{i}] = {g} expected {e}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
133
crates/ml/src/cuda_pipeline/alpha_linear_q.cu
Normal file
133
crates/ml/src/cuda_pipeline/alpha_linear_q.cu
Normal file
@@ -0,0 +1,133 @@
|
||||
// crates/ml/src/cuda_pipeline/alpha_linear_q.cu
|
||||
//
|
||||
// Phase E.1 Task 12a (2026-05-15): linear Q-network kernels for the
|
||||
// H=600 DQN smoke. Single linear layer:
|
||||
//
|
||||
// Q[b, i] = sum_j W[i, j] * X[b, j] + b[i]
|
||||
//
|
||||
// No hidden layer — the smoke's purpose is to validate whether ANY signal
|
||||
// propagates through the env + reward shaping + kill-criteria pipeline.
|
||||
// If linear Q can't learn, no architecture upgrade will save it. The
|
||||
// 10-dim state has meaningful direct features (alpha_logit, spread_bps,
|
||||
// position, ...) so a linear model has real expressive power for this
|
||||
// kill-criteria gate.
|
||||
//
|
||||
// Three kernels in this file:
|
||||
// 1. alpha_linear_q_forward_kernel — batched Q forward pass
|
||||
// 2. alpha_linear_q_grad_kernel — gradient of MSE-style TD loss
|
||||
// w.r.t. W and b (sparse over actions)
|
||||
// 3. alpha_linear_q_sgd_step_kernel — plain SGD update (W -= lr · dW)
|
||||
//
|
||||
// **GPU contract:** all three are simple element-wise / inner-loop work.
|
||||
// No atomicAdd (per `feedback_no_atomicadd`). No host branches inside
|
||||
// kernels (per `pearl_no_host_branches_in_captured_graph`). Tiny grid /
|
||||
// block configs fit per-step launch overhead.
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
// ----------------------------------------------------------------------
|
||||
// (1) Forward: Q = X · W^T + b
|
||||
// ----------------------------------------------------------------------
|
||||
//
|
||||
// Threads: one per (b, i) where b ∈ [0, batch) and i ∈ [0, n_actions).
|
||||
// Grid: (ceil(batch * n_actions / 256), 1, 1); Block: (256, 1, 1).
|
||||
//
|
||||
// Each thread loops over state_dim accumulating W[i, j] * X[b, j].
|
||||
extern "C" __global__ void alpha_linear_q_forward_kernel(
|
||||
const float* __restrict__ W, // [n_actions, state_dim] row-major
|
||||
const float* __restrict__ b, // [n_actions]
|
||||
const float* __restrict__ X, // [batch, state_dim] row-major
|
||||
float* __restrict__ Q, // [batch, n_actions] row-major
|
||||
int batch,
|
||||
int state_dim,
|
||||
int n_actions
|
||||
) {
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int total = batch * n_actions;
|
||||
if (idx >= total) return;
|
||||
|
||||
const int sample = idx / n_actions;
|
||||
const int action = idx % n_actions;
|
||||
|
||||
float acc = b[action];
|
||||
const int w_row_base = action * state_dim;
|
||||
const int x_row_base = sample * state_dim;
|
||||
#pragma unroll 4
|
||||
for (int j = 0; j < state_dim; ++j) {
|
||||
acc += W[w_row_base + j] * X[x_row_base + j];
|
||||
}
|
||||
Q[idx] = acc;
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------
|
||||
// (2) Gradient: ∂L/∂W and ∂L/∂b for sparse TD-style loss
|
||||
// ----------------------------------------------------------------------
|
||||
//
|
||||
// Per the standard DQN gradient, only the action taken at sample b
|
||||
// contributes to the gradient. For each (i, j):
|
||||
//
|
||||
// dW[i, j] = scale · sum_{b: actions[b] == i} (Q[b, i] − target[b]) · X[b, j]
|
||||
// db[i] = scale · sum_{b: actions[b] == i} (Q[b, i] − target[b])
|
||||
//
|
||||
// `scale` defaults to 1/batch (standard MSE batching) but is passed as
|
||||
// a kernel arg so a controller could ISV-drive it (e.g., IS-weighted
|
||||
// prioritized replay).
|
||||
//
|
||||
// Threads: 1 per (i, j) for dW, 1 per i for db. Combined index into a
|
||||
// single (n_actions * state_dim + n_actions)-thread grid for simplicity.
|
||||
// At batch=64 and (9, 10): 99 threads — single block, fine.
|
||||
extern "C" __global__ void alpha_linear_q_grad_kernel(
|
||||
const float* __restrict__ Q, // [batch, n_actions] — current Q from forward
|
||||
const float* __restrict__ target, // [batch] — TD target for the taken action
|
||||
const int* __restrict__ actions, // [batch] — action taken per sample
|
||||
const float* __restrict__ X, // [batch, state_dim] — input states
|
||||
float* __restrict__ dW, // [n_actions, state_dim] — output, overwritten
|
||||
float* __restrict__ db, // [n_actions] — output, overwritten
|
||||
int batch,
|
||||
int state_dim,
|
||||
int n_actions,
|
||||
float scale // typ. 1/batch
|
||||
) {
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int total_w = n_actions * state_dim;
|
||||
|
||||
if (idx < total_w) {
|
||||
const int i = idx / state_dim; // action
|
||||
const int j = idx % state_dim; // input feature dim
|
||||
float grad = 0.0f;
|
||||
for (int b_idx = 0; b_idx < batch; ++b_idx) {
|
||||
// Sparse: only contribute when sample b's action matches i.
|
||||
if (actions[b_idx] == i) {
|
||||
const float td_err = Q[b_idx * n_actions + i] - target[b_idx];
|
||||
grad += td_err * X[b_idx * state_dim + j];
|
||||
}
|
||||
}
|
||||
dW[idx] = grad * scale;
|
||||
} else if (idx < total_w + n_actions) {
|
||||
const int i = idx - total_w;
|
||||
float grad = 0.0f;
|
||||
for (int b_idx = 0; b_idx < batch; ++b_idx) {
|
||||
if (actions[b_idx] == i) {
|
||||
grad += Q[b_idx * n_actions + i] - target[b_idx];
|
||||
}
|
||||
}
|
||||
db[i] = grad * scale;
|
||||
}
|
||||
}
|
||||
|
||||
// ----------------------------------------------------------------------
|
||||
// (3) SGD step: W -= lr · dW (element-wise, flat over W or b)
|
||||
// ----------------------------------------------------------------------
|
||||
//
|
||||
// Generic enough to update either the weight matrix or the bias vector
|
||||
// — caller passes `n = numel` of the buffer.
|
||||
extern "C" __global__ void alpha_linear_q_sgd_step_kernel(
|
||||
float* __restrict__ params, // [n] — overwritten
|
||||
const float* __restrict__ grad, // [n]
|
||||
float lr,
|
||||
int n
|
||||
) {
|
||||
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= n) return;
|
||||
params[idx] -= lr * grad[idx];
|
||||
}
|
||||
Reference in New Issue
Block a user