Adds labels 47 (train_vsn_out), 48 (step_entry_w_in_W), 49 (step_entry_w_in_b).
10-run smoke at seed=16962, n=10, b=128: in ALL failing runs the FIRST
NaN is train_vsn_out at first_bad_idx=40 — the FIRST broadcast context
feature in the [B, K, ENCODER_INPUT_DIM=56] window tensor.
ROOT CAUSE: VSN kernel stride mismatch.
`window_tensor_d` is allocated `[B, K, ENCODER_INPUT_DIM=56]` and
`encoder_context_broadcast` writes to offset `idx * 56 + 40`. But the
VSN forward kernel `variable_selection_fwd` at variable_selection.cu:41
reads with `x_row = x + row * VSN_FEATURE_DIM=40` — stride 40 against
a stride-56 buffer.
For row 0 the stride happens to align. For row 1+ VSN reads mixed data
across the [B, K, 56] row boundaries — broadcast context features bleed
into VSN's "snap features" view, and snap features bleed across K
boundaries. The bug has existed since context broadcast was wired but
only produces NaN when the broadcast values get large enough to cause
softmax overflow — typically after a few steps of trade-context
accumulation (unrealized_R growing with position drift).
The "step-4" character is a side effect: clean weights + accumulated
trade-context magnitude → first row that overflows softmax → NaN
propagates through gradients → all encoder weights corrupted by
step 5 (confirmed: step_entry_w_in_W label 48 fires at step=5).
Confirmed by the chain of established scans:
- Labels 38-42 (top of step): all weights clean
- Label 47 (train_vsn_out post-VSN): NaN at idx=40 (first broadcast feature)
- Labels 43/44/45 (Mamba2 L1 projections pre-scan): NaN cascading
- Labels 36/37 (saved fwd state at L1 save): NaN cascading
- Labels 19, 31, 32, 34 (LN_a bwd inputs): NaN cascading
- Label 13 (vsn_W in action-selection fwd graph): fires at step=5 from AdamW
having applied NaN gradient (consequence, not cause)
Fix scope: change VSN_FEATURE_DIM stride to ENCODER_INPUT_DIM in
variable_selection_fwd and variable_selection_bwd, OR pass row_stride
as a kernel argument. Single-PR scope, no architecture change.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
4e629830b6