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

fix(cuda): fix GPU hotpath guard violations + dynamic shmem in backtest forward kernel

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-11 23:01:29 +01:00
parent a22e6420b6
commit a334ca288f
7 changed files with 18 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

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

8
crates/ml/src/ensemble/adapters/diffusion.rs
View File

@@ -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<f32> = 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;

6
crates/ml/src/ensemble/adapters/dqn.rs
View File

@@ -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<Vec<f32>> = 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;

6
crates/ml/src/ensemble/adapters/kan.rs
View File

@@ -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<Vec<f32>>
let raw_2d: Vec<Vec<f32>> = 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;

6
crates/ml/src/ensemble/adapters/liquid.rs
View File

@@ -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<Vec<f32>> = 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;

6
crates/ml/src/ensemble/adapters/ppo.rs
View File

@@ -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<Vec<f32>> = 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;

6
crates/ml/src/ensemble/adapters/tggn.rs
View File

@@ -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<Vec<f32>>
let raw_2d: Vec<Vec<f32>> = 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;