From 107649c03b969633da4d3bc3e92bd9a540d359f2 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Tue, 26 May 2026 18:54:32 +0200
Subject: [PATCH] perf(rl): eliminate 2 host-blocking syncs per step from
 forward_encoder
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Split forward_only into forward_only_dispatch (no sync, no download)
and forward_only (sync + download). forward_encoder now calls the
dispatch-only path — it only needs h_t_d on device, stream-ordered.

Was: 2 × (raw_stream_sync + probs download) per step from the two
forward_encoder calls in step_with_lobsim. Each sync drained the
entire GPU pipeline while CPU did host staging for the next call.
GPU was 98.5% idle per nsys.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
---
 crates/ml-alpha/src/trainer/perception.rs | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

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<Vec<f32>> {
+    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<Vec<f32>> {
+        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)