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.