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>
fbc299fa27