feat(sp22-vnext): Phase C-1 — K=3 softmax → per-env 3-slot cache (producer)
First half of Phase C. Lands the producer side of the K=3 trade- outcome aux head's state bridge: new kernel populates a per-env 3-slot cache from the K=3 softmax tile every rollout step. The consumer side (state gather reading from this cache → state slots [121..124)) lands in Phase C-2. Mirrors the K=2 head's existing aux_softmax_to_per_env_kernel exactly at K=3: - K=2: prev_aux_dir_prob[env] = 2*softmax[env, 1] - 1 (recentered) - K=3: prev_aux_outcome_probs[env, k] = softmax[env, k] for k in [0, 3) Changes: - state_layout.rs: 3 new constants AUX_OUTCOME_PROFIT_INDEX = 121, AUX_OUTCOME_STOP_INDEX = 122, AUX_OUTCOME_TIMEOUT_INDEX = 123. PROFIT_INDEX aliases AUX_DIR_PROB_INDEX (same value, different semantic). Phase C-2 flips slot 121's meaning from K=2's recentered p_up to K=3's p_Profit. - aux_outcome_softmax_to_per_env_kernel.cu: new kernel + cubin. - gpu_dqn_trainer.rs: new SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN embed. - gpu_experience_collector.rs: 2 new struct fields (cache buffer + kernel handle); cubin load + alloc in constructor; struct-init; per-step launch in rollout loop after K=3 forward. - build.rs: kernel registered. Encoding shift K=2 → K=3: K=2 used recentered [-1, +1] to match "no signal = 0" baseline of every other slot. K=3 keeps raw softmax probabilities [0, 1]. Cold-start sentinel 0.0 for all 3 slots = "no prediction yet" (mask). The 3-slot natural distribution is more informative than a scalar. Dead-code status: producer populates cache every step but experience_state_gather doesn't read from it yet — state slot 121 still receives K=2's prev_aux_dir_prob write. Phase C-2 swaps the state gather's source from K=2 cache to K=3 cache (3-slot write). Why split C into C-1 + C-2: experience_state_gather is a hot-path kernel with many consumers. Updating it touches training collector, eval-side backtest evaluator, Rust launcher arg list. C-2 lands that as an atomic state-semantic flip; C-1 lands the GPU-side scaffolding independently so the producer chain can be validated first. Verification: - cargo check -p ml clean. - cargo test -p ml --lib → 1016/0 green. Audit: docs/dqn-wire-up-audit.md Phase C-1 section. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -841,6 +841,15 @@ fn main() {
|
||||
// this commit — Phase B5b allocates buffers + wires launches +
|
||||
// flips SH2_TOTAL=262 in forward/backward.
|
||||
"aux_to_input_concat_kernel.cu",
|
||||
// SP22 H6 vNext Phase C (2026-05-14): per-env 3-slot softmax
|
||||
// extractor for the K=3 trade-outcome aux head. Copies
|
||||
// aux_outcome_softmax [n_envs, K=3] → prev_aux_outcome_probs
|
||||
// [n_envs, 3]. Consumed at the NEXT step's state assembly via
|
||||
// experience_state_gather (Phase C-2 follow-up). For Phase C
|
||||
// (this commit) the producer fires every step but no consumer
|
||||
// reads — dead-code scaffolding parallel to the K=2 head's
|
||||
// aux_softmax_to_per_env_kernel.
|
||||
"aux_outcome_softmax_to_per_env_kernel.cu",
|
||||
// SP22 H6 (2026-05-12): per-env p_up extractor that copies
|
||||
// `aux_softmax[env, 1]` → `prev_aux_dir_prob[env]` after each aux
|
||||
// forward. The cache buffer is read by `experience_state_gather`
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
// crates/ml/src/cuda_pipeline/aux_outcome_softmax_to_per_env_kernel.cu
|
||||
//
|
||||
// SP22 H6 vNext Phase C (2026-05-14): aux→policy state bridge — per-env
|
||||
// trade-outcome probability extractor.
|
||||
//
|
||||
// Role
|
||||
// ────
|
||||
// At rollout step t, after the K=3 trade-outcome aux head's softmax
|
||||
// forward produces the per-env tile `aux_outcome_softmax [n_envs, K=3]`,
|
||||
// this kernel copies all three probability columns into the per-env
|
||||
// cache buffer `prev_aux_outcome_probs [n_envs, 3]`. The cache is read
|
||||
// by `experience_state_gather` at step t+1 to populate state slots:
|
||||
//
|
||||
// state[AUX_OUTCOME_PROFIT_INDEX = 121] ← prev_aux_outcome_probs[env, 0]
|
||||
// state[AUX_OUTCOME_STOP_INDEX = 122] ← prev_aux_outcome_probs[env, 1]
|
||||
// state[AUX_OUTCOME_TIMEOUT_INDEX = 123] ← prev_aux_outcome_probs[env, 2]
|
||||
//
|
||||
// Replaces the K=2 head's single-slot bridge
|
||||
// (`aux_softmax_to_per_env_kernel` → `prev_aux_dir_prob` → state[121] =
|
||||
// `2*p_up - 1`). The K=3 head's 3-slot bridge supersedes that semantic
|
||||
// at step t+1 once the state gather consumer is updated (Phase C-2,
|
||||
// follow-up commit). For Phase C (this commit) the buffer is populated
|
||||
// but no consumer reads it yet — dead-code scaffolding.
|
||||
//
|
||||
// Encoding
|
||||
// ────────
|
||||
// Slots [121..124) hold raw softmax probabilities in [0, 1] (sum to 1
|
||||
// per row). Unlike the K=2 head's slot 121 which used a recentered
|
||||
// `2*p - 1` encoding to land in [-1, +1], the K=3 head's three slots
|
||||
// keep the natural [0, 1] range:
|
||||
// - No structural bound violation (softmax outputs are inherently
|
||||
// [0, 1] and sum to 1)
|
||||
// - Three-slot probabilistic representation is more informative than
|
||||
// a single "p_up − p_down" scalar
|
||||
// - Cold-start sentinel: all three slots = 0.0. The policy reads
|
||||
// "no signal" for all three classes, which is the conservative
|
||||
// baseline (vs. the K=2 head's 0.0 = neutral encoding). After
|
||||
// first rollout step the producer overwrites with real K=3 probs.
|
||||
//
|
||||
// Discipline
|
||||
// ──────────
|
||||
// - No atomicAdd (no reduction, just a strided gather)
|
||||
// - GPU-only per `feedback_cpu_is_read_only.md`
|
||||
// - Stream-ordered with the K=3 forward producer (same-stream barrier
|
||||
// suffices — no explicit sync needed)
|
||||
//
|
||||
// Launch
|
||||
// ──────
|
||||
// grid = ((n_envs + 255) / 256), block = (256). One thread per env.
|
||||
// Each thread writes 3 consecutive output elements (K=3 unrolled).
|
||||
//
|
||||
// Cost
|
||||
// ────
|
||||
// Three coalesced 4-byte loads + three 4-byte stores per env. n_envs is
|
||||
// O(1k); ~1-2µs total. Negligible vs the existing K=2 sibling.
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
extern "C" __global__ void aux_outcome_softmax_to_per_env_kernel(
|
||||
/* Per-env K=3 softmax tile [n_envs, K=3]. Written by
|
||||
* `aux_trade_outcome_forward` during the captured rollout-step
|
||||
* graph; valid after same-stream sync. Layout row-major; this
|
||||
* kernel reads `aux_outcome_softmax[env * K + 0/1/2]` per env. */
|
||||
const float* __restrict__ aux_outcome_softmax,
|
||||
/* Per-env cache buffer [n_envs, K=3] row-major. Written here, read
|
||||
* by `experience_state_gather` (Phase C-2) at the next step's
|
||||
* state assembly. Slot 0 = p_Profit, slot 1 = p_Stop, slot 2 =
|
||||
* p_Timeout per the K=3 outcome label convention from
|
||||
* `trade_outcome_label_kernel`. */
|
||||
float* __restrict__ prev_aux_outcome_probs,
|
||||
/* Number of envs (rollout-time batch width = `alloc_episodes`). */
|
||||
int n_envs,
|
||||
/* Softmax K (= AUX_OUTCOME_K = 3 in Phase A3+). Kernel reads slots
|
||||
* [0, K). Passed runtime so the kernel signature stays stable if
|
||||
* the head's K flips in a future redesign. */
|
||||
int K
|
||||
) {
|
||||
int env = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (env >= n_envs) return;
|
||||
|
||||
/* Phase C (2026-05-14): write raw softmax probabilities [0, 1]
|
||||
* into 3 per-env slots. The state gather consumer (Phase C-2)
|
||||
* will copy these into state[121..124). Unlike the K=2 head's
|
||||
* recentered encoding, the K=3 head's natural softmax range is
|
||||
* already in [0, 1] and the three-slot representation preserves
|
||||
* the full categorical distribution.
|
||||
*
|
||||
* The kernel writes K slots (passed runtime) so the loop tolerates
|
||||
* future K flips; in production K=3 always so the loop trips
|
||||
* exactly 3 times. Bounds: K can be any positive value, but the
|
||||
* cache buffer must be sized [n_envs, K] by the caller. */
|
||||
const size_t src_base = (size_t)env * (size_t)K;
|
||||
const size_t dst_base = (size_t)env * (size_t)K;
|
||||
for (int k = 0; k < K; ++k) {
|
||||
prev_aux_outcome_probs[dst_base + (size_t)k] =
|
||||
aux_outcome_softmax[src_base + (size_t)k];
|
||||
}
|
||||
}
|
||||
@@ -685,6 +685,14 @@ pub(crate) static SP14_Q_DISAGREEMENT_CUBIN: &[u8] =
|
||||
pub(crate) static SP22_AUX_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/aux_softmax_to_per_env_kernel.cubin"));
|
||||
|
||||
/// SP22 H6 vNext Phase C (2026-05-14): per-env K=3 trade-outcome softmax
|
||||
/// extractor cubin. Loaded by `GpuExperienceCollector` (and eventually
|
||||
/// `GpuBacktestEvaluator` in the eval-side mirror). Producer of the 3-
|
||||
/// slot per-env cache that Phase C-2 wires into state[121..124).
|
||||
#[allow(dead_code)]
|
||||
pub(crate) static SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN: &[u8] =
|
||||
include_bytes!(concat!(env!("OUT_DIR"), "/aux_outcome_softmax_to_per_env_kernel.cubin"));
|
||||
|
||||
/// SP22 H6 Phase 3 α structural-prior init (atom-shift design 2026-05-13).
|
||||
/// One-time GPU init kernel that writes the per-action prior
|
||||
/// `[-0.5, 0.0, +0.5, 0.0]` for `[Short, Hold, Long, Flat]` into the
|
||||
|
||||
@@ -1297,6 +1297,32 @@ pub struct GpuExperienceCollector {
|
||||
/// consumer.
|
||||
exp_aux_to_label_per_sample: cudarc::driver::CudaSlice<i32>,
|
||||
|
||||
/// SP22 H6 vNext Phase C (2026-05-14) — per-env K=3 trade-outcome
|
||||
/// softmax cache. Sized `[alloc_episodes × AUX_OUTCOME_K=3]` f32,
|
||||
/// row-major: env-major then class. Written end-of-step by
|
||||
/// `aux_outcome_softmax_to_per_env_kernel` (Phase C kernel) from
|
||||
/// `exp_aux_to_softmax_buf`. Read start-of-next-step by
|
||||
/// `experience_state_gather` (Phase C-2 wireup) to populate state
|
||||
/// slots `[AUX_OUTCOME_PROFIT_INDEX..AUX_OUTCOME_TIMEOUT_INDEX+1)`
|
||||
/// = `[121, 122, 123]`.
|
||||
///
|
||||
/// Replaces the K=2 head's single-slot `prev_aux_dir_prob` semantic
|
||||
/// at state slot 121. Both buffers coexist in this commit (K=2 head
|
||||
/// remains active; its prev_aux_dir_prob still feeds state[121] via
|
||||
/// the current state gather). Phase C-2 flips the state gather to
|
||||
/// read from this 3-slot K=3 buffer instead.
|
||||
///
|
||||
/// Cold-start sentinel = 0.0 across all 3 slots. The state gather
|
||||
/// reads "no Profit/Stop/Timeout signal" on step 0, then this
|
||||
/// kernel writes real K=3 softmax probs for step 1+ state assembly.
|
||||
prev_aux_outcome_probs: cudarc::driver::CudaSlice<f32>,
|
||||
/// SP22 H6 vNext Phase C (2026-05-14) — kernel handle for
|
||||
/// `aux_outcome_softmax_to_per_env_kernel`. Loaded from the Phase
|
||||
/// C cubin. Launched on the collector's stream right after the
|
||||
/// K=3 `aux_trade_outcome_forward` populates
|
||||
/// `exp_aux_to_softmax_buf`.
|
||||
exp_aux_outcome_softmax_to_per_env_kernel: CudaFunction,
|
||||
|
||||
/// SP22 H6 Phase 2 (2026-05-12) — per-env aux directional probability
|
||||
/// cache, RECENTERED encoding. Sized `[alloc_episodes]` f32 (one slot
|
||||
/// per env). Written end-of-step by `exp_aux_softmax_to_per_env_kernel`
|
||||
@@ -2764,6 +2790,31 @@ impl GpuExperienceCollector {
|
||||
"sp22-vnext B4b: alloc exp_aux_to_label_per_sample ({} i32): {e}",
|
||||
alloc_episodes * alloc_timesteps
|
||||
)))?;
|
||||
|
||||
// SP22 H6 vNext Phase C (2026-05-14): K=3 per-env softmax cache
|
||||
// + producer kernel handle. Sized [alloc_episodes × AUX_OUTCOME_K=3]
|
||||
// f32. Cold-start sentinel 0.0 across all 3 slots = "no
|
||||
// Profit/Stop/Timeout signal" (the consumer state gather will
|
||||
// read this baseline on step 0; subsequent steps populated
|
||||
// from the K=3 forward's softmax).
|
||||
let prev_aux_outcome_probs = stream
|
||||
.alloc_zeros::<f32>(alloc_episodes * AUX_OUTCOME_K)
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-vnext C: alloc prev_aux_outcome_probs ({} f32): {e}",
|
||||
alloc_episodes * AUX_OUTCOME_K
|
||||
)))?;
|
||||
let exp_aux_outcome_softmax_to_per_env_kernel = {
|
||||
use super::gpu_dqn_trainer::SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN;
|
||||
let m = stream.context()
|
||||
.load_cubin(SP22_AUX_OUTCOME_SOFTMAX_TO_PER_ENV_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-vnext C: aux_outcome_softmax_to_per_env cubin (collector): {e}"
|
||||
)))?;
|
||||
m.load_function("aux_outcome_softmax_to_per_env_kernel")
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"sp22-vnext C: aux_outcome_softmax_to_per_env_kernel load (collector): {e}"
|
||||
)))?
|
||||
};
|
||||
// SP22 H6 (2026-05-12): aux→policy state bridge. Load the per-env
|
||||
// softmax extractor (`aux_softmax_to_per_env_kernel`) + the shared
|
||||
// `fill_f32` from the epsilon_greedy cubin (for 0.5 sentinel init at
|
||||
@@ -3354,6 +3405,12 @@ impl GpuExperienceCollector {
|
||||
exp_aux_to_softmax_buf,
|
||||
exp_aux_to_label_buf,
|
||||
exp_aux_to_label_per_sample,
|
||||
// SP22 H6 vNext Phase C (2026-05-14): K=3 trade-outcome
|
||||
// per-env softmax cache + producer kernel handle (Phase C
|
||||
// scaffolding — populates `prev_aux_outcome_probs` every
|
||||
// step; state gather consumer wired in Phase C-2).
|
||||
prev_aux_outcome_probs,
|
||||
exp_aux_outcome_softmax_to_per_env_kernel,
|
||||
// SP22 H6 (2026-05-12): aux→policy state bridge — per-env p_up
|
||||
// cache + extractor kernel + fill_f32 handle for 0.5 sentinel.
|
||||
prev_aux_dir_prob,
|
||||
@@ -5887,6 +5944,40 @@ impl GpuExperienceCollector {
|
||||
self.exp_aux_to_logits_buf.raw_ptr(),
|
||||
self.exp_aux_to_softmax_buf.raw_ptr(),
|
||||
)?;
|
||||
|
||||
// ── SP22 H6 vNext Phase C (2026-05-14): K=3 softmax →
|
||||
// per-env 3-slot cache. Mirrors the K=2 sibling's
|
||||
// `aux_softmax_to_per_env_kernel` launch (Phase 2) but
|
||||
// writes 3 slots instead of 1. Reads
|
||||
// `exp_aux_to_softmax_buf [n_envs, 3]` (just produced
|
||||
// by the K=3 forward above on the same stream — stream-
|
||||
// implicit ordering); writes `prev_aux_outcome_probs
|
||||
// [n_envs, 3]`. Phase C-2 wires `experience_state
|
||||
// _gather` to read this cache and populate state slots
|
||||
// [121..124).
|
||||
{
|
||||
let n_envs_i32 = n as i32;
|
||||
let k_i32 = AUX_OUTCOME_K as i32;
|
||||
let blocks = ((n as u32).div_ceil(256)).max(1);
|
||||
let softmax_ptr = self.exp_aux_to_softmax_buf.raw_ptr();
|
||||
let cache_ptr = self.prev_aux_outcome_probs.raw_ptr();
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.exp_aux_outcome_softmax_to_per_env_kernel)
|
||||
.arg(&softmax_ptr)
|
||||
.arg(&cache_ptr)
|
||||
.arg(&n_envs_i32)
|
||||
.arg(&k_i32)
|
||||
.launch(LaunchConfig {
|
||||
grid_dim: (blocks, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
})
|
||||
.map_err(|e| MLError::ModelError(format!(
|
||||
"aux_outcome_softmax_to_per_env_kernel t={t}: {e}"
|
||||
)))?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── 3. Convert logits → expected Q-values ───────────────────
|
||||
|
||||
Reference in New Issue
Block a user