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]>

crates/production/rtx-inference/src/engine/forward_pass.rs

@@ -68,7 +68,10 @@ impl ForwardPass {
             .or_else(|| model_weights.get("wte.weight"))
             .ok_or_else(|| InferenceError::model_not_found("embedding weights"))?;
 
-        let embeddings = embedding_weight.gather(0, input_ids)?;
+        // Embedding lookup: [vocab, hidden] indexed by [batch, seq] ->
+        // [batch, seq, hidden]. (gather would return the indices' shape and
+        // drop the hidden dim.)
+        let embeddings = embedding_weight.embedding_lookup(input_ids)?;
 
         debug!(
             "Applied embedding layer: input shape {:?} -> output shape {:?}",
@@ -140,15 +143,18 @@ impl ForwardPass {
     ) -> InferenceResult<Tensor> {
         let head_dim = config.hidden_size / config.num_heads;
 
-        // Get Q, K, V projection weights
-        let q_weight = model_weights
-            .get(&format!("layers.{layer_idx}.self_attn.q_proj.weight"))
+        // Get Q, K, V projection weights. Accept both the HF-LLaMA naming
+        // (`self_attn.`) and the plain `attention.` prefix.
+        let get_attn_weight = |proj: &str| {
+            model_weights
+                .get(&format!("layers.{layer_idx}.self_attn.{proj}.weight"))
+                .or_else(|| model_weights.get(&format!("layers.{layer_idx}.attention.{proj}.weight")))
+        };
+        let q_weight = get_attn_weight("q_proj")
             .ok_or_else(|| InferenceError::model_not_found("Q projection weights"))?;
-        let k_weight = model_weights
-            .get(&format!("layers.{layer_idx}.self_attn.k_proj.weight"))
+        let k_weight = get_attn_weight("k_proj")
             .ok_or_else(|| InferenceError::model_not_found("K projection weights"))?;
-        let v_weight = model_weights
-            .get(&format!("layers.{layer_idx}.self_attn.v_proj.weight"))
+        let v_weight = get_attn_weight("v_proj")
             .ok_or_else(|| InferenceError::model_not_found("V projection weights"))?;
 
         // Project to Q, K, V
@@ -192,8 +198,7 @@ impl ForwardPass {
                 .transpose(1, 2)?
                 .view([batch_size, seq_len, config.hidden_size])?;
 
-        let o_weight = model_weights
-            .get(&format!("layers.{layer_idx}.self_attn.o_proj.weight"))
+        let o_weight = get_attn_weight("o_proj")
             .ok_or_else(|| InferenceError::model_not_found("O projection weights"))?;
 
         let output = attention_concat.matmul(o_weight)?;
@@ -270,8 +275,20 @@ impl ForwardPass {
         model_weights: &HashMap<String, Tensor>,
         layer_name: &str,
     ) -> InferenceResult<Tensor> {
+        // Accept common aliases for the final norm across model conventions:
+        // "final.weight" (ours), "norm.weight" (LLaMA), "ln_f.weight" (GPT-2).
         let weight = model_weights
             .get(&format!("{layer_name}.weight"))
+            .or_else(|| {
+                if layer_name == "final" {
+                    model_weights
+                        .get("norm.weight")
+                        .or_else(|| model_weights.get("model.norm.weight"))
+                        .or_else(|| model_weights.get("ln_f.weight"))
+                } else {
+                    None
+                }
+            })
             .ok_or_else(|| {
                 InferenceError::model_not_found(format!("{layer_name} layer norm weights"))
             })?;