feat(demos,inference): wire simulation demos to real compute; fix embedding lookup and weight-name aliases

Demos:
- rtx-distllm-demo: real rtx-tensor weights per shard, real
  scaled-dot-product attention forward, metrics measured (Instant)
  instead of hardcoded constants; network topology remains a documented
  simulation fed by real tensor byte sizes.
- rtx-model-zoo: MockInferenceEngine deleted; RealInferenceEngine loads
  a tiny real transformer into rtx_inference::InferenceEngine and runs
  genuine engine.infer per request; domain outputs are explicitly-
  labeled toy proxies derived from real output tokens.
- rtx-inference-profiler: mock models deleted; profiles real
  matmul/softmax pipelines on rtx-tensor with measured latency/memory.

Inference-path bugs the demos surfaced (fixed here):
- ForwardPass::apply_embedding misused Tensor::gather for the embedding
  lookup — gather returns the indices' shape, silently dropping the
  hidden dim and breaking every downstream broadcast. Now uses the
  existing Tensor::embedding_lookup ([vocab,hidden] x [batch,seq] ->
  [batch,seq,hidden]).
- Attention weight lookup accepts both self_attn. (HF-LLaMA) and
  attention. prefixes; final layer norm accepts norm.weight /
  model.norm.weight / ln_f.weight aliases.
- Integration fixture gains the final norm weight; the previously
  always-failing engine tests now pass (8/8 model_loading_test).

End-to-end inference through the real engine now works for the first
time — verified via model_zoo_demo producing real forward-pass outputs
across all categories.

Co-Authored-By: Claude Fable 5 <[email protected]>

demos/rtx-model-zoo/src/real_inference.rs

@@ -0,0 +1,521 @@
+//! Real inference engine backed by `rtx_inference::InferenceEngine`.
+//!
+//! # What is real here
+//!
+//! - A tiny transformer (`ModelConfig` with a handful of layers/heads) is built with
+//!   randomly-initialized weights and loaded into a genuine `rtx_inference::InferenceEngine`.
+//! - Every call to [`RealInferenceEngine::run_inference`] drives the engine's real
+//!   `infer(...)` path: the input string is tokenized (byte-level, modulo the tiny
+//!   vocabulary) into `input_tokens`, and the engine performs a real forward pass
+//!   (real tensor matmuls through the transformer layers) to autoregressively produce
+//!   `output_tokens`.
+//! - `inference_time_ms` is a genuine `Instant::elapsed()` measurement wrapped tightly
+//!   around the `engine.infer(...)` call — not a random formula.
+//! - `device` is read from the real `rtx_tensor::Device` the engine was constructed with.
+//!
+//! # What is NOT real (and is labeled as such in the output)
+//!
+//! The tiny transformer loaded here has **randomly-initialized, untrained weights** and
+//! no domain-specific head (no ImageNet classifier, no COCO detection head, no ASR
+//! decoder, no sentiment classifier). `rtx-inference` provides a generic transformer
+//! serving runtime, not trained vision/audio/NLP models. So:
+//!
+//! - For [`ModelCategory::TextGeneration`], the real engine output (`output_tokens`) is
+//!   reported directly as the generated token sequence — this is a real forward pass,
+//!   but the tokens are meaningless because the weights are untrained and there is no
+//!   real vocabulary/tokenizer wired up for this demo model.
+//! - For [`ModelCategory::ImageClassification`], [`ModelCategory::ObjectDetection`],
+//!   [`ModelCategory::Segmentation`], [`ModelCategory::SpeechRecognition`], and
+//!   [`ModelCategory::NLP`], the demo still drives the *same* real engine and real
+//!   forward pass (so `inference_time_ms` reflects genuine compute), but the domain
+//!   output (class names, bounding boxes, transcripts, sentiment, ...) is derived
+//!   deterministically from the real output token ids via simple, explicitly-labeled
+//!   post-processing. This is a **toy real-compute proxy, not a trained domain model** —
+//!   it proves the compute path is real without pretending to be a calibrated
+//!   classifier/detector/ASR model. Every such output JSON includes a
+//!   `"note"` field saying so.
+//!
+//! # Upstream inference path (fixed 2026-07-09)
+//!
+//! Earlier versions of `rtx-inference` failed every `infer()` call because
+//! `ForwardPass::apply_embedding` misused `Tensor::gather` for the embedding lookup
+//! (dropping the hidden dimension) and required a `final.weight` norm entry that
+//! common fixtures did not provide. Both are fixed upstream (embedding_lookup is used
+//! now, and the final norm accepts `norm.weight`/`ln_f.weight` aliases), so
+//! [`RealInferenceEngine::run_inference`] performs the full real forward pass
+//! end-to-end: real engine, real loaded weights, real `Instant`-measured elapsed
+//! time, real output token ids feeding the per-category toy proxies.
+
+use crate::error::ModelZooError;
+use model_zoo_shared::{InferenceRequest, InferenceResult, ModelCategory};
+use rtx_inference::{InferenceEngine, InferenceEngineConfig, ModelConfig};
+use rtx_tensor::{Device, Tensor};
+use std::collections::HashMap;
+use std::time::Instant;
+use tokio::runtime::Runtime;
+
+/// Name under which the tiny demo transformer is registered with the engine.
+const DEMO_MODEL_NAME: &str = "model_zoo_demo_transformer";
+
+/// Tiny transformer dimensions — kept small so the demo builds/runs fast while still
+/// exercising a real multi-layer forward pass.
+const VOCAB_SIZE: usize = 256;
+const HIDDEN_SIZE: usize = 64;
+const NUM_LAYERS: usize = 2;
+const NUM_HEADS: usize = 4;
+const MAX_POSITION_EMBEDDINGS: usize = 128;
+
+/// Real inference engine wrapper for the model zoo demo.
+///
+/// Wraps a genuine `rtx_inference::InferenceEngine` (async, tokio-based) behind a
+/// synchronous API matching the previous `MockInferenceEngine` call sites, using an
+/// internal single-threaded tokio runtime to drive the async engine calls.
+pub struct RealInferenceEngine {
+    runtime: Runtime,
+    engine: InferenceEngine,
+    device_label: String,
+}
+
+impl RealInferenceEngine {
+    /// Build a tiny real transformer and load it into a real `InferenceEngine`.
+    pub fn new() -> Result<Self, ModelZooError> {
+        let runtime = Runtime::new().map_err(|e| ModelZooError::ConfigError {
+            message: format!("Failed to start tokio runtime: {e}"),
+        })?;
+
+        let (mut engine, device_label) = runtime.block_on(async {
+            let config = InferenceEngineConfig::default();
+            let engine =
+                InferenceEngine::new(config)
+                    .await
+                    .map_err(|e| ModelZooError::ConfigError {
+                        message: format!("Failed to create InferenceEngine: {e}"),
+                    })?;
+            let device_label = format!("{:?}", Device::cpu());
+            Ok::<_, ModelZooError>((engine, device_label))
+        })?;
+
+        let (weights, model_config) = Self::build_tiny_model();
+        runtime
+            .block_on(engine.load_model(DEMO_MODEL_NAME, &weights, &model_config))
+            .map_err(|e| ModelZooError::ConfigError {
+                message: format!("Failed to load demo model: {e}"),
+            })?;
+
+        Ok(Self {
+            runtime,
+            engine,
+            device_label,
+        })
+    }
+
+    /// Build a tiny transformer's worth of randomly-initialized weights, matching the
+    /// layout `rtx-inference`'s model loader expects (see
+    /// `crates/production/rtx-inference/tests/model_loading_test.rs::create_test_model`).
+    fn build_tiny_model() -> (HashMap<String, Tensor>, ModelConfig) {
+        let config = ModelConfig {
+            vocab_size: VOCAB_SIZE,
+            hidden_size: HIDDEN_SIZE,
+            num_layers: NUM_LAYERS,
+            num_heads: NUM_HEADS,
+            max_position_embeddings: MAX_POSITION_EMBEDDINGS,
+            layer_norm_epsilon: 1e-6,
+        };
+
+        let mut weights = HashMap::new();
+        let device = Device::cpu();
+
+        let embedding = Tensor::randn(&[config.vocab_size, config.hidden_size], &device)
+            .expect("failed to allocate embedding tensor");
+        weights.insert("embedding.weight".to_string(), embedding);
+
+        for layer_idx in 0..config.num_layers {
+            let attn_q = Tensor::randn(&[config.hidden_size, config.hidden_size], &device)
+                .expect("failed to allocate q_proj tensor");
+            let attn_k = Tensor::randn(&[config.hidden_size, config.hidden_size], &device)
+                .expect("failed to allocate k_proj tensor");
+            let attn_v = Tensor::randn(&[config.hidden_size, config.hidden_size], &device)
+                .expect("failed to allocate v_proj tensor");
+            let attn_o = Tensor::randn(&[config.hidden_size, config.hidden_size], &device)
+                .expect("failed to allocate o_proj tensor");
+
+            weights.insert(format!("layers.{layer_idx}.attention.q_proj.weight"), attn_q);
+            weights.insert(format!("layers.{layer_idx}.attention.k_proj.weight"), attn_k);
+            weights.insert(format!("layers.{layer_idx}.attention.v_proj.weight"), attn_v);
+            weights.insert(format!("layers.{layer_idx}.attention.o_proj.weight"), attn_o);
+
+            let mlp_gate = Tensor::randn(&[config.hidden_size, config.hidden_size * 4], &device)
+                .expect("failed to allocate gate_proj tensor");
+            let mlp_up = Tensor::randn(&[config.hidden_size, config.hidden_size * 4], &device)
+                .expect("failed to allocate up_proj tensor");
+            let mlp_down = Tensor::randn(&[config.hidden_size * 4, config.hidden_size], &device)
+                .expect("failed to allocate down_proj tensor");
+
+            weights.insert(format!("layers.{layer_idx}.mlp.gate_proj.weight"), mlp_gate);
+            weights.insert(format!("layers.{layer_idx}.mlp.up_proj.weight"), mlp_up);
+            weights.insert(format!("layers.{layer_idx}.mlp.down_proj.weight"), mlp_down);
+
+            let ln1 = Tensor::ones(&[config.hidden_size], &device)
+                .expect("failed to allocate input_layernorm tensor");
+            let ln2 = Tensor::ones(&[config.hidden_size], &device)
+                .expect("failed to allocate post_attention_layernorm tensor");
+            weights.insert(format!("layers.{layer_idx}.input_layernorm.weight"), ln1);
+            weights.insert(
+                format!("layers.{layer_idx}.post_attention_layernorm.weight"),
+                ln2,
+            );
+        }
+
+        let final_norm = Tensor::ones(&[config.hidden_size], &device)
+            .expect("failed to allocate final norm tensor");
+        weights.insert("norm.weight".to_string(), final_norm);
+
+        let lm_head = Tensor::randn(&[config.hidden_size, config.vocab_size], &device)
+            .expect("failed to allocate lm_head tensor");
+        weights.insert("lm_head.weight".to_string(), lm_head);
+
+        (weights, config)
+    }
+
+    /// Byte-level tokenization into the tiny demo vocabulary. Deterministic and real
+    /// (no randomness), but not a trained tokenizer.
+    fn tokenize(input: &str) -> Vec<i32> {
+        input
+            .bytes()
+            .map(|b| i32::from(b) % VOCAB_SIZE as i32)
+            .collect()
+    }
+
+    /// Run a real forward pass through the engine and return the raw output tokens plus
+    /// the measured wall-clock time of the `infer` call itself.
+    fn real_forward(
+        &self,
+        input: &str,
+        max_new_tokens: usize,
+    ) -> Result<(Vec<i32>, f64), ModelZooError> {
+        let mut tokens = Self::tokenize(input);
+        if tokens.is_empty() {
+            tokens.push(0);
+        }
+        // Keep prompts within the tiny model's context window.
+        tokens.truncate(MAX_POSITION_EMBEDDINGS / 2);
+
+        let request = rtx_inference::InferenceRequest::new(
+            DEMO_MODEL_NAME.to_string(),
+            tokens,
+            max_new_tokens,
+        );
+
+        let start = Instant::now();
+        let outcome = self.runtime.block_on(self.engine.infer(request));
+        let elapsed_ms = start.elapsed().as_secs_f64() * 1000.0;
+
+        let result = outcome.map_err(|e| ModelZooError::InferenceFailed {
+            reason: format!(
+                "real engine.infer() call failed after {elapsed_ms:.4}ms (measured): {e}. \
+                 See RealInferenceEngine module docs — this is a known pre-existing upstream \
+                 bug in rtx-inference/rtx-tensor's embedding gather, not a bug in this demo."
+            ),
+        })?;
+
+        Ok((result.output_tokens, elapsed_ms))
+    }
+
+    /// Fetch real model metadata (parameter count, memory usage, layer count) from the
+    /// loaded demo model via the real engine. Model loading works correctly today —
+    /// only the generation forward pass hits the upstream bug documented above.
+    pub fn model_info(&self) -> Result<rtx_inference::ModelInfo, ModelZooError> {
+        self.runtime
+            .block_on(self.engine.get_model_info(DEMO_MODEL_NAME))
+            .map_err(|e| ModelZooError::ConfigError {
+                message: format!("{e}"),
+            })
+    }
+
+    /// Run inference for the given request/category, matching the previous
+    /// `MockInferenceEngine::run_inference` signature and error behavior.
+    pub fn run_inference(
+        &self,
+        request: &InferenceRequest,
+        category: ModelCategory,
+    ) -> Result<InferenceResult, ModelZooError> {
+        if request.input_data.is_empty() {
+            return Err(ModelZooError::InvalidInput {
+                reason: "Input data is empty".to_string(),
+            });
+        }
+
+        let (output_tokens, inference_time_ms) = self.real_forward(&request.input_data, 8)?;
+
+        let output = match category {
+            ModelCategory::TextGeneration => {
+                Self::text_generation_output(&request.input_data, &output_tokens)
+            }
+            ModelCategory::ImageClassification => Self::classification_output(&output_tokens),
+            ModelCategory::ObjectDetection => Self::detection_output(&output_tokens),
+            ModelCategory::Segmentation => Self::segmentation_output(&output_tokens),
+            ModelCategory::SpeechRecognition => Self::speech_recognition_output(&output_tokens),
+            ModelCategory::NLP => Self::nlp_output(&request.input_data, &output_tokens),
+        };
+
+        Ok(InferenceResult {
+            model_id: request.model_id.clone(),
+            output,
+            inference_time_ms,
+            device: self.device_label.clone(),
+        })
+    }
+
+    /// Real-compute path: reports the actual generated token ids from the real forward
+    /// pass. Tokens are meaningless (untrained weights, no real vocabulary) but the
+    /// compute and timing are genuine.
+    fn text_generation_output(input: &str, output_tokens: &[i32]) -> String {
+        format!(
+            "{{\"generated_token_ids\": {output_tokens:?}, \"prompt\": \"{input}\", \"note\": \"real transformer forward pass with untrained demo weights; token ids are not a trained vocabulary\"}}"
+        )
+    }
+
+    /// Toy real-compute proxy: derives a class + pseudo-confidence deterministically
+    /// from real output token ids. Not a trained classifier.
+    fn classification_output(output_tokens: &[i32]) -> String {
+        let classes = [
+            "cat", "dog", "bird", "car", "person", "horse", "airplane", "boat", "bottle", "chair",
+        ];
+        let sum: i64 = output_tokens.iter().map(|&t| i64::from(t)).sum();
+        let idx = (sum.unsigned_abs() as usize) % classes.len();
+        let class = classes[idx];
+        let confidence = 0.5 + (sum.unsigned_abs() as f64 % 50.0) / 100.0;
+        format!(
+            "{{\"class\": \"{class}\", \"confidence\": {confidence:.4}, \"note\": \"toy real-compute proxy derived from real transformer output tokens; not a trained image classifier\"}}"
+        )
+    }
+
+    /// Toy real-compute proxy: derives detections deterministically from real output
+    /// token ids. Not a trained detector.
+    fn detection_output(output_tokens: &[i32]) -> String {
+        let classes = ["person", "car", "dog", "cat", "chair", "bottle"];
+        let mut detections = Vec::new();
+        for (i, &tok) in output_tokens.iter().enumerate() {
+            let t = i64::from(tok).unsigned_abs() as usize;
+            let class = classes[t % classes.len()];
+            let confidence = 0.7 + (t % 28) as f64 / 100.0;
+            let x = (t % 80) as f64 / 100.0;
+            let y = ((t / 2) % 80) as f64 / 100.0;
+            let w = 0.1 + (t % 20) as f64 / 100.0;
+            let h = 0.1 + ((t / 3) % 20) as f64 / 100.0;
+            detections.push(format!(
+                "{{\"class\": \"{class}\", \"confidence\": {confidence:.4}, \"bbox\": [{x:.4}, {y:.4}, {w:.4}, {h:.4}]}}"
+            ));
+            if i >= 4 {
+                break;
+            }
+        }
+        format!(
+            "{{\"detections\": [{}], \"note\": \"toy real-compute proxy derived from real transformer output tokens; not a trained object detector\"}}",
+            detections.join(", ")
+        )
+    }
+
+    /// Toy real-compute proxy: derives pseudo-segments deterministically from real
+    /// output token ids. Not a trained segmentation model.
+    fn segmentation_output(output_tokens: &[i32]) -> String {
+        let classes = [
+            "background",
+            "person",
+            "car",
+            "road",
+            "building",
+            "vegetation",
+            "sky",
+        ];
+        let mut segments = Vec::new();
+        for (i, &tok) in output_tokens.iter().enumerate() {
+            let t = i64::from(tok).unsigned_abs() as usize;
+            let class = classes[i % classes.len()];
+            let pixel_count = 1000 + (t % 49000);
+            let percentage = 5.0 + (t % 20) as f64;
+            segments.push(format!(
+                "{{\"class\": \"{class}\", \"pixels\": {pixel_count}, \"percentage\": {percentage:.2}}}"
+            ));
+        }
+        format!(
+            "{{\"segments\": [{}], \"note\": \"toy real-compute proxy derived from real transformer output tokens; not a trained segmentation model\"}}",
+            segments.join(", ")
+        )
+    }
+
+    /// Toy real-compute proxy: selects a canned transcript deterministically from real
+    /// output token ids. Not a trained ASR model.
+    fn speech_recognition_output(output_tokens: &[i32]) -> String {
+        let transcripts = [
+            "Hello, how are you doing today?",
+            "The weather is beautiful this morning.",
+            "I would like to schedule a meeting for tomorrow.",
+            "Can you please help me with this task?",
+            "Thank you very much for your assistance.",
+            "The quick brown fox jumps over the lazy dog.",
+        ];
+        let sum: i64 = output_tokens.iter().map(|&t| i64::from(t)).sum();
+        let idx = (sum.unsigned_abs() as usize) % transcripts.len();
+        let confidence = 0.85 + (sum.unsigned_abs() as f64 % 14.0) / 100.0;
+        format!(
+            "{{\"transcript\": \"{}\", \"confidence\": {confidence:.4}, \"note\": \"toy real-compute proxy derived from real transformer output tokens; not a trained ASR model\"}}",
+            transcripts[idx]
+        )
+    }
+
+    /// Toy real-compute proxy: derives sentiment/entities deterministically from real
+    /// output token ids. Not a trained NLP model.
+    fn nlp_output(input: &str, output_tokens: &[i32]) -> String {
+        let sentiments = ["positive", "negative", "neutral"];
+        let sum: i64 = output_tokens.iter().map(|&t| i64::from(t)).sum();
+        let sentiment = sentiments[(sum.unsigned_abs() as usize) % sentiments.len()];
+        let score = 0.6 + (sum.unsigned_abs() as f64 % 35.0) / 100.0;
+        let entities_count = (sum.unsigned_abs() as usize) % 5;
+        format!(
+            "{{\"sentiment\": \"{sentiment}\", \"score\": {score:.4}, \"entities\": {entities_count}, \"tokens\": {}, \"note\": \"toy real-compute proxy derived from real transformer output tokens; not a trained sentiment model\"}}",
+            input.split_whitespace().count()
+        )
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn make_engine() -> RealInferenceEngine {
+        RealInferenceEngine::new().expect("engine should build and load the demo model")
+    }
+
+    #[test]
+    fn test_engine_creation() {
+        let _engine = make_engine();
+    }
+
+    #[test]
+    fn test_inference_empty_input() {
+        let engine = make_engine();
+        let request = InferenceRequest {
+            model_id: "resnet50".to_string(),
+            input_data: String::new(),
+            input_type: "image/jpeg".to_string(),
+        };
+        let result = engine.run_inference(&request, ModelCategory::ImageClassification);
+        assert!(result.is_err());
+        match result {
+            Err(ModelZooError::InvalidInput { reason }) => {
+                assert_eq!(reason, "Input data is empty");
+            }
+            _ => panic!("Expected InvalidInput error"),
+        }
+    }
+
+    // The upstream embedding/final-norm bugs are fixed (see module docs), so these
+    // tests assert real end-to-end success: a genuine engine.infer() forward pass
+    // producing output that feeds the per-category proxies.
+
+    #[test]
+    fn test_inference_image_classification_real_engine_succeeds() {
+        let engine = make_engine();
+        let request = InferenceRequest {
+            model_id: "resnet50".to_string(),
+            input_data: "base64encodedimage".to_string(),
+            input_type: "image/jpeg".to_string(),
+        };
+        let result = engine
+            .run_inference(&request, ModelCategory::ImageClassification)
+            .expect("real engine inference should succeed");
+        assert!(result.inference_time_ms > 0.0);
+        assert!(result.output.contains("\"class\""));
+    }
+
+    #[test]
+    fn test_inference_text_generation_real_engine_succeeds() {
+        let engine = make_engine();
+        let request = InferenceRequest {
+            model_id: "gpt2_small".to_string(),
+            input_data: "Machine learning".to_string(),
+            input_type: "text/plain".to_string(),
+        };
+        let result = engine
+            .run_inference(&request, ModelCategory::TextGeneration)
+            .expect("real engine inference should succeed");
+        assert!(result.inference_time_ms > 0.0);
+        assert!(!result.output.is_empty());
+    }
+
+    #[test]
+    fn test_multiple_inference_attempts_same_engine_are_consistent() {
+        let engine = make_engine();
+        for _ in 0..5 {
+            let request = InferenceRequest {
+                model_id: "resnet50".to_string(),
+                input_data: "base64encodedimage".to_string(),
+                input_type: "image/jpeg".to_string(),
+            };
+            // Every attempt is a genuine, independent call through the real engine.
+            let result = engine
+                .run_inference(&request, ModelCategory::ImageClassification)
+                .expect("real engine inference should succeed");
+            assert!(result.inference_time_ms > 0.0);
+        }
+    }
+
+    #[test]
+    fn test_model_info_reports_real_parameters_and_memory() {
+        let engine = make_engine();
+        let info = engine
+            .model_info()
+            .expect("model loading + metadata retrieval works even though generation is broken");
+        assert_eq!(info.name, DEMO_MODEL_NAME);
+        assert_eq!(info.config.vocab_size, VOCAB_SIZE);
+        assert_eq!(info.config.hidden_size, HIDDEN_SIZE);
+        assert!(info.parameter_count > 0, "should report real parameter count");
+        assert!(info.memory_usage > 0, "should report real memory usage");
+    }
+
+    #[test]
+    fn test_classification_output_format() {
+        let output = RealInferenceEngine::classification_output(&[1, 2, 3]);
+        assert!(output.contains("class"));
+        assert!(output.contains("confidence"));
+        assert!(output.contains("note"));
+    }
+
+    #[test]
+    fn test_detection_output_format() {
+        let output = RealInferenceEngine::detection_output(&[1, 2, 3]);
+        assert!(output.contains("detections"));
+        assert!(output.contains("bbox"));
+    }
+
+    #[test]
+    fn test_segmentation_output_format() {
+        let output = RealInferenceEngine::segmentation_output(&[1, 2, 3]);
+        assert!(output.contains("segments"));
+        assert!(output.contains("pixels"));
+    }
+
+    #[test]
+    fn test_text_generation_output_format() {
+        let output = RealInferenceEngine::text_generation_output("Machine learning", &[1, 2, 3]);
+        assert!(output.contains("generated_token_ids"));
+        assert!(output.contains("Machine learning"));
+    }
+
+    #[test]
+    fn test_speech_recognition_output_format() {
+        let output = RealInferenceEngine::speech_recognition_output(&[1, 2, 3]);
+        assert!(output.contains("transcript"));
+        assert!(output.contains("confidence"));
+    }
+
+    #[test]
+    fn test_nlp_output_format() {
+        let output = RealInferenceEngine::nlp_output("This is a test sentence", &[1, 2, 3]);
+        assert!(output.contains("sentiment"));
+        assert!(output.contains("score"));
+        assert!(output.contains("entities"));
+        assert!(output.contains("tokens"));
+    }
+}