Corrective commit on top of a0e81fbdf addressing two load-bearing issues
flagged in the prior DONE_WITH_CONCERNS report.
Issue 1 (USER-FLAGGED, primary): GPU↔CPU roundtrip per event.
The previous `forward_step` did GPU compute → memcpy_dtod to mapped-pinned
host buffer → stream.synchronize() → CPU read → return [f32; N_HORIZONS]
→ harness stored in last_probs → broadcast_alpha memcpy_htod'd it back to
GPU. The 4-5 probs round-tripped the CPU twice per event for nothing.
The per-event stream.synchronize() defeated CUDA graph capture downstream
and throttled the event rate.
Per feedback_cpu_is_read_only and feedback_no_htod_htoh_only_mapped_pinned:
no compute data round-trips the CPU.
Fix:
- New `PerceptionTrainer::forward_step_into(snapshot, &mut alpha_probs_dst)`
signature. The GRN heads kernel writes per-horizon probs directly into
the caller's device buffer (`LobSimCuda::alpha_probs_d_mut`).
- Removed `probs_step_d`, `probs_step_host` fields. Removed the
DtoD-to-host-staging, the stream.synchronize, and the CPU read.
- New `LobSimCuda::alpha_probs_d_mut()` accessor exposes the on-device
decision-input buffer so the trainer writes directly into it.
- Harness loop now: `forward_step_into(&raw, sim.alpha_probs_d_mut())`
then (stride-gated) `step_decision_with_latency`. broadcast_alpha is
no longer called on the hot path — the probs were never on host.
- Conviction logging moved on-device: new `record_max_conviction_to_slot`
kernel writes one f32/decision into `LobSimCuda::convictions_d` (5M
capacity); `LobSimCuda::read_convictions(n)` DtoH's once at end of
run during `write_artifacts`. Replaces the host-side max-of-5 loop on
`self.last_probs` per decision. `last_probs` field deleted.
- Test-only helper `forward_step_into_returning(snap) -> [f32; N]`
preserves the prior test API shape with one DtoH; not exposed to
production callers. Existing forward_step_golden.rs tests retargeted
to this helper.
Acceptance check (per spec) passes — no memcpy_htod/dtoh/dtov/synchronize
inside forward_step_into or its callees. Only memcpy_dtod_async (DtoD).
Issue 2 (prior report concern #1): bit-identical seed from forward_only.
Previously the convergence test passed only at tolerance 0.15 because
forward_only seeds CfC's h_old from the attention pool over the K-window;
the step path starts from h=0 and the attention pool is dropped. Per memo
§4.5 Option (a) — extract terminal state from forward_only and seed
forward_step from it.
Fix:
- New `Mamba2Block::step_advance_from_seq_row(a_proj_ptr, b_proj_ptr,
scratch)` helper: launches scan_fwd_step against pre-computed
a_proj/b_proj from the seq path. Bit-identical x_state by construction
(same arithmetic, same per-step order). Skips the W_in/W_a/W_b GEMMs
which would otherwise differ from the seq path's batched GEMM at the
bit level.
- New `PerceptionTrainer::seed_step_state_from_forward_only(window)`:
1. Run forward_only(window) — populates mamba2 L1/L2 a_proj/b_proj
+ h_new_per_k_d via the regular seq path.
2. Reset step scratches' x_state to zero.
3. For k in 0..K: launch scan_fwd_step on step_scratch_l1 reading
row k of mamba2_fwd_scratch.a_proj/b_proj. Same for L2.
4. DtoD copy h_new_per_k_d[K-1] (the cfc h_new that would feed
position K if there were one) → cfc_h_state_step_d.
- New test forward_step_bit_identical_after_seed_from_forward_only at
1e-5 tolerance. Asserts forward_step_into on snapshot K (after seeding
from window [0..K-1]) matches forward_only's last-position prediction
on window [1..=K].
Residual structural caveat documented in the test: A and B see different
attn_context inputs (forward_only over [1..K+1] vs warmup [0..K]) and B
has one extra cfc iteration in its chain. The seed pins SSM x_state +
cfc h_state to forward_only's terminal values bit-identically; what
remains is cfc trajectory divergence after that pin. Asserting at 1e-5
exposes the gap at review rather than hiding it under a loose tolerance.
Per pearls: no host branches in captured graph (none added; kernel-only
work), no atomicAdd (block-tree-free single-thread kernel),
mapped-pinned-only for any CPU↔GPU contact (none on hot path; only the
constructor's weight upload + setup paths). cargo check workspace clean.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
92f8b10ed2