From a334ca288ff44fbf02622fcfe5b49eb51099ce0b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 11 Mar 2026 23:01:29 +0100 Subject: [PATCH] fix(cuda): fix GPU hotpath guard violations + dynamic shmem in backtest forward kernel MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move // gpu-ok: annotations to same line as violation patterns so the guard script's grep -v filter actually suppresses them. Fixes 6 false positives in ensemble adapters (dqn, ppo, liquid, kan, tggn, diffusion). Replace hardcoded 48KB shmem limit in compile_forward_kernel() with GPU-aware query (max_shared_memory_kb) — matches gpu_experience_collector pattern. H100 now gets 128-row tiles (was 64), eliminating tile loops for ≤128-dim layers in the backtest forward kernel. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs | 7 +++++-- crates/ml/src/ensemble/adapters/diffusion.rs | 8 +++----- crates/ml/src/ensemble/adapters/dqn.rs | 6 ++---- crates/ml/src/ensemble/adapters/kan.rs | 6 ++---- crates/ml/src/ensemble/adapters/liquid.rs | 6 ++---- crates/ml/src/ensemble/adapters/ppo.rs | 6 ++---- crates/ml/src/ensemble/adapters/tggn.rs | 6 ++---- 7 files changed, 18 insertions(+), 27 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index 2ef92e9c6..cf9ddc2cf 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -1212,9 +1212,12 @@ impl GpuBacktestEvaluator { let state_dim = self.feature_dim + self.portfolio_dim; let shmem_max_in_dim = state_dim.max(shared_h1).max(shared_h2); let shmem_tile_rows = { - let max_tile = 49152 / (4 * (shmem_max_in_dim + 1)); + let gpu_caps = crate::gpu::capabilities::cached_capabilities(); + let shmem_kb = crate::gpu::capabilities::max_shared_memory_kb(&gpu_caps.device_name); + let shmem_limit_bytes = (shmem_kb * 1024 * 80 / 100).max(49152); + let max_tile = shmem_limit_bytes / (4 * (shmem_max_in_dim + 1)); let pow2 = (max_tile as u32).next_power_of_two() >> 1; - pow2.max(16) as usize + pow2.clamp(16, 256) as usize }; let common_src = include_str!("common_device_functions.cuh"); diff --git a/crates/ml/src/ensemble/adapters/diffusion.rs b/crates/ml/src/ensemble/adapters/diffusion.rs index 313c3e0ad..e570df2ec 100644 --- a/crates/ml/src/ensemble/adapters/diffusion.rs +++ b/crates/ml/src/ensemble/adapters/diffusion.rs @@ -206,14 +206,13 @@ impl ModelInferenceAdapter for DiffusionInferenceAdapter { flat.extend_from_slice(&self.pad_features(&fv.values)); } - // gpu-ok: single batched inference input tensorization - let input = Tensor::from_vec(flat, (n, self.data_dim), &self.device) + let input = Tensor::from_vec(flat, (n, self.data_dim), &self.device) // gpu-ok: single batched inference input tensorization .map_err(|e| MLError::ModelError(format!("Diffusion batch input: {e}")))?; let input = crate::dqn::mixed_precision::ensure_training_dtype(&input) .map_err(|e| MLError::ModelError(format!("Diffusion batch dtype: {e}")))?; // Timestep t=1 for all samples in the batch - let t = Tensor::from_vec(vec![1_u32; n], (n,), &self.device) + let t = Tensor::from_vec(vec![1_u32; n], (n,), &self.device) // gpu-ok: batch timestep constant .map_err(|e| MLError::ModelError(format!("Diffusion batch timestep: {e}")))?; let model = self @@ -227,11 +226,10 @@ impl ModelInferenceAdapter for DiffusionInferenceAdapter { let means = output .mean(1) .map_err(|e| MLError::ModelError(format!("Diffusion batch mean: {e}")))?; - // gpu-ok: single batched scalar extraction let mean_vec: Vec = means .to_dtype(candle_core::DType::F32) .map_err(|e| MLError::ModelError(format!("Diffusion means dtype: {e}")))? - .to_vec1() + .to_vec1() // gpu-ok: single batched scalar extraction .map_err(|e| MLError::ModelError(format!("Diffusion means extract: {e}")))?; let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64; diff --git a/crates/ml/src/ensemble/adapters/dqn.rs b/crates/ml/src/ensemble/adapters/dqn.rs index 07c62905e..f8b65e433 100644 --- a/crates/ml/src/ensemble/adapters/dqn.rs +++ b/crates/ml/src/ensemble/adapters/dqn.rs @@ -166,8 +166,7 @@ impl ModelInferenceAdapter for DqnInferenceAdapter { } } - // gpu-ok: single batched inference input tensorization - let input = Tensor::from_vec(flat, (n, feature_dim), &self.device) + let input = Tensor::from_vec(flat, (n, feature_dim), &self.device) // gpu-ok: single batched inference input tensorization .map_err(|e| MLError::ModelError(format!("Batch input tensor: {e}")))?; let input = ensure_training_dtype(&input) .map_err(|e| MLError::ModelError(format!("Batch dtype cast: {e}")))?; @@ -180,11 +179,10 @@ impl ModelInferenceAdapter for DqnInferenceAdapter { let q_output = model.forward(&input)?; drop(model); - // gpu-ok: single batched Q-values extraction let q_2d: Vec> = q_output .to_dtype(candle_core::DType::F32) .map_err(|e| MLError::ModelError(format!("Q-values dtype cast: {e}")))? - .to_vec2() + .to_vec2() // gpu-ok: single batched Q-values extraction .map_err(|e| MLError::ModelError(format!("Q-values extraction: {e}")))?; let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64; diff --git a/crates/ml/src/ensemble/adapters/kan.rs b/crates/ml/src/ensemble/adapters/kan.rs index bc55d93e3..b05ed28e3 100644 --- a/crates/ml/src/ensemble/adapters/kan.rs +++ b/crates/ml/src/ensemble/adapters/kan.rs @@ -182,8 +182,7 @@ impl ModelInferenceAdapter for KanInferenceAdapter { flat.extend_from_slice(&self.pad_features(&fv.values)); } - // gpu-ok: single batched inference input tensorization - let input = Tensor::from_vec(flat, (n, self.input_dim), &self.device) + let input = Tensor::from_vec(flat, (n, self.input_dim), &self.device) // gpu-ok: single batched inference input tensorization .map_err(|e| MLError::ModelError(format!("KAN batch input: {e}")))?; let input = crate::dqn::mixed_precision::ensure_training_dtype(&input) .map_err(|e| MLError::ModelError(format!("KAN batch dtype: {e}")))?; @@ -195,9 +194,8 @@ impl ModelInferenceAdapter for KanInferenceAdapter { let output = model.forward(&input)?; drop(model); - // gpu-ok: single batched extraction [N, 1] → Vec> let raw_2d: Vec> = output - .to_vec2() + .to_vec2() // gpu-ok: single batched extraction [N, 1] .map_err(|e| MLError::ModelError(format!("KAN batch extract: {e}")))?; let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64; diff --git a/crates/ml/src/ensemble/adapters/liquid.rs b/crates/ml/src/ensemble/adapters/liquid.rs index 43fcaf663..37d90522e 100644 --- a/crates/ml/src/ensemble/adapters/liquid.rs +++ b/crates/ml/src/ensemble/adapters/liquid.rs @@ -198,8 +198,7 @@ impl ModelInferenceAdapter for LiquidInferenceAdapter { } // CfC expects 3D: [N, seq_len=1, features] - // gpu-ok: single batched inference input tensorization - let input = Tensor::from_vec(flat, (n, 1, self.input_size), &self.device) + let input = Tensor::from_vec(flat, (n, 1, self.input_size), &self.device) // gpu-ok: single batched inference input tensorization .map_err(|e| MLError::ModelError(format!("Liquid batch input: {e}")))?; let model = self @@ -210,9 +209,8 @@ impl ModelInferenceAdapter for LiquidInferenceAdapter { drop(model); // Output: [N, output_size] → extract all at once - // gpu-ok: single batched extraction let raw_2d: Vec> = output - .to_vec2() + .to_vec2() // gpu-ok: single batched extraction .map_err(|e| MLError::ModelError(format!("Liquid batch extract: {e}")))?; let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64; diff --git a/crates/ml/src/ensemble/adapters/ppo.rs b/crates/ml/src/ensemble/adapters/ppo.rs index c05ea477b..0dbb53e40 100644 --- a/crates/ml/src/ensemble/adapters/ppo.rs +++ b/crates/ml/src/ensemble/adapters/ppo.rs @@ -159,8 +159,7 @@ impl ModelInferenceAdapter for PpoInferenceAdapter { flat.extend_from_slice(&self.pad_features(&fv.values)); } - // gpu-ok: single batched inference input tensorization - let input = Tensor::from_vec(flat, (n, self.state_dim), &self.device) + let input = Tensor::from_vec(flat, (n, self.state_dim), &self.device) // gpu-ok: single batched inference input tensorization .map_err(|e| MLError::ModelError(format!("PPO batch input tensor: {e}")))?; // One forward pass: [N, state_dim] → [N, num_actions] @@ -171,11 +170,10 @@ impl ModelInferenceAdapter for PpoInferenceAdapter { let probs_tensor = model.actor.action_probabilities(&input)?; drop(model); - // gpu-ok: single batched probability extraction let probs_2d: Vec> = probs_tensor .to_dtype(candle_core::DType::F32) .map_err(|e| MLError::ModelError(format!("PPO probs dtype cast: {e}")))? - .to_vec2() + .to_vec2() // gpu-ok: single batched probability extraction .map_err(|e| MLError::ModelError(format!("PPO probs extraction: {e}")))?; let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64; diff --git a/crates/ml/src/ensemble/adapters/tggn.rs b/crates/ml/src/ensemble/adapters/tggn.rs index 2879d2a98..cd81f7f07 100644 --- a/crates/ml/src/ensemble/adapters/tggn.rs +++ b/crates/ml/src/ensemble/adapters/tggn.rs @@ -216,8 +216,7 @@ impl ModelInferenceAdapter for TggnInferenceAdapter { flat.extend_from_slice(&self.pad_features(&fv.values)); } - // gpu-ok: single batched inference input tensorization - let input = Tensor::from_vec(flat, (n, self.input_dim), &self.device) + let input = Tensor::from_vec(flat, (n, self.input_dim), &self.device) // gpu-ok: single batched inference input tensorization .map_err(|e| MLError::ModelError(format!("TGGN batch input: {e}")))?; let model = self @@ -227,9 +226,8 @@ impl ModelInferenceAdapter for TggnInferenceAdapter { let output = model.forward(&input)?; drop(model); - // gpu-ok: single batched extraction [N, 1] → Vec> let raw_2d: Vec> = output - .to_vec2() + .to_vec2() // gpu-ok: single batched extraction [N, 1] .map_err(|e| MLError::ModelError(format!("TGGN batch extract: {e}")))?; let latency_us = start.elapsed().as_micros() as u64 / n.max(1) as u64;