foxhunt/docs/h_s2_consumers_audit.md

# h_s2 Consumer Audit (Plan 4 Task 2a)

Source of truth for which trunk-`h_s2` consumers tolerate signed input
(post-GRN swap) vs. which require non-negativity. Drives Task 2c
mitigation strategy.

## Methodology

For each consumer site:
1. Identify the kernel/file/line where `h_s2` (or its device pointer) is read.
2. Document the arithmetic applied to it (multiplication, division, log,
   exp, sigmoid, comparison, etc.).
3. Classify: **safe-with-signed** / **assumes-non-negative** / **needs-verification**.
4. Propose mitigation if assumes-non-negative.

Today the trunk produces `h_s2 = ReLU(W_s2 · h_s1 + b_s2)` (line 488 of
`crates/ml/src/cuda_pipeline/batched_forward.rs`), so the buffer is
guaranteed `>= 0`. Post-Task-2c the trunk will produce the LayerNorm
output of a GRN block — zero-mean, signed, and roughly `O(1)` after
the LN's affine transform. Each row below evaluates whether the
consumer's arithmetic still does what the design intended once `h_s2`
can be negative.

## Inventory

| # | Consumer site | File:line | Operation | Classification | Mitigation if needed |
|---|---|---|---|---|---|
| 1 | Value head L1 GEMM `h_v = ReLU(h_s2 · W_v1^T + b_v1)` | `batched_forward.rs:609`, `batched_forward.rs:941` (cuBLAS `sgemm_f32_fused_relu_bias`) | Linear projection followed by ReLU on output. Sign of `h_s2` only matters insofar as W_v1 is randomly Xavier-initialised; ReLU on the output absorbs negatives downstream. | safe-with-signed | None — the Xavier-initialised linear layer is sign-agnostic; output ReLU keeps downstream features non-negative. |
| 2 | Direction branch FC `h_b0 = ReLU(h_s2 · W_b0_fc^T + b_b0_fc)` (when VSN/GLU disabled, legacy path) | `batched_forward.rs:554`, `batched_forward.rs:676`, `batched_forward.rs:720`, `batched_forward.rs:821`, `batched_forward.rs:862`, `batched_forward.rs:1063` (decoder forward) | Linear projection + bias + ReLU. Identical shape to value head. | safe-with-signed | None. |
| 3 | Magnitude branch FC: input is `[h_s2; Q_dir]` concat (mag_concat) | `batched_forward.rs:544` and `experience_kernels.cu:3653-3656` (`mag_concat_qdir`) | First copies `h_s2` row verbatim, then appends `Q_dir/Δz`. The concat is fed through a linear+ReLU. The kernel comment explicitly says "Normalize by delta_z so Q_dir ~ 1-10 (matches h_s2 post-ReLU scale)" — this is a *scale* assumption, not a sign assumption. LayerNorm output has magnitude ~1, so `Q_dir/Δz ~ 1-10` no longer matches. | needs-verification | Scale-only concern. After the swap, normalise `Q_dir` against the new `h_s2` scale (e.g. divide `Q_dir/Δz` by `~Δz·sqrt(SH2)` so its RMS matches LN's RMS=1). No sign issue. Confirm by RMS-matching in Task 2c. |
| 4 | Order/Urgency branch FC concat path: `[h_s2; OFI3]` | `batched_forward.rs:710` and the `concat_ofi` kernel (no h_s2 nonlinearity, just bytewise scatter) | Same as #3: linear + ReLU on a wider input. | safe-with-signed | None — OFI features are already signed, so the layer already handles signed input. |
| 5 | VSN bottleneck `h_masked[b,i] = h_s2[b,i] · sigmoid(W_vsn1·W_vsn2·h_s2)` | `experience_kernels.cu:3775-3814` (`variable_select_bottleneck`), launched at `batched_forward.rs:1167` | Two linear projections through a `[SH2,R]·[R,SH2]` bottleneck → sigmoid → element-wise multiply with `h_s2`. The projection is sign-agnostic; sigmoid produces `(0,1)`. The element-wise product `h_s2 * gate` simply preserves the sign of `h_s2`. No magnitude or sign assumption is broken. | safe-with-signed | None. The kernel header comment ("At W_vsn2=0 init: sigmoid(0)=0.5 → h_masked = 0.5·h_s2 (stable start)") describes a property that survives the sign change unchanged. |
| 6 | Mamba2 history insert: `h_history[K-1] = [h_s2; ofi_embed]` | `mamba2_temporal_kernel.cu:33`, launched at `gpu_dqn_trainer.rs:5131` (`mamba2_update_history`) | Plain copy of `h_s2[i,j]` into the history buffer slot. No arithmetic. Downstream consumer (mamba2 scan) feeds `h_history` through a learned linear projection `W_A`/`W_B`. | safe-with-signed | None. |
| 7 | Mamba2 enrich: `h_enriched[i,j] = h_s2[i,j] + tw[j] · ctx` | `mamba2_temporal_kernel.cu:142`, launched at `gpu_dqn_trainer.rs:5108` (`mamba2_scan_projected_fwd`) | Pure additive residual: temporal context `tw·ctx` (signed) added to `h_s2`. Result is then copied back into `save_h_s2` for downstream branches. | safe-with-signed | None — additive residuals are sign-agnostic. |
| 8 | Multi-head feature attention (post-graph): `attn_in = h_s2`, output overwrites `save_h_s2` | `gpu_dqn_trainer.rs:6712-6743` (`apply_attention_forward`); CUDA kernels `attention_backward_kernel.cu:140-167` (`attn_sdp_fwd`) and `attention_kernel.cu` | Linear Q/K/V projections (sign-agnostic), feature-level attention with sigmoid gate `1/(1+exp(-Q·K/√d))`, output GEMM, residual `state + projected`, LayerNorm. The `attn_concat_input` kernel that builds the [D,B] attention input is also a plain copy of states and OFI — no arithmetic on `h_s2`. The init comment says "Zero W_O makes attention start as identity: output = LN(h_s2 + 0 @ attn) ≈ h_s2" — this property holds for signed `h_s2` because LN of a signed vector is well-defined. | safe-with-signed | None. Attention's existing residual+LayerNorm idempotently absorbs signed input. |
| 9 | IQN dual-head forward (per-quantile): `combined = h_s2 ⊙ embed`, `combined ⊙ relu_embed` | `iqn_dual_head_kernel.cu:191`, `:236-274`, `:301-310` (forward) and `:411` (target path); launched from `gpu_iqn_head.rs::execute_training_pipeline`. The element-wise product is `h_s2 ⊙ ReLU(W_embed·cos(τ))`. | Element-wise product of `h_s2` with the ReLU'd embedding. ReLU makes the embedding non-negative, but multiplying by signed `h_s2` is mathematically fine — `combined` is signed, fed to a per-branch linear `q = W_b · combined + b_b`. No log/sqrt/sigmoid is applied to `h_s2` itself. Backward (`iqn_dual_head_kernel.cu:684`) computes `d_pre = dq · w_sum · h_s2 · (embed_pre > 0)` — multiplies by `h_s2`, fine when signed. | safe-with-signed | None. Note that the kernel today ingests `h_s2` as `float*` even though the header comment says "(bf16 from DQN trunk)" — that's a stale comment; the actual storage is f32 already. |
| 10 | IQN tile + Hadamard-sigmoid path: `out = h_s2_tiled · sigmoid(embed)` | `iqn_dual_head_kernel.cu:1262` (`iqn_tile_hadamard_sigmoid`), `:1284` (backward `d_h_s2_tiled = d_combined · sigmoid(embed)`); tile/reduce kernels at `:1361-1396` | Element-wise product of `h_s2_tiled` with `sigmoid(embed_pre) ∈ (0,1)`. Sign of `h_s2` propagates linearly through the sigmoid-gated multiply. Backward is symmetric. | safe-with-signed | None. |
| 11 | IQN target trunk — `iqn_trunk_forward_kernel` (audit-orig name `iqn_compute_target_h_s2`) | `iqn_dual_head_kernel.cu` kernel function. **CORRECTION (2026-04-25 post-2c.3+4 recon):** the kernel is **already orphaned** — loaded into `IqnHead::trunk_forward_kernel` but never invoked. The actual IQN target trunk runs through cuBLAS `iqn_lt_matmul` calls in `gpu_iqn_head.rs::execute_training_pipeline` lines 803-847 (cached fast-path) and 820-846 (cuBLAS fallback). | The cuBLAS fallback path reproduces the legacy `Linear → ReLU → Linear → ReLU` via `launch_trunk_bias_relu`. Once GRN lands, online and target diverge unless the cuBLAS path is replaced with a `target_encoder_forward_only` call (Task 4-style extraction on the target-net path). | **assumes-non-negative** *(via cuBLAS path's `launch_trunk_bias_relu`)* | (a) Delete the dead `iqn_trunk_forward_kernel` + its `load(...)` line + struct field. (b) Extract `target_encoder_forward_only` mirroring Task 4's online encoder split. (c) Replace the cuBLAS fallback path in `execute_training_pipeline:820-846` with `target_encoder_forward_only`. The cached fast-path stays valid as long as `cached_target_h_s2_ptr` is updated by whoever populates it. |
| 12 | Risk-budget head: SH2 cooperative load, hidden = ReLU(W_risk_fc · [h_s2; ISV; pred_err] + b) | `experience_kernels.cu:5867-5917` (`risk_budget_forward`), launched at `gpu_dqn_trainer.rs:4592-4626` | Concats `h_s2[B,SH2]` with `ISV[12]` and `predicted_error[1]`, runs linear + ReLU + sigmoid. Backward (`experience_kernels.cu:5954-6037`) reads `h_s2[i,k]` as `input_val` for `dW`. | safe-with-signed | None. |
| 13 | Recursive confidence head: `pred = sigmoid(W·h_s2 + b)` | `experience_kernels.cu:6447-6462` (forward) + `:6476-6543` (backward); Rust at `gpu_dqn_trainer.rs:4001-4027`/`:4443-4496` | Linear projection through sigmoid. Backward writes `d_h_s2 += d_sigmoid · w_conf` (signed contribution) and `dW += d_sigmoid · h_s2`. Both are sign-agnostic. | safe-with-signed | None. |
| 14 | Selectivity head: `sel = sigmoid(W_sel · h_s2 + b_sel)` | `experience_kernels.cu:3998-4014` (forward), `:4056-4068` (dW reduce) | Plain linear projection through sigmoid. Backward `dW += d_z · h_s2[b,k]` is sign-agnostic. | safe-with-signed | None. |
| 15 | ISV feature gate (in-place): `h_s2[i,j] *= sigmoid(W·isv_emb + b)` | `experience_kernels.cu:6580-6600` (`isv_feature_gate`); launched at `gpu_dqn_trainer.rs:3943-3969` | Multiplies each `h_s2` cell by a per-feature sigmoid gate `∈ (0,1)`. Sign of `h_s2` is preserved by the multiply. No other arithmetic on `h_s2`. | safe-with-signed | None. |
| 16 | ISV temporal route | `mamba2_temporal_kernel.cu:50-64` (`isv_temporal_route`); launched at `gpu_dqn_trainer.rs:3974-3997` | Reads `isv_embedding`, NOT `h_s2`. Misclassified in the spec's "12+" list — this kernel only consumes `h_s2` indirectly through the mamba2 enrich kernel (#7). | safe-with-signed (no direct read) | None. |
| 17 | Trade-plan / plan head L1 GEMM: `hidden = ReLU(W_fc · h_s2 + b_fc)` | `gpu_dqn_trainer.rs:4029-4097` (`launch_trade_plan_forward`); subsequent `trade_plan_activate` kernel at `experience_kernels.cu:6608-6627` operates on `pre_out`, not `h_s2` | Two cuBLAS GEMMs with bias+ReLU between them (and bias+sigmoid epilogue at the end). The first GEMM consumes `h_s2`. Sign-agnostic. The plan activations live downstream and don't read `h_s2`. | safe-with-signed | None. |
| 18 | Regime-aware dropout (in-place): `h_s2[idx] = 0` or `*= 1/(1-p)` | `experience_kernels.cu:5616-5644` (`regime_dropout`); launched at `gpu_dqn_trainer.rs:4707-4725` | Zeros a fraction of `h_s2` cells, scales survivors by `1/(1-p)`. Both ops are sign-preserving. | safe-with-signed | None. |
| 19 | Predictive coding loss: `L = λ · mean_j(h_s2[i] − h_s2[i+1])²` | `experience_kernels.cu:5713-5736` (forward), `:5766-5800` (backward); launched at `gpu_dqn_trainer.rs:4781-4827` | Squared difference between consecutive samples; output is non-negative regardless of input sign. Backward computes `dL/dh = 2λ/SH2 · (h_self − h_neighbour)`, also sign-symmetric. | safe-with-signed | None. |
| 20 | Mag-concat backward `strided_accumulate` extracts `d_h_s2 += d_mag_concat[:SH2]` | `experience_kernels.cu:3708-3726`; called at `gpu_dqn_trainer.rs:4831+` (`accumulate_d_h_s2_from_concat`) | Pure copy/accumulate; no nonlinearity. | safe-with-signed | None. |
| 21 | Branch-decoder VSN+GLU `launch_vsn_glu_branch` (per branch path) | `batched_forward.rs:1143-1306` calls `vsn_kernel`, then `kan_gate_combine_kernel` (KAN spline gate). The first stage is the VSN above (#5). | The KAN spline gate `kan_gate_combine` in this path operates on the *bottleneck* output of a per-branch FC, not on `h_s2` directly. Its input is `b1_proj` (after the per-branch linear+ReLU). KAN doesn't touch `h_s2`. | safe-with-signed (no direct read) | None — KAN sees only post-FC bottleneck features. |
| 22 | **Trunk-backward ReLU mask: `d_h_s2 *= (save_h_s2 > 0)`** | `batched_backward.rs:1778` (`relu_mask`), `gpu_dqn_trainer.rs:5655` (online trunk grad post-h_s2), `gpu_dqn_trainer.rs:5697` (h_s1 mask, NOT h_s2 — verify before editing), `gpu_dqn_trainer.rs:5872` (target-sync per-target-update). Kernels: `relu_mask_kernel.cu` and `backward_kernels.cu:18`. **Line numbers updated 2026-04-25 post-recon — original audit had stale 5581/5798 anchors.** | This is the **load-bearing ReLU assumption** in the entire pipeline. The mask zeros gradients where `save_h_s2 <= 0`, which is correct mathematics for `h = ReLU(z)`. With LayerNorm output (zero-mean, signed), this would zero ~50% of gradients corresponding to **valid** negative activations — a silently wrong gradient with no NaN/Inf. | **assumes-non-negative** | The `relu_mask` step **must be removed** (or replaced with the GRN's actual jacobian) on the post-Task-2c trunk-out. The 2c.1 GRN kernel module already provides the full LN+GLU+ELU backward; the 2c.2 wrapper sequences it. Each of the 3 sites has different save-buffer lifetime semantics (online per-step, IQN-aux per-step in different trainer, target-sync per-target-update) — verify save-buffer plumbing per site. |
| 23 | Target-network trunk forward: `tg_h_s2_buf` is produced/consumed by the same path as #1, #2, #4 | `batched_forward.rs:980-1090` (`forward_target_raw`); IQN reuses `tg_h_s2_buf` via `gpu_iqn_head.rs::set_cached_target_h_s2_ptr` → `cached_target_h_s2_ptr` | Same arithmetic as the online trunk: linear + ReLU(legacy) → branches via linear+ReLU. After Task 2c, target trunk MUST run the same GRN block as the online trunk. | **assumes-non-negative** *(via shared kernel reuse)* | Target trunk weight set must be migrated alongside online. The relu_mask in `apply_iqn_trunk_gradient` (`gpu_dqn_trainer.rs:5581`) currently masks against `save_h_s2`, the **online** trunk activation; once that buffer is LN-output the mask is wrong (see #22). |
| 24 | Experience-collector `exp_h_s2_f32` storage | `gpu_experience_collector.rs:559,690,1231-1299,3306` | Plain f32 storage of past `h_s2` for replay/episode-level data. Read back by replay/IQN as ordinary tensor data. No nonlinearity at the storage layer. | safe-with-signed | None — storage layout is sign-agnostic. |
| 25 | Q-stats / monitoring path (`tg_h_s2_ptr` exposed for monitoring) | `gpu_dqn_trainer.rs:4936-4938` accessor | The pointer is exposed but the call sites are restricted to monitoring readback. Monitoring kernels we audited do not perform sign-dependent ops on `h_s2`. | safe-with-signed | None. |
| 26 | Cross-branch Q-attention `cross_branch_q_attention` | `experience_kernels.cu:3894-3984` | Reads `Q_raw[B,12]` (per-branch Q-values), NOT `h_s2`. Listed here only to dispel the "12+" spec list; this attention is on Q-values, not features. | safe-with-signed (no `h_s2` read) | None. |
| 27 | Ensemble value head `forward_value_head_for_ensemble` | `gpu_dqn_trainer.rs:11787-11801` calls `cublas_forward.forward_value_head(h_s2, ...)` | Linear+ReLU same shape as #1. | safe-with-signed | None. |
| 28 | Direction-branch backward `bw_d_h_s2` accumulation by branch heads | `batched_backward.rs:1466-1470` — comment "Branch FC layer: `h_bd` gradient → `h_s2` (accumulated)" | Linear backward; all `dX = dY · W` ops are sign-agnostic. | safe-with-signed | None — but the accumulated buffer is masked by #22 downstream. |

## Summary

- **safe-with-signed: 24** consumers (rows 1–10, 12–21, 24–28, plus the indirect/no-direct-read rows 16, 21, 26)
- **assumes-non-negative: 3** consumers — rows **11** (`iqn_compute_target_h_s2` reproduces the legacy LeakyReLU trunk), **22** (the `relu_mask` step on `bw_d_h_s2` in trunk and IQN-aux backward), **23** (target trunk via the same kernel as online).
- **needs-verification: 1** — row **3** (mag_concat scale assumption: `Q_dir/Δz ~ 1-10 matches h_s2 post-ReLU scale`). This is a *scale* mismatch, not a sign mismatch, and is straightforward to fix once the new `h_s2` magnitude is measured.

### What was surprising

- **Most consumers are linear projections** (cuBLAS GEMMs into branch/value/risk/conf/selectivity/plan heads) and are inherently sign-agnostic. The audit confirms the design intuition: the trunk's ReLU was a feature of the trunk itself, not an externalised contract that downstream consumers depend on.
- **The two real failure modes both live in the backward pass**, not the forward pass: the `relu_mask` step (#22, #23) and the legacy LeakyReLU recompute kernel for the target trunk (#11). All forward-direction nonlinearities (sigmoid gates, element-wise products, additive residuals) are sign-tolerant.
- **VSN is safe** despite the kernel's "(0,1) gate × h_s2" framing — the comment about "softmax would give 1/SH2 collapse" describes a *training-init* property, not a sign assumption. Element-wise sigmoid × signed is well-defined.
- **The "12+ consumers" list in the spec partially miscounted**: ISV-temporal-route (#16), KAN-spline (#21), and cross-branch-Q-attention (#26) do **not** read `h_s2` directly. Plan-head, risk-budget, recursive-confidence, selectivity, and ISV-feature-gate **do**. The actual count of direct readers is closer to 22 once you split forward vs backward and online vs target.

## Recommended Task 2c shim strategy

The audit answers the strategic question definitively: **do not add a global `relu(h_s2)` shim post-LayerNorm and do not add per-consumer relu shims**. A global shim would re-introduce the dead-zone the GRN swap is meant to eliminate (negative LN values are *real signal*, not noise to be clamped); per-consumer shims would require auditing 24 sign-tolerant call sites for nothing. Instead, the 2c migration is a **backward-side surgical replacement** at exactly two points:

1. **Replace the `relu_mask(d_h_s2, save_h_s2)` call** in `batched_backward.rs:1778`, plus the two equivalent calls in `gpu_dqn_trainer.rs:5581-5599` and `:5798-5801`, with the actual GRN-block backward (LayerNorm Jacobian + GLU sigmoid backward + ELU' on `Linear_a(h_s1)`). The LayerNorm backward kernels already exist for the attention path (`attn_layer_norm_bwd_*`); the new GRN-block backward should reuse that and add the GLU/ELU stages. Everything upstream of `bw_d_h_s2` (branch heads, value head, risk, conf, selectivity, predictive coding) writes its share without any sign assumption, so the only change is *what comes after the accumulator is full*.
2. **Migrate `iqn_compute_target_h_s2`** (`iqn_dual_head_kernel.cu:1031-1092`) to the new GRN trunk in the same commit, or — preferably — delete the kernel and wire the IQN target path through `BatchedForward::forward_target_raw` so there is exactly one trunk implementation and no chance for online/target divergence. This is a "no-partial-refactor" obligation per `feedback_no_partial_refactor.md` — the trunk topology is a shared contract.

The single scale-tweak in row #3 (RMS-match `Q_dir/Δz` against the new `h_s2` RMS, since LN output has RMS ≈ 1 vs ReLU output ≈ 1–10) is the only forward-side adjustment needed and lives entirely inside the `mag_concat_qdir` kernel's normalisation constant — a one-line adaptive change driven by a measured ISV slot, not a tuned constant.