gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(rl): eliminate 2 host-blocking syncs per step from forward_encoder

Browse Source 
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>
This commit is contained in:
jgrusewski
2026-05-26 18:54:32 +02:00
parent a5b39ecc61
commit 107649c03b
1 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
crates/ml-alpha/src/trainer/perception.rs
View File

@@ -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)