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

feat(dqn-v2): E.4 encoder/decoder Rust API split — additive, no behavior change

Browse Source 
Plan 4 Task 4. Pure Rust API refactor — no kernel changes, no parameter
changes, no checkpoint break, no new ISV slot.

New types:
- StateEncoder: wraps the trunk encoder portion (h_s1 + h_s2 GEMMs)
- ValueDecoder: per-branch decoder (4 instances: Dir, Mag, Ord, Urg)
- EncoderOutput: { h_s2_dev_ptr, h_s2_dim, batch_size }
- BranchDecoderOutput: Q-per-action + V_short + V_long (Plan 2 D.3
  horizon-decomposed value already landed) + num_atoms + batch_size

BatchedForward gains 2 helper methods:
- encoder_forward_only(...) — extracted trunk encoder dispatch
- decoder_forward_only(branch_idx, ...) — single-stream per-branch
  dispatch mirroring the loop body in forward_online_raw

forward_online_raw stays intact and callable; existing call sites
(training_loop, eval, experience collector) unchanged. The trunk
portion of forward_online_raw now routes through encoder_forward_only
internally — lossless extraction (same launch sequence, same
workspace usage). The new wrappers are ADDITIVE attachment points for
Plan 4 Task 1 (Full VSN, attaches pre-encoder), Task 3 (multi-Q IQN,
attaches at decoder), and Task 6 (auxiliary heads, attaches at decoder
peer). Migration of existing callers is deferred — the goal of this
task is establishing a clean type boundary, not migrating call sites.

Smoke (RTX 3050 Ti, 5 epochs, 3 folds, T5 baseline): fold 2 best
val Sharpe = 109.06 at epoch 3 (within RNG noise of T5 landing's
fold-2 best of 95.09; the refactor is behaviorally a no-op since
encoder_forward_only is invoked with byte-identical args).

cargo check --workspace: 11 warnings (matches baseline).

No checkpoint break. No new ISV slot. No new CUDA kernels. No call-site
migration in this commit.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-25 10:40:05 +02:00
parent cfc4ccb72f
commit fbc299fa27
5 changed files with 417 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

154
crates/ml/src/cuda_pipeline/batched_forward.rs
View File

@@ -449,6 +449,141 @@ impl CublasGemmSet {
///
/// Weight pointers are raw u64 into the flat F32 `params_buf` at GOFF_* offsets.
#[allow(clippy::too_many_arguments)]
/// Run only the trunk encoder (h_s1 + h_s2) GEMMs.
///
/// Plan 4 Task 4 (E.4) extraction point: the same dispatch `forward_online_raw`
/// uses for the trunk-encoder portion. Exposed publicly so the
/// `StateEncoder` Rust wrapper can drive the encoder forward in
/// isolation; existing `forward_online_raw` callers route through
/// this method internally (lossless extraction — same launch
/// sequence, same workspace usage).
///
/// `w_ptrs[0]`/`w_ptrs[1]` = `W_s1` / `b_s1`; `w_ptrs[2]`/`w_ptrs[3]` = `W_s2` / `b_s2`.
/// Caller-owned device buffers: `states_ptr`, `h_s1_ptr`, `h_s2_ptr`.
pub fn encoder_forward_only(
&self,
stream: &Arc<CudaStream>,
states_ptr: u64,
w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
h_s1_ptr: u64,
h_s2_ptr: u64,
) -> Result<(), MLError> {
let b = self.batch_size;
let ws = self.handle.lt_workspace_ptr;
let wss = self.handle.lt_workspace_size;
// First layer: ldb = state_dim_padded (CUTLASS K-tile alignment).
// Try fused GEMM+bias+ReLU epilogue, fall back to separate kernels.
if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[0], states_ptr, h_s1_ptr, w_ptrs[1],
self.shared_h1, b, self.s1_input_dim, self.s1_ldb, ws, wss, "h_s1").is_err()
{
self.sgemm_f32_ldb(stream, w_ptrs[0], states_ptr, h_s1_ptr, self.shared_h1, b, self.s1_input_dim, self.s1_ldb, "h_s1")?;
self.launch_add_bias_relu_f32_raw(stream, h_s1_ptr, w_ptrs[1], self.shared_h1, b)?;
}
if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[2], h_s1_ptr, h_s2_ptr, w_ptrs[3],
self.shared_h2, b, self.shared_h1, self.shared_h1, ws, wss, "h_s2").is_err()
{
self.sgemm_f32(stream, w_ptrs[2], h_s1_ptr, h_s2_ptr, self.shared_h2, b, self.shared_h1, "h_s2")?;
self.launch_add_bias_relu_f32_raw(stream, h_s2_ptr, w_ptrs[3], self.shared_h2, b)?;
}
Ok(())
}
/// Run a single advantage branch's decoder (FC head + adv-logits) on the
/// caller-supplied stream, sequentially.
///
/// Plan 4 Task 4 (E.4) attachment point for the per-branch `ValueDecoder`
/// Rust wrapper. **Additive** to `forward_online_raw`: this helper runs
/// the same GEMM sequence the per-branch loop body in `forward_online_raw`
/// runs (when `use_vsn_glu == false`), but on a single stream — no fork-
/// join, no event handshake. It is intended for callers that drive the
/// branches one at a time (e.g. validation paths, future per-branch
/// experiments). Production training stays on the existing multi-stream
/// `forward_online_raw` path until/unless a follow-up explicitly
/// migrates it.
///
/// `branch_idx ∈ [0, 4)` selects which branch's weights/buffers to use.
/// `h_s2_ptr` is the encoder output (read-only). For `branch_idx == 1`
/// (magnitude), pass the pre-built mag_concat pointer to use the wider
/// `[B, SH2 + 3]` input; pass 0 for the legacy `[B, SH2]` input.
/// `branch_h_ptr` is the per-branch hidden activation buffer; the
/// `adv_logits_ptr` is the per-branch slice into the flat
/// `b_logits_buf` at byte offset
/// `sum_{j<branch_idx}(B * branch_size_j * num_atoms * 4)`.
#[allow(clippy::too_many_arguments)]
pub fn decoder_forward_only(
&self,
stream: &Arc<CudaStream>,
branch_idx: usize,
h_s2_ptr: u64,
mag_concat_ptr: u64,
w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
branch_h_ptr: u64,
adv_logits_ptr: u64,
) -> Result<(), MLError> {
if branch_idx >= 4 {
return Err(MLError::ModelError(format!(
"decoder_forward_only: branch_idx {branch_idx} out of range [0, 4)"
)));
}
let b = self.batch_size;
let ws = self.handle.lt_workspace_ptr;
let wss = self.handle.lt_workspace_size;
let na = self.num_atoms;
let branch_w_base = [8_usize, 12, 16, 20];
let w_fc_idx = branch_w_base[branch_idx];
let n_d = self.branch_q_dim_internal(branch_idx);
let use_vsn_glu = self.vsn_kernel.is_some();
if use_vsn_glu {
// VSN bottleneck + GLU gating path (sequential — main stream).
self.launch_vsn_glu_branch(
stream, branch_idx, h_s2_ptr, mag_concat_ptr, w_ptrs,
branch_h_ptr,
ws, wss,
false, "decoder",
)?;
} else {
// Legacy GEMM+bias+ReLU FC head.
let (fc_input, fc_k) = if branch_idx == 1 && mag_concat_ptr != 0 {
(mag_concat_ptr, self.mag_concat_dim)
} else {
(h_s2_ptr, self.shared_h2)
};
if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptr, w_ptrs[w_fc_idx + 1],
self.adv_h, b, fc_k, fc_k, ws, wss, "decoder_h_bd").is_err()
{
self.sgemm_f32(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptr, self.adv_h, b, fc_k, "decoder_h_bd")?;
self.launch_add_bias_relu_f32_raw(stream, branch_h_ptr, w_ptrs[w_fc_idx + 1], self.adv_h, b)?;
}
}
// Output GEMM + bias (linear): adv_logits = h_bd @ W_out^T + b_out
self.sgemm_f32(stream, w_ptrs[w_fc_idx + 2], branch_h_ptr, adv_logits_ptr, n_d * na, b, self.adv_h, "decoder_adv_logits")?;
self.launch_add_bias_f32_raw(stream, adv_logits_ptr, w_ptrs[w_fc_idx + 3], n_d * na, b)?;
Ok(())
}
/// Per-branch advantage dimension (action count for that branch — pre
/// num_atoms multiplication). Internal so `ValueDecoder` can report
/// `q_per_action_dim` without exposing the four `branch_k_size` fields.
pub(crate) fn branch_q_dim_internal(&self, branch_idx: usize) -> usize {
match branch_idx {
0 => self.branch_0_size,
1 => self.branch_1_size,
2 => self.branch_2_size,
3 => self.branch_3_size,
_ => 0,
}
}
/// Distributional atoms per branch action (C51 support size).
pub(crate) fn num_atoms(&self) -> usize { self.num_atoms }
/// Graph-safe forward: takes pre-resolved u64 buffer pointers (no device_ptr calls).
/// Use CachedPtrs from GpuDqnTrainer for all buffer addresses.
#[allow(clippy::too_many_arguments)]
@@ -463,24 +598,13 @@ impl CublasGemmSet {
mag_concat_ptr: u64,
) -> Result<(), MLError> {
let b = self.batch_size;
// First layer: ldb = state_dim_padded (CUTLASS K-tile alignment).
// Try fused GEMM+bias+ReLU epilogue, fall back to separate kernels.
let ws = self.handle.lt_workspace_ptr;
let wss = self.handle.lt_workspace_size;
if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[0], states_ptr, h_s1_ptr, w_ptrs[1],
self.shared_h1, b, self.s1_input_dim, self.s1_ldb, ws, wss, "h_s1").is_err()
{
self.sgemm_f32_ldb(stream, w_ptrs[0], states_ptr, h_s1_ptr, self.shared_h1, b, self.s1_input_dim, self.s1_ldb, "h_s1")?;
self.launch_add_bias_relu_f32_raw(stream, h_s1_ptr, w_ptrs[1], self.shared_h1, b)?;
}
if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[2], h_s1_ptr, h_s2_ptr, w_ptrs[3],
self.shared_h2, b, self.shared_h1, self.shared_h1, ws, wss, "h_s2").is_err()
{
self.sgemm_f32(stream, w_ptrs[2], h_s1_ptr, h_s2_ptr, self.shared_h2, b, self.shared_h1, "h_s2")?;
self.launch_add_bias_relu_f32_raw(stream, h_s2_ptr, w_ptrs[3], self.shared_h2, b)?;
}
// Trunk encoder (h_s1, h_s2) — extracted to `encoder_forward_only` so
// the new `StateEncoder` Rust wrapper (Plan 4 Task 4 / E.4) can reuse
// the identical dispatch sequence. Behaviorally a no-op refactor.
self.encoder_forward_only(stream, states_ptr, w_ptrs, h_s1_ptr, h_s2_ptr)?;
if self.sgemm_f32_fused_relu_bias(stream, w_ptrs[4], h_s2_ptr, h_v_ptr, w_ptrs[5],
self.value_h, b, self.shared_h2, self.shared_h2, ws, wss, "h_v").is_err()

2
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -29,6 +29,8 @@ pub mod gpu_dqn_trainer;
pub mod shared_cublas_handle;
pub mod cublas_algo_deterministic;
pub mod batched_forward;
pub mod state_encoder;
pub mod value_decoder;
pub mod batched_backward;
pub mod gpu_her;
#[cfg(test)]

106
crates/ml/src/cuda_pipeline/state_encoder.rs Normal file
View File

@@ -0,0 +1,106 @@
//! Plan 4 Task 4 (E.4): explicit state-encoder boundary for the DQN trunk.
//!
//! `StateEncoder` is a thin Rust API wrapper around `BatchedForward`'s trunk
//! encoder portion (the two shared layers `h_s1` + `h_s2`). It owns no
//! parameters and no buffers — it borrows a `BatchedForward` and forwards
//! caller-owned device pointers through the existing
//! `BatchedForward::encoder_forward_only` dispatch.
//!
//! Design intent (E.4 spec):
//! - Provide a clean type boundary so Plan 4 Task 1 (Full VSN, attaches
//! pre-encoder), Task 3 (multi-Q IQN, attaches at decoder), and Task 6
//! (auxiliary heads, attaches at decoder peer) have a stable attachment
//! point.
//! - **Additive** to the existing monolithic `forward_online_raw` — that
//! function still routes through `encoder_forward_only` internally, so
//! call sites remain unchanged.
//! - Pure Rust API refactor: no kernel changes, no parameter changes, no
//! checkpoint break.
use std::sync::Arc;
use cudarc::driver::CudaStream;
use crate::MLError;
use crate::cuda_pipeline::batched_forward::CublasGemmSet;
/// Output of one `StateEncoder::forward` call.
///
/// Carries the device pointer the encoder wrote to (`h_s2_dev_ptr`), the
/// shared-trunk output dimension, and the batch size. `h_s2_dev_ptr` is
/// always equal to the pointer the caller passed in — the type exists to
/// document the contract for downstream `ValueDecoder` consumers.
#[derive(Debug, Clone, Copy)]
pub struct EncoderOutput {
pub h_s2_dev_ptr: u64,
pub h_s2_dim: usize,
pub batch_size: usize,
}
/// Plan 4 Task 4 (E.4) state-encoder façade.
///
/// Borrows a `CublasGemmSet` (`BatchedForward`) and exposes the trunk
/// encoder (`h_s1` + `h_s2` GEMMs) as a separately-callable forward step.
/// The encoder runs on the caller's stream; callers own all device
/// buffers (`states_ptr`, `h_s1_ptr`, `h_s2_ptr`).
///
/// Owns no parameters of its own — references the trainer's params via
/// the `weight_ptrs` slice the caller passes in (positions 0..=3 are the
/// trunk encoder weights `W_s1, b_s1, W_s2, b_s2`).
#[allow(missing_debug_implementations)] // CublasGemmSet borrow is not Debug
pub struct StateEncoder<'a> {
batched: &'a CublasGemmSet,
}
impl<'a> StateEncoder<'a> {
/// Wrap an existing `BatchedForward` (`CublasGemmSet`) without taking
/// ownership of any GPU resources. The returned encoder is a borrow,
/// so multiple decoders can run after one encoder forward without
/// conflicting with the trainer's primary call sites.
pub fn new(batched: &'a CublasGemmSet) -> Self {
Self { batched }
}
/// Forward-pass through the trunk encoder (`h_s1` + `h_s2` GEMMs).
///
/// Reuses `BatchedForward::encoder_forward_only`, which is the same
/// dispatch the existing `forward_online_raw` runs for the encoder
/// portion. Behavior is byte-identical to the legacy path.
///
/// `weight_ptrs` is the full DQN weight pointer array; only positions
/// `[0, 1, 2, 3]` (W_s1 / b_s1 / W_s2 / b_s2) are read.
pub fn forward(
&self,
stream: &Arc<CudaStream>,
state_dev_ptr: u64,
weight_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
h_s1_dev_ptr: u64,
h_s2_dev_ptr: u64,
batch_size: usize,
) -> Result<EncoderOutput, MLError> {
if batch_size != self.batched.batch_size() {
return Err(MLError::ModelError(format!(
"StateEncoder::forward: batch_size {} != BatchedForward batch_size {}",
batch_size,
self.batched.batch_size(),
)));
}
self.batched.encoder_forward_only(
stream,
state_dev_ptr,
weight_ptrs,
h_s1_dev_ptr,
h_s2_dev_ptr,
)?;
Ok(EncoderOutput {
h_s2_dev_ptr,
h_s2_dim: self.batched.shared_h2(),
batch_size,
})
}
/// Output activation dimension (`shared_h2`).
pub fn h_s2_dim(&self) -> usize {
self.batched.shared_h2()
}
}

164
crates/ml/src/cuda_pipeline/value_decoder.rs Normal file
View File

@@ -0,0 +1,164 @@
//! Plan 4 Task 4 (E.4): explicit per-branch value-decoder boundary.
//!
//! `ValueDecoder` is a thin Rust API wrapper around one of the four DQN
//! advantage branches (Direction, Magnitude, Order, Urgency). Each
//! instance forwards an encoder's `h_s2` activation through its branch's
//! FC + adv-logits heads via
//! `BatchedForward::decoder_forward_only`. It owns no parameters of its
//! own — references the trainer's per-branch weights via the
//! `weight_ptrs` slice the caller passes in.
//!
//! Plan 2 D.3 (V_short + V_long horizon-decomposed value head) is
//! already landed on the value-head side; the decoder API surfaces the
//! `BranchDecoderOutput::v_short_dev_ptr` / `v_long_dev_ptr` fields so
//! Plan 4 Task 3 (multi-Q IQN) and Task 6 (auxiliary heads) can attach
//! cleanly. For decoders that do not own a value head, callers may set
//! the `v_short` / `v_long` pointers to 0 — they are not consumed by
//! `decoder_forward_only` itself today (they pass through into the
//! struct so downstream code sees a uniform shape).
//!
//! Design intent (E.4 spec):
//! - **Additive** to `BatchedForward::forward_online_raw` — that
//! monolithic dispatch keeps the existing multi-stream branch
//! fork/join. `ValueDecoder::forward` is for callers that need to
//! drive a single branch on a single stream (validation paths,
//! per-branch experiments, future migrations). No call-site
//! migration required for production training.
//! - Pure Rust API refactor: no kernel changes, no parameter changes,
//! no checkpoint break.
use std::sync::Arc;
use cudarc::driver::CudaStream;
use crate::MLError;
use crate::cuda_pipeline::batched_forward::CublasGemmSet;
/// Identifies one of the four DQN advantage branches. The `repr(usize)` is
/// chosen so `Branch::Dir as usize == 0` etc., matching the
/// `branch_idx` convention used throughout `BatchedForward`.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(usize)]
pub enum Branch {
Dir = 0,
Mag = 1,
Ord = 2,
Urg = 3,
}
impl Branch {
/// Branch index in `[0, 4)`, matching `BatchedForward`'s internal layout.
pub fn idx(self) -> usize { self as usize }
/// Short label, useful for diagnostics.
pub fn label(self) -> &'static str {
match self {
Branch::Dir => "dir",
Branch::Mag => "mag",
Branch::Ord => "ord",
Branch::Urg => "urg",
}
}
}
/// Output of one `ValueDecoder::forward` call.
///
/// Carries the device pointer for the per-action Q-distribution
/// (`q_per_action_dev_ptr`, shape `[B, branch_q_dim * num_atoms]` in
/// f32) plus the Plan 2 D.3 horizon-decomposed value pointers. The
/// value pointers are pass-through fields — the decoder does not own
/// them, the caller supplies them so downstream consumers see a uniform
/// type. Pass `0` for either pointer if the caller does not maintain a
/// horizon-decomposed value head for this branch.
#[derive(Debug, Clone, Copy)]
pub struct BranchDecoderOutput {
pub q_per_action_dev_ptr: u64,
pub q_per_action_dim: usize,
/// Number of distributional atoms per action (C51 support size).
pub num_atoms: usize,
pub v_short_dev_ptr: u64,
pub v_long_dev_ptr: u64,
pub batch_size: usize,
}
/// Plan 4 Task 4 (E.4) per-branch value-decoder façade.
///
/// One instance per branch. Borrows a `CublasGemmSet` (`BatchedForward`)
/// and dispatches `decoder_forward_only` for the configured branch.
#[allow(missing_debug_implementations)] // CublasGemmSet borrow is not Debug
pub struct ValueDecoder<'a> {
batched: &'a CublasGemmSet,
branch: Branch,
}
impl<'a> ValueDecoder<'a> {
/// Construct a decoder bound to a specific branch.
pub fn new(batched: &'a CublasGemmSet, branch: Branch) -> Self {
Self { batched, branch }
}
/// Branch this decoder targets.
pub fn branch(&self) -> Branch { self.branch }
/// Number of actions per branch (pre-num_atoms multiplication).
pub fn q_per_action_dim(&self) -> usize {
self.batched.branch_q_dim_internal(self.branch.idx())
}
/// Forward `h_s2` through this branch's decoder heads.
///
/// Reuses `BatchedForward::decoder_forward_only`, which mirrors the
/// per-branch loop body in `forward_online_raw` (sequential
/// single-stream variant). Behaviorally a no-op for callers that
/// keep using `forward_online_raw` directly.
///
/// Pointer ownership (caller-supplied):
/// - `h_s2_dev_ptr` : `[B, shared_h2]` encoder output (read-only)
/// - `mag_concat_dev_ptr` : `[B, shared_h2 + 3]` magnitude branch
/// wider input; pass `0` for legacy `[B, shared_h2]` input
/// - `branch_h_dev_ptr` : `[B, adv_h]` per-branch hidden scratch
/// - `q_per_action_dev_ptr`: `[B, branch_q_dim * num_atoms]` output
/// - `v_short_dev_ptr`,
/// `v_long_dev_ptr` : optional horizon-decomposed value
/// pointers, pass `0` if not used.
#[allow(clippy::too_many_arguments)]
pub fn forward(
&self,
stream: &Arc<CudaStream>,
h_s2_dev_ptr: u64,
mag_concat_dev_ptr: u64,
weight_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS],
branch_h_dev_ptr: u64,
q_per_action_dev_ptr: u64,
v_short_dev_ptr: u64,
v_long_dev_ptr: u64,
batch_size: usize,
) -> Result<BranchDecoderOutput, MLError> {
if batch_size != self.batched.batch_size() {
return Err(MLError::ModelError(format!(
"ValueDecoder::forward: batch_size {} != BatchedForward batch_size {}",
batch_size,
self.batched.batch_size(),
)));
}
self.batched.decoder_forward_only(
stream,
self.branch.idx(),
h_s2_dev_ptr,
mag_concat_dev_ptr,
weight_ptrs,
branch_h_dev_ptr,
q_per_action_dev_ptr,
)?;
Ok(BranchDecoderOutput {
q_per_action_dev_ptr,
q_per_action_dim: self.batched.branch_q_dim_internal(self.branch.idx()),
num_atoms: self.batched.num_atoms(),
v_short_dev_ptr,
v_long_dev_ptr,
batch_size,
})
}
}

8
docs/dqn-wire-up-audit.md
View File

@@ -86,6 +86,8 @@
| `cuda_pipeline/gpu_dqn_trainer.rs` (`GpuDqnTrainer`) | `fused_training.rs`, `trainer/mod.rs` | Wired | Core GPU-side DQN weight / forward / grad ops | — |
| `cuda_pipeline/fused_training.rs` wrapper (`FusedTrainingCtx`) | `trainer/training_loop.rs` | Wired | Wires all GPU ops into epoch loop | — |
| `cuda_pipeline/batched_forward.rs` | `gpu_dqn_trainer.rs` — forward graph node | Wired | Batched SGEMM forward pass | — |
| `cuda_pipeline/state_encoder.rs` (`StateEncoder`) | Borrows `CublasGemmSet` and routes through `BatchedForward::encoder_forward_only`; coexists with the monolithic `forward_online_raw` path used by `gpu_dqn_trainer.rs` | Wired (additive — coexists with `forward_online_raw` monolithic path) | Plan 4 Task 4 E.4 — explicit trunk-encoder boundary; attachment point for Plan 4 Task 1 (Full VSN, pre-encoder) and follow-up per-branch experiments | — |
| `cuda_pipeline/value_decoder.rs` (`ValueDecoder`, `Branch`, `BranchDecoderOutput`) | Borrows `CublasGemmSet` and routes through `BatchedForward::decoder_forward_only`; one instance per branch (Dir/Mag/Ord/Urg). Coexists with the multi-stream branch fork/join in `forward_online_raw` | Wired (additive — coexists with `forward_online_raw` monolithic path) | Plan 4 Task 4 E.4 — per-branch value-decoder boundary; attachment point for Plan 4 Task 3 (multi-Q IQN) and Task 6 (auxiliary heads). Surfaces Plan 2 D.3 V_short / V_long pointers via `BranchDecoderOutput` | — |
| `cuda_pipeline/batched_backward.rs` | `gpu_dqn_trainer.rs` — backward graph node | Wired | Batched SGEMM backward pass (KAN gate included) | — |
| `cuda_pipeline/shared_cublas_handle.rs` | `fused_training.rs`, `gpu_iqn_head.rs`, `gpu_iql_trainer.rs`, `gpu_attention.rs`, `gpu_curiosity_trainer.rs` (10 consumers) | Wired | Shared cuBLAS/cuBLASLt handle | — |
| `cuda_pipeline/cublas_algo_deterministic.rs` | `gpu_dqn_trainer.rs`, `gpu_curiosity_trainer.rs`, `gpu_iqn_head.rs` (7 consumers) | Wired | Deterministic cuBLASLt algo selection | — |
@@ -273,14 +275,16 @@ Plan 2 Task 6B D.3 (2026-04-24): IQL value head widened from 1 to 2 outputs (V_s
Plan 4 Task 5 Mode A E.5 (2026-04-24): `attention_focus_ema_kernel.cu` + `AttentionMonitor` added. 3 new ISV slots [87] VSN_MAG_EMA, [88] VSN_DIR_EMA, [89] MAMBA2_RETENTION_EMA tail-appended; fingerprint shifted [85..87) → [90..92); `ISV_TOTAL_DIM` 87 → 92. Light, ISV-diagnostic-only; no model parameters added, no checkpoint break. Single-thread cold-path EMA kernel takes 3 host scalars by value at the per-epoch HEALTH_DIAG site. New `mamba2_retention_mean(batch_size)` accessor on `GpuDqnTrainer`/`FusedTrainingCtx` mirrors `per_branch_vsn_mean` (cold-path host-side dtoh + mean abs of `mamba2_h_enriched`). Adaptive α convention matches Plan 3 producers. All three slots FoldReset to 0.0 (mirror constructor cold-start). 2 new Wired rows. Mode B (full per-feature-group VSN ISV) blocks on Plan 4 Task 1 (E.1).
Plan 4 Task 4 E.4 (2026-04-24): `state_encoder.rs` + `value_decoder.rs` Rust API split. PURE Rust API refactor — no kernel changes, no parameter changes, no checkpoint break, no new ISV slot. New types: `StateEncoder` (wraps trunk encoder `h_s1` + `h_s2`), `ValueDecoder` (per-branch FC + adv-logits, one instance per `Branch::{Dir,Mag,Ord,Urg}`), `EncoderOutput`, `BranchDecoderOutput` (carries `q_per_action_dev_ptr` + Plan 2 D.3 `v_short_dev_ptr` / `v_long_dev_ptr` pass-through fields). `BatchedForward` gains `encoder_forward_only(...)` and `decoder_forward_only(branch_idx, ...)` helpers, both lossless extractions of the existing dispatch sequence in `forward_online_raw`. The monolithic `forward_online_raw` stays callable and now routes through `encoder_forward_only` internally for the trunk portion (multi-stream branch fork/join unchanged). The new types are ADDITIVE attachment points for Plan 4 Task 1 (Full VSN, pre-encoder), Task 3 (multi-Q IQN, decoder), and Task 6 (auxiliary heads, decoder peer). 2 new Wired rows.
| Classification | Count |
|---|---|
| Wired | 86 |
| Wired | 88 |
| Partial | 10 |
| Orphan (held for follow-up) | 3 |
| Ghost | 0 |
| OUT-of-DQN-scope | 17 |
| **Total** | **114** |
| **Total** | **116** |
The 3 remaining Orphan rows are:
- `cuda_pipeline/gpu_statistics.rs` + `statistics_kernel.cu` — held for Plan 2 D.2 wire-or-delete decision.