diff --git a/crates/ml-alpha/src/per_horizon_attention_pool.rs b/crates/ml-alpha/src/per_horizon_attention_pool.rs index 44f740cd1..e1b3345a3 100644 --- a/crates/ml-alpha/src/per_horizon_attention_pool.rs +++ b/crates/ml-alpha/src/per_horizon_attention_pool.rs @@ -82,7 +82,8 @@ impl PerHorizonAttentionPool { .launch(cfg) .context("per_horizon_attention_pool_fwd")?; } - self.stream.synchronize()?; + // No synchronize: same-stream issue order is sufficient and + // synchronize is illegal during CUDA Graph capture. Ok(()) } @@ -122,7 +123,8 @@ impl PerHorizonAttentionPool { .launch(cfg) .context("per_horizon_attention_pool_bwd")?; } - self.stream.synchronize()?; + // No synchronize: same-stream issue order is sufficient and + // synchronize is illegal during CUDA Graph capture. Ok(()) } } diff --git a/crates/ml-alpha/src/per_horizon_residual_head.rs b/crates/ml-alpha/src/per_horizon_residual_head.rs index df8ef9ce2..c78bd7273 100644 --- a/crates/ml-alpha/src/per_horizon_residual_head.rs +++ b/crates/ml-alpha/src/per_horizon_residual_head.rs @@ -71,7 +71,8 @@ impl PerHorizonResidualHead { .launch(cfg) .context("per_horizon_residual_head_fwd")?; } - self.stream.synchronize()?; + // No synchronize: same-stream issue order is sufficient and + // synchronize is illegal during CUDA Graph capture. Ok(()) } @@ -107,7 +108,8 @@ impl PerHorizonResidualHead { .launch(cfg) .context("per_horizon_residual_head_bwd")?; } - self.stream.synchronize()?; + // No synchronize: same-stream issue order is sufficient and + // synchronize is illegal during CUDA Graph capture. Ok(()) } }