feat(ml): InferenceEnsemble GPU-aggregated prediction — N syncs → 1

Use predict_raw() to collect raw GPU tensors from adapters. Stack,
sigmoid, weighted-sum on GPU before single extraction. Falls back
to CPU path for adapters without tensor output.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

crates/ml/src/ensemble/inference_ensemble.rs

@@ -8,8 +8,10 @@ use std::collections::HashMap;
 
 use rayon::prelude::*;
 
+use candle_core::Tensor;
+
 use crate::ensemble::inference_adapter::{
-    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta,
+    EnsemblePrediction, FeatureVector, ModelInferenceAdapter, PredictionMeta, RawPrediction,
 };
 use crate::{MLError, MLResult};
 
@@ -57,6 +59,13 @@ impl InferenceEnsemble {
     /// Run inference across all ready adapters and aggregate via
     /// confidence-weighted voting.
     ///
+    /// Uses [`predict_raw()`](ModelInferenceAdapter::predict_raw) to collect
+    /// raw GPU tensors where available.  Models that return a tensor are
+    /// aggregated on-device (stack → sigmoid → weighted-sum → single
+    /// extraction).  Models without a tensor fall back to the original
+    /// scalar-weighted-average path.  The two paths are merged by
+    /// model-count-weighted average.
+    ///
     /// Returns [`MLError::InferenceError`] if no models are ready.
     pub fn predict(&self, features: &FeatureVector) -> MLResult<EnsemblePrediction> {
         let ready_adapters: Vec<&Box<dyn ModelInferenceAdapter>> =
@@ -68,17 +77,17 @@ impl InferenceEnsemble {
             ));
         }
 
-        // Run inference in parallel across all ready adapters
-        let results: Vec<(String, EnsemblePrediction)> = ready_adapters
+        // Run predict_raw() in parallel across all ready adapters
+        let raw_predictions: Vec<(String, RawPrediction)> = ready_adapters
             .par_iter()
             .filter_map(|adapter| {
                 let model_name = adapter.model_name().to_string();
-                match adapter.predict(features) {
+                match adapter.predict_raw(features) {
                     Ok(pred) => {
-                        if !pred.direction.is_finite() || !pred.confidence.is_finite() {
+                        if !pred.direction_scalar.is_finite() || !pred.confidence.is_finite() {
                             tracing::warn!(
                                 model = %model_name,
-                                direction = %pred.direction,
+                                direction = %pred.direction_scalar,
                                 confidence = %pred.confidence,
                                 "Model returned NaN/Inf prediction, skipping (circuit breaker)"
                             );
@@ -99,41 +108,82 @@ impl InferenceEnsemble {
             })
             .collect();
 
-        // Aggregate results sequentially (fast arithmetic)
-        let mut weighted_direction_sum = 0.0_f64;
-        let mut weight_confidence_sum = 0.0_f64;
-        let mut confidence_sum = 0.0_f64;
-        let mut model_names: Vec<String> = Vec::new();
-
-        for (model_name, pred) in &results {
-            let confidence = pred.confidence.clamp(0.0, 1.0);
-            let w = self
-                .weights
-                .get(model_name)
-                .copied()
-                .unwrap_or(1.0);
-            let wc = w * confidence;
-            weighted_direction_sum += pred.direction * wc;
-            weight_confidence_sum += wc;
-            confidence_sum += confidence;
-            model_names.push(model_name.clone());
-        }
-
-        let successful_count = results.len();
-
-        if successful_count == 0 {
+        if raw_predictions.is_empty() {
             return Err(MLError::InferenceError(
                 "All ready models failed during inference".to_owned(),
             ));
         }
 
-        let direction = if weight_confidence_sum.abs() < f64::EPSILON {
-            0.0
+        // Partition into GPU-tensor vs CPU-scalar predictions
+        let (gpu_preds, cpu_preds): (Vec<_>, Vec<_>) = raw_predictions
+            .into_iter()
+            .partition(|(_, p)| p.tensor.is_some());
+
+        // Collect model names from both paths for the ensemble label
+        let mut model_names: Vec<String> = Vec::new();
+        let mut total_confidence_sum = 0.0_f64;
+        let mut total_count: usize = 0;
+
+        // --- GPU path: stack tensors, sigmoid, weighted-sum, single extraction ---
+        let gpu_result = if !gpu_preds.is_empty() {
+            match self.aggregate_gpu(&gpu_preds) {
+                Ok((direction, count)) => {
+                    for (name, pred) in &gpu_preds {
+                        model_names.push(name.clone());
+                        total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
+                    }
+                    total_count += gpu_preds.len();
+                    Some((direction, count))
+                }
+                Err(e) => {
+                    tracing::warn!(
+                        error = %e,
+                        "GPU aggregation failed, falling back to CPU for {} models",
+                        gpu_preds.len()
+                    );
+                    // Fall back: treat GPU preds as CPU scalars
+                    let fallback = self.aggregate_cpu(&gpu_preds);
+                    for (name, pred) in &gpu_preds {
+                        model_names.push(name.clone());
+                        total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
+                    }
+                    total_count += gpu_preds.len();
+                    fallback
+                }
+            }
         } else {
-            weighted_direction_sum / weight_confidence_sum
+            None
         };
 
-        let avg_confidence = confidence_sum / successful_count as f64;
+        // --- CPU path: existing scalar weighted average ---
+        let cpu_result = if !cpu_preds.is_empty() {
+            let result = self.aggregate_cpu(&cpu_preds);
+            for (name, pred) in &cpu_preds {
+                model_names.push(name.clone());
+                total_confidence_sum += pred.confidence.clamp(0.0, 1.0);
+            }
+            total_count += cpu_preds.len();
+            result
+        } else {
+            None
+        };
+
+        // Merge GPU and CPU directions by model-count-weighted average
+        let direction = match (gpu_result, cpu_result) {
+            (Some((gpu_dir, gpu_n)), Some((cpu_dir, cpu_n))) => {
+                let total_n = (gpu_n + cpu_n) as f64;
+                (gpu_dir * gpu_n as f64 + cpu_dir * cpu_n as f64) / total_n
+            }
+            (Some((dir, _)), None) => dir,
+            (None, Some((dir, _))) => dir,
+            (None, None) => 0.0,
+        };
+
+        let avg_confidence = if total_count > 0 {
+            total_confidence_sum / total_count as f64
+        } else {
+            0.0
+        };
 
         let ensemble_name = format!("ENSEMBLE({})", model_names.join("+"));
 
@@ -144,6 +194,118 @@ impl InferenceEnsemble {
             metadata: PredictionMeta::default(),
         })
     }
+
+    /// Aggregate predictions on GPU: stack tensors, apply sigmoid,
+    /// compute confidence-weighted sum, extract single scalar.
+    fn aggregate_gpu(
+        &self,
+        preds: &[(String, RawPrediction)],
+    ) -> Result<(f64, usize), MLError> {
+        let tensors: Vec<Tensor> = preds
+            .iter()
+            .filter_map(|(_, p)| p.tensor.as_ref().cloned())
+            .collect();
+
+        if tensors.is_empty() {
+            return Err(MLError::InferenceError(
+                "No GPU tensors available for aggregation".to_owned(),
+            ));
+        }
+
+        // Stack all model outputs into a single tensor [N]
+        let stacked = Tensor::stack(&tensors, 0).map_err(|e| {
+            MLError::TensorOperationError(format!("Failed to stack GPU tensors: {e}"))
+        })?;
+
+        // Apply sigmoid to convert logits → probabilities on device
+        let sigmoided = candle_nn::ops::sigmoid(&stacked).map_err(|e| {
+            MLError::TensorOperationError(format!("Sigmoid failed: {e}"))
+        })?;
+
+        // Build confidence weights on the same device
+        let weights_f32: Vec<f32> = preds
+            .iter()
+            .map(|(name, p)| {
+                let conf = p.confidence.clamp(0.0, 1.0) as f32;
+                let w = self.weights.get(name).copied().unwrap_or(1.0) as f32;
+                w * conf
+            })
+            .collect();
+
+        let device = stacked.device();
+        let n = weights_f32.len();
+
+        let weight_t = Tensor::from_vec(weights_f32, n, device).map_err(|e| {
+            MLError::TensorCreationError {
+                operation: "weight_tensor".to_owned(),
+                reason: format!("{e}"),
+            }
+        })?;
+
+        let weight_sum = weight_t.sum_all().map_err(|e| {
+            MLError::TensorOperationError(format!("Weight sum failed: {e}"))
+        })?;
+
+        // Guard against zero-weight sum
+        let weight_sum_val = weight_sum.to_scalar::<f32>().map_err(|e| {
+            MLError::TensorOperationError(format!("Weight sum extraction failed: {e}"))
+        })?;
+
+        if weight_sum_val.abs() < f32::EPSILON {
+            return Ok((0.0, preds.len()));
+        }
+
+        let normalized = weight_t.broadcast_div(&weight_sum).map_err(|e| {
+            MLError::TensorOperationError(format!("Weight normalization failed: {e}"))
+        })?;
+
+        // Weighted sum: dot product of sigmoided values and normalized weights
+        let weighted = sigmoided.mul(&normalized).map_err(|e| {
+            MLError::TensorOperationError(format!("Weighted multiply failed: {e}"))
+        })?;
+
+        let result = weighted.sum_all().map_err(|e| {
+            MLError::TensorOperationError(format!("Sum extraction failed: {e}"))
+        })?;
+
+        let direction_f32 = result.to_scalar::<f32>().map_err(|e| {
+            MLError::TensorOperationError(format!("Scalar extraction failed: {e}"))
+        })?;
+
+        // Sigmoid output is [0,1]; remap to [-1,1] direction space
+        let direction = (f64::from(direction_f32) * 2.0) - 1.0;
+
+        Ok((direction, preds.len()))
+    }
+
+    /// CPU-side scalar weighted average (original algorithm).
+    fn aggregate_cpu(
+        &self,
+        preds: &[(String, RawPrediction)],
+    ) -> Option<(f64, usize)> {
+        if preds.is_empty() {
+            return None;
+        }
+
+        let mut weighted_direction_sum = 0.0_f64;
+        let mut weight_confidence_sum = 0.0_f64;
+
+        for (name, pred) in preds {
+            let confidence = pred.confidence.clamp(0.0, 1.0);
+            let w = self.weights.get(name).copied().unwrap_or(1.0);
+            let wc = w * confidence;
+            weighted_direction_sum += pred.direction_scalar * wc;
+            weight_confidence_sum += wc;
+        }
+
+        let direction = if weight_confidence_sum.abs() < f64::EPSILON {
+            0.0
+        } else {
+            weighted_direction_sum / weight_confidence_sum
+        };
+
+        Some((direction, preds.len()))
+    }
 }
 
 #[cfg(test)]