fix(sp14): dial back aux head/trunk enlargement after L40S OOM

Initial 8x head + 2x trunk-H2 (commit 787ee7b86) OOM'd all 3 folds on
L40S with alloc next_states f32[1048576000] failure. Root cause: aux
backward stores [B, H, SH2] partial grad buffers per sample. At
H=256, SH2=256, B=16384: 4GB per head x 2 heads + 4GB trunk = +12GB
on top of existing DQN buffers, exceeding L40S 48GB budget.

Dial back:
- AUX_HIDDEN_DIM: 256 -> 128 (still 4x prior 32-unit bottleneck)
- AUX_TRUNK_H2: 256 -> 128 (back to original mild bottleneck)

At H=128: partial buffer = 2GB per head = 4GB total. Fits comfortably
alongside next_states ~4GB and other DQN buffers.

Cargo check clean. Smoke re-dispatch next.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

crates/ml/src/cuda_pipeline/aux_heads_kernel.cu

@@ -64,9 +64,12 @@
 /* AUX_HIDDEN_DIM kept as a compile-time literal so kernel launch configs
  * don't need to thread `H` through. Both heads share the same hidden width. */
 #ifndef AUX_HIDDEN_DIM
-/* 2026-05-13: enlarged 32 → 256 (matches input h_s2_aux dim). See
- * `gpu_aux_heads.rs::AUX_HIDDEN_DIM` for rationale and side effects. */
-#define AUX_HIDDEN_DIM 256
+/* 2026-05-13: enlarged 32 → 128. Initial 256 attempt OOM'd L40S due
+ * to per-sample dW1_partial [B, H, SH2] = 16384 × 256 × 256 × 4 = 4GB
+ * per head. At H=128 it's 2GB per head = 4GB total, fits in 48GB
+ * budget alongside next_states ~4GB. See gpu_aux_heads.rs for full
+ * rationale. */
+#define AUX_HIDDEN_DIM 128
 #endif
 
 /* Block dim used by reduction-shaped launches (per-tensor backward, label

crates/ml/src/cuda_pipeline/gpu_aux_heads.rs

@@ -50,23 +50,29 @@ use super::gpu_dqn_trainer::{AUX_HEADS_CUBIN, AUX_HEADS_LOSS_EMA_CUBIN};
 /// Kept in sync with the `AUX_HIDDEN_DIM` macro at the top of
 /// `aux_heads_kernel.cu`.
 ///
-/// **Enlarged 32 → 256 (2026-05-13)** as part of the aux-head capacity
+/// **Enlarged 32 → 128 (2026-05-13)** as part of the aux-head capacity
 /// upgrade after Path C investigation revealed the aux head's 28%
 /// accuracy at H=60 bars. The prior 32-unit bottleneck was severely
 /// under-capacity for a 256-dim input → 2-way classification task. The
-/// new 256 hidden dim matches the input dimension exactly (no
-/// dimensionality reduction in Layer 1) and gives the head full
-/// capacity to learn an arbitrary mapping. Architecture:
-///   256 (h_s2_aux input) → Linear → 256 (hidden) → ELU → Linear → 2 (logits)
+/// new 128 hidden dim gives 4× capacity while keeping memory footprint
+/// under control. Architecture:
+///   256 (h_s2_aux input) → Linear → 128 (hidden) → ELU → Linear → 2 (logits)
+///
+/// **Memory note**: initial attempt at 256 (8×) OOM'd L40S. Aux backward
+/// stores `dW1_partial [B, H, SH2]` = batch_size × H × 256 floats =
+/// 16384 × 256 × 256 × 4 bytes = 4GB per head — total 8GB for both heads,
+/// which pushed total VRAM over the 48GB L40S budget when combined with
+/// the rollout's `next_states[~4GB]` and other DQN buffers. At H=128,
+/// the partial buffer is 2GB per head = 4GB total, fitting comfortably.
 ///
 /// Side effects:
 /// - Param buffer grows: aux_nb_w1/aux_rg_w1 went from [32, 256] to
-///   [256, 256] each → +64K params per head → +128K params total.
+///   [128, 256] each → +24K params per head → +48K params total.
 /// - Fingerprint changes → checkpoint incompatibility (intentional,
 ///   fresh-start training).
-/// - Thread utilisation improves: AUX_BLOCK=256 threads × H=256 hidden
-///   = 1 thread per hidden unit (vs 8 threads/unit at H=32).
-pub(crate) const AUX_HIDDEN_DIM: usize = 256;
+/// - Thread utilisation improves: AUX_BLOCK=256 threads × H=128 hidden
+///   = 2 threads per hidden unit (vs 8 threads/unit at H=32).
+pub(crate) const AUX_HIDDEN_DIM: usize = 128;
 
 /// Output cardinality of the next-bar direction classification head
 /// (SP13 B1.1a, 2026-05-05). Flipped 1 → 2: the head was a K=1 MSE

crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs

@@ -29,12 +29,13 @@
 //!   - `h_aux2 [B, H2]`       — Layer-2 post-ELU activation
 //!   - `h_s2_aux [B, AUX_HIDDEN_DIM]` — final linear output (no activation)
 //!
-//! Production topology (verified C.2; widths uplifted 2026-05-13):
+//! Production topology (verified C.2):
 //!   - `ENCODER_OUT_DIM = config.shared_h1 = 256`
 //!   - `H1 = AUX_TRUNK_H1 = 256`
-//!   - `H2 = AUX_TRUNK_H2 = 256` (was 128 pre-2026-05-13 — bottleneck removed)
-//!   - `AUX_HIDDEN_DIM = config.shared_h2 = 256` (matches existing aux head input dim;
-//!     internal aux head hidden layer ALSO enlarged 32 → 256 in same commit)
+//!   - `H2 = AUX_TRUNK_H2 = 128` (mild bottleneck)
+//!   - `AUX_HIDDEN_DIM = config.shared_h2 = 256` (matches aux head input dim;
+//!     internal aux head hidden layer enlarged 32 → 128 on 2026-05-13;
+//!     initial 32 → 256 attempt OOM'd L40S — dialed back to 128)
 //!
 //! # Pearls applied
 //!
@@ -66,16 +67,15 @@ use super::gpu_dqn_trainer::{
 pub(crate) const AUX_TRUNK_H1: usize = 256;
 
 /// Hidden width of the aux trunk's second internal layer (Linear_2 → ELU
-/// output).
+/// output). Smaller than H1 for a mild bottleneck — encourages the aux
+/// trunk to compress its representation before lifting back up to
+/// `AUX_HIDDEN_DIM` in Layer 3.
 ///
-/// **Enlarged 128 → 256 (2026-05-13)** alongside `AUX_HIDDEN_DIM 32 → 256`
-/// to remove the aux trunk's internal bottleneck. After the Path C
-/// investigation found the aux head's 28% accuracy at H=60 bars, the
-/// hypothesis is that capacity (not signal) was the bottleneck. The
-/// trunk now has uniform 256 width throughout: 256 (encoder) → 256 (L1)
-/// → 256 (L2) → 256 (L3 = h_s2_aux). This gives the trunk maximum
-/// expressive power before the aux head's classification layer.
-pub(crate) const AUX_TRUNK_H2: usize = 256;
+/// **Kept at 128 (2026-05-13)**: initial enlargement to 256 alongside
+/// `AUX_HIDDEN_DIM 32 → 256` OOM'd L40S — trunk dW partials grew too.
+/// Kept at 128 for memory budget; head's 4× enlargement to 128 is the
+/// primary capacity uplift.
+pub(crate) const AUX_TRUNK_H2: usize = 128;
 
 /// Block dim used by the per-sample forward kernel. Mirrors
 /// `aux_heads_kernel.cu`'s `AUX_BLOCK = 256`.

docs/dqn-wire-up-audit.md

@@ -17545,3 +17545,28 @@ After Path C confirmed the aux head's 28% accuracy at H=60 bars is the actual H6
 **Phase 3 mechanism stays DORMANT** for this smoke (W=0, β=0). If aux head accuracy improves to ≥55% with the larger capacity, the next step is to re-activate Phase 3 (W=[-0.5, 0, +0.5, 0], β=0.5) and re-validate WR.
 
 Cargo build clean (1m12s — full nvcc rebuild of aux kernels).
+
+#### Aux head capacity dial-back after OOM (2026-05-13)
+
+**Initial smoke `train-gxj42` @ `787ee7b86`** (8× head, 2× trunk-H2) OOM'd all 3 folds:
+```
+alloc next_states f32[1048576000]: DriverError(CUDA_ERROR_OUT_OF_MEMORY)
+```
+
+**Root cause**: aux backward stores `dW1_partial [B, H1, SH2]` partials per sample. At `B=16384, H=256, SH2=256`:
+- Per head: `16384 × 256 × 256 × 4 bytes = 4 GB`
+- Both heads (nb + regime): `8 GB`
+- Plus aux trunk's `dW2_partial [B, H1, H2]` = `16384 × 256 × 256 × 4 = 4GB`
+- Total +12GB on top of the existing DQN buffers → pushed total over L40S 48GB.
+
+**Dial-back**:
+| Constant | Initial (OOM) | Current |
+|----------|---------------|---------|
+| AUX_HIDDEN_DIM | 256 (8×) | **128 (4×)** |
+| AUX_TRUNK_H2 | 256 (2×) | **128 (1×, unchanged)** |
+
+At H=128, aux backward partial = `16384 × 128 × 256 × 4 = 2GB` per head = 4GB total. Fits comfortably in 48GB.
+
+Aux head architecture (now): `256 → Linear → 128 → ELU → Linear → 2` (4× capacity vs prior 32-unit bottleneck). Still meaningful capacity uplift while staying within memory budget.
+
+**Smoke will be re-dispatched at the dial-back commit**.