diff --git a/crates/ml-alpha/src/trainer/perception.rs b/crates/ml-alpha/src/trainer/perception.rs index 3d026203e..2297f4e0a 100644 --- a/crates/ml-alpha/src/trainer/perception.rs +++ b/crates/ml-alpha/src/trainer/perception.rs @@ -3831,9 +3831,8 @@ impl PerceptionTrainer { // Drive the encoder via the existing captured-graph forward // path. The returned probs vec is discarded — RL callers only // care about the encoder representation. - let _probs = self - .forward_only(snapshots) - .context("forward_encoder: forward_only")?; + self.forward_only_dispatch(snapshots) + .context("forward_encoder: forward_only_dispatch")?; // DtoD copy h_new_per_k_d at slot K-1 into the dedicated h_t_d // buffer. Layout in h_new_per_k_d is [K, B, HIDDEN_DIM] @@ -6293,7 +6292,7 @@ impl PerceptionTrainer { /// snapshots vary per call; the captured DtoD copies the latest /// staged data into the trunk's pre-bound device pointers. Mirrors /// the three-state machine already used by `step_batched`. - pub fn forward_only(&mut self, snapshots: &[Mbp10RawInput]) -> Result> { + fn forward_only_dispatch(&mut self, snapshots: &[Mbp10RawInput]) -> Result<()> { let b_sz = self.cfg.n_batch; let k_seq = self.cfg.seq_len; anyhow::ensure!( @@ -6390,7 +6389,14 @@ impl PerceptionTrainer { self.forward_graph = Some(graph); } - // ── 3. Sync + dtoh probs (OUTSIDE capture; download requires sync). + // ── 3. No sync here — callers that need host-side probs call + // forward_only() which syncs + downloads. forward_encoder() + // skips both (only needs h_t_d on device, stream-ordered). + Ok(()) + } + + pub fn forward_only(&mut self, snapshots: &[Mbp10RawInput]) -> Result> { + self.forward_only_dispatch(snapshots)?; unsafe { raw_stream_sync(self.raw_stream) } .map_err(|e| anyhow::anyhow!("forward_only end-sync: {:?}", e))?; download(&self.stream, &self.probs_per_k_d)