Files
foxhunt/crates/ml-core/src/state_layout.rs
jgrusewski 3286dc7dee 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>
2026-05-14 02:26:39 +02:00

279 lines
18 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
/// Canonical state layout — single source of truth.
///
/// Training and validation MUST produce identical state vectors.
/// All systems use these constants instead of configurable state_dim fields.
///
/// Layout (grouped by update frequency, fast-changing first):
/// [0..42) Market features (OHLCV + technicals)
/// [42..74) OFI (32-dim order flow from MBP-10 — see OFI slot map below)
/// [74..90) TLOB (16-dim attention-compressed OFI — D.8 Plan 2 Task 6C; zeros on assembly, GpuTlob::forward overwrites post-gather)
/// [90..106) MTF (4 lookbacks x 4 features)
/// [106..114) Portfolio base (8 features)
/// [114..121) Plan/ISV (7 features — D.6 Plan 2 Task 6A added remaining_fraction at index 6)
/// [121..128) Zero padding (7 zeros for 8-alignment; D.8 Plan 2 Task 6C reclaimed from 16 → 7)
///
/// OFI slot map (OFI_DIM = 32, indices relative to OFI_START):
/// [0..8) raw OFI — ofi_level1, ofi_level5, depth_imbalance, vpin,
/// kyle_lambda, bid_slope, ask_slope, trade_imbalance
/// [8..16) lag-1 deltas of slots [0..8)
/// [16] book_aggression (MBP-10 center-of-mass asymmetry)
/// [17] log_bar_duration (ln(bar_secs) / 10)
/// [18] ofi_acceleration (EMA of ΔOFI_L1 — MicrostructureState)
/// [19] toxicity_gradient (EMA of ΔVPIN — MicrostructureState)
/// [20] ofi_trajectory (online-linreg slope of OFI_L1 — MicrostructureState)
/// [21] realized_variance (Σ log-returns² — MicrostructureState)
/// [22] hawkes_intensity (Hawkes trade process — MicrostructureState)
/// [23] book_pressure (weighted exp(-0.3·l) 10-level — MicrostructureState)
/// [24] spread_dynamics ((max-min)/mean — MicrostructureState)
/// [25] aggression_ratio (trades_at_ask / total — MicrostructureState)
/// [26] queue_depletion_asymmetry (EMA bid/ask depletion — MicrostructureState)
/// [27] order_count_flux (ct_inc / (ct_inc + ct_dec) — MicrostructureState)
/// [28] intra_bar_momentum (sign-weighted half-bar product — MicrostructureState)
/// [29] regime_score (sigmoid of aggr·spread·thin_book — MicrostructureState)
/// [30] order_count_imbalance ((Σbid_ct Σask_ct) / Σ(bid_ct + ask_ct) — Mbp10Snapshot)
/// [31] microprice_residual ((weighted_mid mid) / mid — Mbp10Snapshot)
// D.8 Plan 2 Task 6C: TLOB attention (16-dim) inserted after OFI.
// Layout: Market(42)+OFI(32)+TLOB(16)+MTF(16)+Portfolio(8)+PlanISV(7)+Padding(7) = 128.
// STATE_DIM 112→128 (matches STATE_DIM_PADDED; padding reclaimed to 7 for 8-alignment).
pub const STATE_DIM: usize = 128;
pub const STATE_DIM_PADDED: usize = 128; // pad128(STATE_DIM) for cuBLAS K-tile alignment
pub const MARKET_DIM: usize = 42;
pub const OFI_DIM: usize = 32;
pub const TLOB_DIM: usize = 16;
pub const MTF_DIM: usize = 16;
pub const PORTFOLIO_BASE_DIM: usize = 8;
pub const PORTFOLIO_PLAN_DIM: usize = 7;
pub const PADDING_DIM: usize = 7;
pub const MARKET_START: usize = 0;
pub const OFI_START: usize = MARKET_START + MARKET_DIM; // 42
pub const TLOB_START: usize = OFI_START + OFI_DIM; // 74
pub const MTF_START: usize = TLOB_START + TLOB_DIM; // 90
pub const PORTFOLIO_START: usize = MTF_START + MTF_DIM; // 106
pub const PLAN_ISV_START: usize = PORTFOLIO_START + PORTFOLIO_BASE_DIM; // 114
pub const PADDING_START: usize = PLAN_ISV_START + PORTFOLIO_PLAN_DIM; // 121
// SP22 H6 Phase 2 (2026-05-12) — aux directional probability from previous
// step, RECENTERED. Lives in the first padding slot. PADDING_DIM stays 7
// (no STATE_DIM change); `assemble_state` writes aux_dir_prob to
// [AUX_DIR_PROB_INDEX] and zeros the remaining 6 padding slots
// [PADDING_START+1..STATE_DIM).
// Encoding: `2*p_up - 1` ∈ [-1, +1] (structurally bounded — softmax
// components in [0, 1] sum to 1). +1 = up with full conviction;
// -1 = down with full conviction; 0 = neutral.
// Sentinel = 0.0 at cold-start and FoldReset (matches every other state
// slot's "no signal = 0" baseline per `pearl_first_observation_bootstrap`).
pub const AUX_DIR_PROB_INDEX: usize = PADDING_START; // 121
// SP22 H6 vNext Phase C (2026-05-14) — trade-outcome aux head's K=3
// softmax probabilities, raw [0, 1] range. Replaces the K=2 head's
// single-slot recentered `2*p_up - 1` at slot 121 (aliased — same index
// value, different semantic). The K=3 head's three-slot representation
// is more informative than a single directional scalar:
// - p_Profit: probability the trade closes via profit_target hit
// - p_Stop: probability the trade closes via stop_loss hit
// - p_Timeout:probability the trade closes via target_bars elapsed
// Softmax outputs are inherently in [0, 1] and sum to 1 per row, so no
// recentering needed for "no signal = 0" baseline (cold-start all 3
// slots = 0.0 means "all-mask, no trade-close prediction yet" — also
// the FoldReset sentinel).
//
// Phase C (this commit): constants defined + producer kernel
// (`aux_outcome_softmax_to_per_env_kernel`) populates the per-env
// cache. Phase C-2 (follow-up): `experience_state_gather` reads from
// the K=3 cache and writes these 3 slots, deprecating the K=2 head's
// single-slot input to the policy.
pub const AUX_OUTCOME_PROFIT_INDEX: usize = PADDING_START; // 121
pub const AUX_OUTCOME_STOP_INDEX: usize = PADDING_START + 1; // 122
pub const AUX_OUTCOME_TIMEOUT_INDEX: usize = PADDING_START + 2; // 123
const _: () = assert!(PADDING_START + PADDING_DIM == STATE_DIM, "layout must sum to STATE_DIM");
const _: () = assert!(STATE_DIM % 8 == 0, "STATE_DIM must be 8-aligned for tensor cores");
const _: () = assert!(STATE_DIM_PADDED % 128 == 0, "STATE_DIM_PADDED must be 128-aligned for cuBLAS");
const _: () = assert!(STATE_DIM <= STATE_DIM_PADDED, "STATE_DIM must fit within padded width");
// ────────────────────────────────────────────────────────────────────────────
// Mirrors of crates/ml/src/cuda_pipeline/state_layout.cuh PS_* constants.
// Task 4A of Plan 1. Invariant 8.
// ────────────────────────────────────────────────────────────────────────────
pub const PS_POSITION: usize = 0; // current contract position (signed)
pub const PS_CASH: usize = 1; // cash balance
pub const PS_PORTFOLIO_VALUE: usize = 2; // mark-to-market total (cash + pos*price)
pub const PS_DSR_A: usize = 3; // EMA of realised trade returns (DSR)
pub const PS_DSR_B: usize = 4; // EMA of squared trade returns (DSR)
pub const PS_DSR_TRADE_COUNT: usize = 5; // completed trades counter (DSR warmup)
pub const PS_PREV_CLOSE: usize = 6; // prev bar raw_close for DSR per-bar returns
pub const PS_PEAK_EQUITY: usize = 7; // high-water mark (init to initial_capital)
pub const PS_FLAT_COUNTER: usize = 8; // consecutive flat steps
pub const PS_PREV_EQUITY: usize = 9; // equity at previous step
pub const PS_HOLD_TIME: usize = 10; // consecutive steps with position
pub const PS_REALIZED_PNL: usize = 11; // cumulative realised PnL
pub const PS_ENTRY_PRICE: usize = 12; // price when trade was entered
pub const PS_TRADE_START_PNL: usize = 13; // realized_pnl snapshot at trade entry
pub const PS_KELLY_WIN_COUNT: usize = 14; // Kelly: number of profitable trade exits
pub const PS_KELLY_LOSS_COUNT: usize = 15; // Kelly: number of losing trade exits
pub const PS_KELLY_SUM_WINS: usize = 16; // Kelly: cumulative profit from winners
pub const PS_KELLY_SUM_LOSSES: usize = 17; // Kelly: cumulative |loss| from losers
pub const PS_KELLY_SUM_RETURNS: usize = 18; // Kelly: cumulative net returns (for mu)
pub const PS_KELLY_SUM_SQ_RETURNS: usize = 19; // Kelly: cumulative squared returns (sigma^2)
pub const PS_INTRA_TRADE_MAX_DD: usize = 20; // worst unrealised drawdown during trade (v8)
pub const PS_INTRA_TRADE_MAX_PNL: usize = 21; // best unrealised P&L during trade (hindsight)
pub const PS_PREV_MID: usize = 22; // prev bar MBP-10 mid-price
pub const PS_PLAN_TARGET_BARS: usize = 23; // plan: max hold bars (>0.5 = plan active)
pub const PS_PLAN_PROFIT_TARGET: usize = 24; // plan: raw profit threshold
pub const PS_PLAN_STOP_LOSS: usize = 25; // plan: raw stop threshold
pub const PS_PLAN_SCALE_AGGRESSION: usize = 26; // plan: position ramp speed
pub const PS_PLAN_CONVICTION: usize = 27; // plan: conviction at entry [0,1]
pub const PS_PLAN_ASYMMETRY: usize = 28; // plan: profit/stop asymmetry
pub const PS_PLAN_ENTRY_REGIME: usize = 29; // plan: regime_stability at entry
pub const PS_OFI_PREV_BASE: usize = 30; // first OFI prev-bar slot (stride 8: [30..37])
pub const PS_PEAK_PNL_BAR: usize = 38; // Plan 3 C.4: hold_time at which MAX_PNL was hit (temporal timing bonus)
pub const PS_INTRA_TRADE_MIN_PNL: usize = 39; // Plan 3 D.4a: worst unrealised PnL during trade (symmetric to MAX_PNL at 21)
pub const PS_REGIME_SHIFT_BAR: usize = 40; // Plan 3 D.4b: hold_time at first regime-shift detection (0 if no shift)
pub const PS_PRE_ENTRY_CONVICTION_EMA: usize = 41; // Plan 3 D.4c: EMA of conviction_core during Flat bars
pub const PS_PRE_ENTRY_CONVICTION_VAR_EMA: usize = 42; // Plan 3 D.4c: EMA of (conviction_core mean)² during Flat bars
pub const PS_STRIDE: usize = 43; // total portfolio state stride (== PORTFOLIO_STRIDE), grown 41→43 by Plan 3 D.4c conviction consistency
// ────────────────────────────────────────────────────────────────────────────
// Mirrors of PLAN_ISV_* constants (plan_isv[0..PLAN_ISV_DIM=7)).
// Task 4B of Plan 1. Invariant 8.
// D.6 Plan 2 Task 6A: added PLAN_ISV_REMAINING_FRACTION at index 6.
// ────────────────────────────────────────────────────────────────────────────
pub const PLAN_ISV_PROGRESS: usize = 0; // hold_time / target_bars (clamped to 2.0)
pub const PLAN_ISV_PNL_VS_TARGET: usize = 1; // unrealized / (profit_target × equity)
pub const PLAN_ISV_PNL_VS_STOP: usize = 2; // -unrealized / (stop_loss × equity)
pub const PLAN_ISV_ENTRY_CONVICTION: usize = 3; // conviction at trade entry [0, 1]
pub const PLAN_ISV_CONVICTION_DRIFT: usize = 4; // current conviction entry conviction
pub const PLAN_ISV_REGIME_SHIFT: usize = 5; // |isv[11] entry_regime_stability|
pub const PLAN_ISV_REMAINING_FRACTION: usize = 6; // max(0, min(1, (target_bars hold_time) / target_bars)) when active else 0
pub const PLAN_ISV_DIM: usize = 7;
// ────────────────────────────────────────────────────────────────────────────
// Mirrors of PLAN_PARAM_* constants (plan_params[0..PLAN_PARAM_DIM=6)).
// Task 4C of Plan 1. Invariant 8.
// ────────────────────────────────────────────────────────────────────────────
pub const PLAN_PARAM_TARGET_BARS: usize = 0; // max hold bars
pub const PLAN_PARAM_PROFIT_TARGET: usize = 1; // raw profit threshold %
pub const PLAN_PARAM_STOP_LOSS: usize = 2; // raw stop loss threshold %
pub const PLAN_PARAM_SCALE_AGGRESSION: usize = 3; // position ramp speed
pub const PLAN_PARAM_CONVICTION: usize = 4; // position size fraction [0, 1]
pub const PLAN_PARAM_ASYMMETRY: usize = 5; // profit/stop ratio
pub const PLAN_PARAM_DIM: usize = 6;
// ────────────────────────────────────────────────────────────────────────────
// Mirrors of BRANCH_* constants (4-branch factored action).
// Task 4D of Plan 1. Invariant 8.
// ────────────────────────────────────────────────────────────────────────────
pub const BRANCH_DIR: usize = 0; // direction branch (4 actions: Short/Hold/Long/Flat)
pub const BRANCH_MAG: usize = 1; // magnitude branch (3 actions: Quarter/Half/Full)
pub const BRANCH_ORD: usize = 2; // order-type branch (3 actions)
pub const BRANCH_URG: usize = 3; // urgency branch (3 actions)
pub const NUM_BRANCHES: usize = 4;
// ────────────────────────────────────────────────────────────────────────────
// Mirrors of DIR_* constants (direction action sub-indices).
// Task 4E of Plan 1. Invariant 8.
// ────────────────────────────────────────────────────────────────────────────
pub const DIR_SHORT: usize = 0; // open/maintain short position
pub const DIR_HOLD: usize = 1; // keep current position (no-op)
pub const DIR_LONG: usize = 2; // open/maintain long position
pub const DIR_FLAT: usize = 3; // close all positions to zero
pub const NUM_DIRECTIONS: usize = 4;
// ────────────────────────────────────────────────────────────────────────────
// Mirrors of MAG_* constants (magnitude action sub-indices).
// Task 4F of Plan 1. Invariant 8.
// ────────────────────────────────────────────────────────────────────────────
pub const MAG_QUARTER: usize = 0; // 0.25× max_position
pub const MAG_HALF: usize = 1; // 0.50× max_position
pub const MAG_FULL: usize = 2; // 1.00× max_position
pub const NUM_MAGNITUDES: usize = 3;
// ────────────────────────────────────────────────────────────────────────────
// Feature-group ranges — Plan 4 Task 1A (E.1 VSN prerequisite).
//
// Mirrors the `SL_*_GROUP_BEGIN/END` macros in
// `crates/ml/src/cuda_pipeline/state_layout.cuh`. Ranges are half-open
// `[begin, end)`. The last group (`plan_isv`) ends at `PADDING_START` —
// the 7-element zero-pad block is *not* a feature group.
//
// Consumers:
// * Plan 4 Task 1 (E.1) Variable Selection Network — gates each group's
// features by a learned softmax-over-groups mask before the trunk.
// * Plan 4 Task 5 (E.5) attention-weight ISV diagnostics.
// * Plan 4 Task 4 (E.4) encoder-decoder separation — group-aware
// `StateEncoder` interface.
// ────────────────────────────────────────────────────────────────────────────
pub const SL_NUM_FEATURE_GROUPS: usize = 6;
/// Largest per-group dimension (currently `MARKET_DIM = 42`). Used for kernel-
/// side fixed-size scratch arrays so every group fits a uniform stride.
pub const SL_MAX_FEATURE_GROUP_DIM: usize = MARKET_DIM;
/// `[begin, end)` ranges per feature group, in declaration order.
/// Order matches `FEATURE_GROUP_NAMES`.
pub const FEATURE_GROUP_RANGES: [(usize, usize); SL_NUM_FEATURE_GROUPS] = [
(MARKET_START, OFI_START), // [0] market — 42 dims
(OFI_START, TLOB_START), // [1] ofi — 32 dims
(TLOB_START, MTF_START), // [2] tlob — 16 dims
(MTF_START, PORTFOLIO_START), // [3] mtf — 16 dims
(PORTFOLIO_START, PLAN_ISV_START), // [4] portfolio — 8 dims
(PLAN_ISV_START, PADDING_START), // [5] plan_isv — 7 dims
];
/// Group names — index-aligned with `FEATURE_GROUP_RANGES`. Used for ISV
/// diagnostic labels (`VSN_MASK_GROUP_<NAME>_EMA`) and audit-doc rows.
pub const FEATURE_GROUP_NAMES: [&str; SL_NUM_FEATURE_GROUPS] = [
"market",
"ofi",
"tlob",
"mtf",
"portfolio",
"plan_isv",
];
// ── Compile-time invariants ───────────────────────────────────────────────
const _: () = assert!(
FEATURE_GROUP_RANGES[0].0 == 0,
"feature groups must start at offset 0"
);
const _: () = assert!(
FEATURE_GROUP_RANGES[SL_NUM_FEATURE_GROUPS - 1].1 == PADDING_START,
"feature groups must end at PADDING_START — padding is not a group"
);
const _: () = {
// Every adjacent pair (g, g+1) must satisfy ranges[g].end == ranges[g+1].begin
// so the groups contiguously cover [0, PADDING_START) with no gaps.
let mut g = 0;
while g + 1 < SL_NUM_FEATURE_GROUPS {
assert!(
FEATURE_GROUP_RANGES[g].1 == FEATURE_GROUP_RANGES[g + 1].0,
"feature group ranges must be contiguous (no gaps)"
);
g += 1;
}
};
const _: () = {
// Each group's dim must fit within SL_MAX_FEATURE_GROUP_DIM so kernel
// scratch arrays sized to the maximum can hold any group.
let mut g = 0;
while g < SL_NUM_FEATURE_GROUPS {
let (b, e) = FEATURE_GROUP_RANGES[g];
assert!(e >= b, "feature group end must not precede begin");
assert!(
e - b <= SL_MAX_FEATURE_GROUP_DIM,
"feature group dim must not exceed SL_MAX_FEATURE_GROUP_DIM"
);
g += 1;
}
};