Expand L4/L5 test coverage: golden data, seq2seq generation, ci_skip_reason

scripts/generate_dashboard.py

@@ -83,6 +83,7 @@ class ModelInfo:
     l5_test_case_skipped: bool = False
     yaml_test_case_file: str | None = None
     yaml_test_case_skip_reason: str | None = None
+    yaml_test_case_ci_skip_reason: str | None = None
     yaml_min_token_match_ratio: float | None = None
     # L3 synthetic parity status: "pass", "xfail", "skip", or None
     l3_status: str | None = None
@@ -413,6 +414,7 @@ def _scan_yaml_test_cases(models: dict[str, ModelInfo]) -> None:
         # Skip test cases that are explicitly skipped — they don't count as coverage,
         # but we still record them so the dashboard can show "skipped" status.
         skip_reason = data.get("skip_reason")
+        ci_skip_reason = data.get("ci_skip_reason")
         min_token_match_ratio = data.get("min_token_match_ratio")
         if skip_reason:
             matched_types = model_id_to_types.get(model_id, [])
@@ -435,6 +437,8 @@ def _scan_yaml_test_cases(models: dict[str, ModelInfo]) -> None:
         for model_type in matched_types:
             if model_type in models:
                 models[model_type].yaml_test_case_file = rel_path
+                if ci_skip_reason:
+                    models[model_type].yaml_test_case_ci_skip_reason = ci_skip_reason
                 if min_token_match_ratio is not None:
                     models[model_type].yaml_min_token_match_ratio = float(
                         min_token_match_ratio
@@ -711,6 +715,7 @@ def _render_html(
                 "l3_reason": info.l3_status_reason,
                 "yaml_case": info.yaml_test_case_file,
                 "yaml_skip_reason": info.yaml_test_case_skip_reason,
+                "yaml_ci_skip_reason": info.yaml_test_case_ci_skip_reason,
                 "min_token_match_ratio": info.yaml_min_token_match_ratio,
                 "code_paths": sorted(info.code_paths),
                 "config_overrides": _json_safe(info.config_overrides),

scripts/generate_golden.py

@@ -230,6 +230,15 @@ def _generate_encoder(case: TestCase, json_path: Path, device: str) -> None:
 
     model, tokenizer = load_torch_encoder_model(case.model_id, device=device)
 
+    # CLIP-like multimodal models wrap a text sub-model that can be
+    # called with text-only inputs (pixel_values not required).
+    if hasattr(model, "text_model"):
+        model = model.text_model
+
+    # X-MOD requires setting a default language before inference.
+    if hasattr(model, "set_default_language"):
+        model.set_default_language("en_XX")
+
     encoded = tokenizer(case.prompts[0], return_tensors="np", padding=False)
     input_ids = encoded["input_ids"]
     attention_mask = encoded["attention_mask"]
@@ -267,7 +276,11 @@ def _generate_seq2seq(case: TestCase, json_path: Path, device: str) -> None:
     input_ids = encoded["input_ids"]
 
     # Prepare decoder input (pad token for autoregressive start)
-    decoder_start = np.array([[model.config.decoder_start_token_id or 0]], dtype=np.int64)
+    decoder_start_id = getattr(model.config, "decoder_start_token_id", None)
+    if decoder_start_id is None:
+        generation_config = getattr(model, "generation_config", None)
+        decoder_start_id = getattr(generation_config, "decoder_start_token_id", None) or 0
+    decoder_start = np.array([[decoder_start_id]], dtype=np.int64)
 
     # L4: single forward pass through full model
     torch_device = next(model.parameters()).device
@@ -735,10 +748,20 @@ def _generate_image_classification(case: TestCase, json_path: Path, device: str)
 
     # Forward pass → last_hidden_state
     hidden_states = torch_vision_forward(model, pixel_values)
-    # Use the last patch token rather than the CLS token (index 0) because
-    # patch-based ViT models aggregate spatial context into trailing tokens;
-    # the last token provides a stable, architecture-neutral summary vector.
-    last_hidden = hidden_states[0, -1, :]  # (hidden_size,)
+    # Vision models return different output shapes:
+    # - ViT-like: [B, seq_len, hidden] → select first token (CLS)
+    # - CNN-like (CvT, MobileViT, PVT): [B, C, H, W] → flatten feature map
+    # - Classification head: [B, num_classes] → 1-D logits
+    batch_hidden = hidden_states[0]  # drop batch dim
+    if batch_hidden.ndim == 2:
+        # (seq_len, hidden) — take CLS token
+        last_hidden = batch_hidden[0]
+    elif batch_hidden.ndim >= 3:
+        # (C, H, W) feature map — flatten
+        last_hidden = batch_hidden.reshape(-1)
+    else:
+        # 1-D logits or already flat
+        last_hidden = batch_hidden
     golden = _extract_logits_golden(last_hidden)
 
     # Image classification is L4-only (no generation)
@@ -765,6 +788,11 @@ def _generate_image_classification(case: TestCase, json_path: Path, device: str)
     "speech-to-text": _generate_speech_to_text,
     "speech-language": _generate_speech_language,
     "audio-feature-extraction": _generate_audio_feature_extraction,
+    # Vision tasks that produce last_hidden_state — reuse image classification.
+    "depth-estimation": _generate_image_classification,
+    "image-segmentation": _generate_image_classification,
+    "image-to-image": _generate_image_classification,
+    "object-detection": _generate_image_classification,
 }
 
 

src/mobius/_configs.py

@@ -271,6 +271,7 @@ class Gemma4AudioConfig(AudioConfig):
     hidden_size: int = 1024
     subsampling_conv_channels: list[int] | None = None
     use_causal_chunked_attn: bool = False
+    output_proj_dims: int | None = None
 
 
 def _first_not_none(*values, default=None):
@@ -627,6 +628,7 @@ def _extract_audio_config(config, parent_config, model_type: str) -> dict:
                     ),
                     use_causal_chunked_attn=getattr(ac, "use_causal_chunked_attn", False),
                     output_dim=getattr(ac, "output_dim", None),
+                    output_proj_dims=getattr(ac, "output_proj_dims", None),
                     audio_token_id=getattr(composite, "audio_token_id", None),
                 )
             }

src/mobius/_testing/generation.py

@@ -3,8 +3,9 @@
 
 """Autoregressive text generation using ONNX Runtime.
 
-Provides a self-contained generation loop with KV cache management,
-without depending on onnxruntime-genai.
+Provides self-contained generation loops with KV cache management,
+without depending on onnxruntime-genai.  Supports causal-LM (decoder-
+only) and seq2seq (encoder-decoder) architectures.
 """
 
 from __future__ import annotations
@@ -164,6 +165,130 @@ def generate(
         return all_ids
 
 
+class OnnxSeq2SeqGenerator:
+    """Greedy generation for encoder-decoder (seq2seq) ONNX models.
+
+    Runs the encoder once, then autoregressively decodes using the
+    decoder with cross-attention to encoder hidden states.  Manages
+    both self-attention and cross-attention KV caches.
+
+    Example::
+
+        enc_session = OnnxModelSession(pkg["encoder"])
+        dec_session = OnnxModelSession(pkg["decoder"])
+        gen = OnnxSeq2SeqGenerator(enc_session, dec_session, config)
+        output_ids = gen.generate(input_ids, max_new_tokens=20)
+    """
+
+    def __init__(
+        self,
+        enc_session: OnnxModelSession,
+        dec_session: OnnxModelSession,
+        config: ArchitectureConfig,
+    ):
+        self.enc_session = enc_session
+        self.dec_session = dec_session
+        self.config = config
+
+    def generate(
+        self,
+        input_ids: np.ndarray,
+        max_new_tokens: int = 20,
+        eos_token_id: int | None = None,
+        decoder_start_token_id: int = 0,
+    ) -> np.ndarray:
+        """Generate tokens from encoder input using greedy decoding.
+
+        Args:
+            input_ids: [batch, src_seq_len] int64 encoder input tokens.
+            max_new_tokens: Maximum number of tokens to generate.
+            eos_token_id: If set, stop when this token is produced.
+            decoder_start_token_id: Token to seed the decoder.
+
+        Returns:
+            [batch, generated_len] int64 array of generated token IDs
+            (decoder start token + generated, no encoder input).
+        """
+        batch_size = input_ids.shape[0]
+        src_seq_len = input_ids.shape[1]
+
+        # Step 1: Run encoder once
+        enc_feeds = {
+            "input_ids": input_ids,
+            "attention_mask": np.ones_like(input_ids),
+        }
+        enc_outputs = self.enc_session.run(enc_feeds)
+
+        # Extract encoder hidden states
+        enc_hidden = None
+        for key in ("encoder_hidden_states", "last_hidden_state"):
+            if key in enc_outputs:
+                enc_hidden = enc_outputs[key]
+                break
+        if enc_hidden is None:
+            raise KeyError(
+                f"Encoder output missing hidden states. Keys: {sorted(enc_outputs.keys())}"
+            )
+
+        # Step 2: Initialize decoder KV caches
+        num_kv_heads = self.config.num_key_value_heads
+        head_dim = self.config.head_dim
+        past_kv: dict[str, np.ndarray] = {}
+
+        for name in self.dec_session.input_names:
+            if not name.startswith("past_key_values."):
+                continue
+            if ".cross." in name:
+                # Cross-attention cache: starts empty, populated on first step
+                past_kv[name] = np.zeros(
+                    (batch_size, num_kv_heads, 0, head_dim),
+                    dtype=np.float32,
+                )
+            else:
+                # Self-attention cache: grows each step
+                past_kv[name] = np.zeros(
+                    (batch_size, num_kv_heads, 0, head_dim),
+                    dtype=np.float32,
+                )
+
+        # Step 3: Autoregressive decode loop
+        cur_dec_ids = np.full((batch_size, 1), decoder_start_token_id, dtype=np.int64)
+        all_ids = cur_dec_ids.copy()
+
+        for _step in range(max_new_tokens):
+            dec_feeds: dict[str, np.ndarray] = {
+                "input_ids": cur_dec_ids,
+                "encoder_hidden_states": enc_hidden,
+                "attention_mask": np.ones((batch_size, src_seq_len), dtype=np.int64),
+                **past_kv,
+            }
+
+            outputs = self.dec_session.run(dec_feeds)
+
+            # Extract logits and take argmax of last token
+            logits = outputs["logits"]  # [batch, dec_seq_len, vocab]
+            next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True)
+
+            all_ids = np.concatenate([all_ids, next_token], axis=1)
+
+            # Check EOS
+            if eos_token_id is not None and np.all(next_token == eos_token_id):
+                break
+
+            # Update KV caches from present outputs
+            for name in list(past_kv.keys()):
+                # past_key_values.N.self.key → present.N.self.key
+                layer_suffix = name.replace("past_key_values.", "")
+                present_name = f"present.{layer_suffix}"
+                if present_name in outputs:
+                    past_kv[name] = outputs[present_name]
+
+            # Next step: only the new token
+            cur_dec_ids = next_token.astype(np.int64)
+
+        return all_ids
+
+
 def torch_generate_greedy(
     model,
     input_ids: np.ndarray,

src/mobius/_testing/golden.py

@@ -100,6 +100,11 @@ class GoldenTestCase:
     skip_reason: str | None
     """If set, the test runner should skip with this message."""
 
+    ci_skip_reason: str | None
+    """If set, the test is skipped in CI (GITHUB_ACTIONS=true) but runs
+    locally.  Used for models that are too large for CI hardware but can
+    be tested on a local GPU.  Does NOT block golden data generation."""
+
     min_token_match_ratio: float | None
     """Per-case override for the L5 token match tolerance (0-1).
     When ``None``, the global tolerance from ``default_tolerances.yaml`` is used.
@@ -225,6 +230,7 @@ def load_test_case(yaml_path: Path) -> GoldenTestCase:
         generation_params=generation,
         trust_remote_code=data.get("trust_remote_code", False),
         skip_reason=data.get("skip_reason"),
+        ci_skip_reason=data.get("ci_skip_reason"),
         min_token_match_ratio=data.get("min_token_match_ratio"),
         yaml_path=yaml_path,
     )

src/mobius/_testing/torch_reference.py

@@ -374,12 +374,33 @@ def load_torch_vision_model(
     processor = transformers.AutoImageProcessor.from_pretrained(
         model_id, trust_remote_code=trust_remote_code
     )
-    model = transformers.AutoModel.from_pretrained(
-        model_id,
-        torch_dtype=dtype,
-        device_map=device,
-        trust_remote_code=trust_remote_code,
-    )
+
+    # Some vision models (DepthAnything, Segformer) aren't loadable via
+    # AutoModel.  Try progressively more specific Auto classes.
+    model = None
+    auto_classes = [
+        transformers.AutoModel,
+        transformers.AutoModelForImageClassification,
+    ]
+    if hasattr(transformers, "AutoModelForDepthEstimation"):
+        auto_classes.append(transformers.AutoModelForDepthEstimation)
+    if hasattr(transformers, "AutoModelForSemanticSegmentation"):
+        auto_classes.append(transformers.AutoModelForSemanticSegmentation)
+    if hasattr(transformers, "AutoModelForImageToImage"):
+        auto_classes.append(transformers.AutoModelForImageToImage)
+    for auto_cls in auto_classes:
+        try:
+            model = auto_cls.from_pretrained(
+                model_id,
+                torch_dtype=dtype,
+                device_map=device,
+                trust_remote_code=trust_remote_code,
+            )
+            break
+        except (ValueError, TypeError):
+            continue
+    if model is None:
+        raise ValueError(f"Could not load {model_id} with any AutoModel variant")
     model.eval()
 
     # Multi-modal models (CLIP, SigLIP) wrap a vision sub-model that
@@ -398,13 +419,35 @@ def torch_vision_forward(
     """Run a single forward pass on a HuggingFace vision model.
 
     Returns:
-        last_hidden_state as numpy array [batch, seq_len, hidden_size].
+        Feature tensor as numpy array.  Shape varies by model:
+        [B, seq_len, hidden] for ViT-like, [B, C, H, W] for CNN-like,
+        or [B, num_classes] for classification heads.
     """
     device = next(model.parameters()).device
     dtype = next(model.parameters()).dtype
     pv = torch.from_numpy(pixel_values).to(device=device, dtype=dtype)
     outputs = model(pixel_values=pv)
-    return outputs.last_hidden_state.cpu().numpy()
+    # Prefer last_hidden_state; fall back to logits or first tensor output.
+    if hasattr(outputs, "last_hidden_state") and outputs.last_hidden_state is not None:
+        return outputs.last_hidden_state.cpu().numpy()
+    if hasattr(outputs, "logits") and outputs.logits is not None:
+        return outputs.logits.cpu().numpy()
+    # Some models (DepthAnything) return predicted_depth or similar.
+    if hasattr(outputs, "predicted_depth") and outputs.predicted_depth is not None:
+        return outputs.predicted_depth.cpu().numpy()
+    # Generic fallback: accept dict-like ModelOutput, tuple, or list returns.
+    if hasattr(outputs, "cpu"):
+        return outputs.cpu().numpy()
+    values = None
+    if hasattr(outputs, "values"):
+        values = outputs.values()
+    elif isinstance(outputs, (tuple, list)):
+        values = outputs
+    if values is not None:
+        for v in values:
+            if hasattr(v, "cpu"):
+                return v.cpu().numpy()
+    raise ValueError(f"No usable tensor in model outputs: {type(outputs)}")
 
 
 # ---------------------------------------------------------------------------

testdata/cases/causal-lm/apertus-8b.yaml

@@ -13,5 +13,6 @@ generation:
   max_new_tokens: 20
   do_sample: false
 
-skip_reason: "Model is 8B — too large for CI golden data generation."
+skip_reason: "Requires xielu activation not yet in ACT2FN mapping."
+ci_skip_reason: "Model is 8B — too large for CI."
 notes: "Apertus 8B. Uses xIELU (Softplus) activation; small FP accumulation differences vs HF (atol=0.02)."

testdata/cases/causal-lm/arcee-afm-4.5b.yaml

@@ -13,5 +13,4 @@ generation:
   max_new_tokens: 20
   do_sample: false
 
-skip_reason: "Model is 4.5B — too large for CI golden data generation."
 notes: "Arcee Foundation Model 4.5B. Llama-compatible architecture."

testdata/cases/causal-lm/cohere-r7b.yaml

@@ -13,5 +13,5 @@ generation:
   max_new_tokens: 20
   do_sample: false
 
+skip_reason: "Gated repo not accessible with current credentials."
 notes: "Cohere Command R 7B. Llama-compatible with layernorm differences."
-skip_reason: "Gated repo, requires HF token."

testdata/cases/causal-lm/cohere2-r7b.yaml

@@ -13,5 +13,5 @@ generation:
   max_new_tokens: 20
   do_sample: false
 
-skip_reason: "Gated repo, requires HF token."
+skip_reason: "Gated repo not accessible with current credentials."
 notes: "Cohere2 Command R 7B. LayerNorm1P with weight+1 scaling, logit scale 1/16."

testdata/cases/causal-lm/ernie4_5-21b-moe.yaml

@@ -13,5 +13,6 @@ generation:
   max_new_tokens: 20
   do_sample: false
 
-skip_reason: "Model is 21B — too large for CI golden data generation."
+skip_reason: "UngatedSharedMoELayer requires config.shared_expert_intermediate_size."
+ci_skip_reason: "Model is 21B — too large for CI."
 notes: "ERNIE 4.5 MoE (3B active). e_score_correction_bias + shared expert."

testdata/cases/causal-lm/flex-olmo-2x7b.yaml

@@ -13,5 +13,5 @@ generation:
   max_new_tokens: 20
   do_sample: false
 
-skip_reason: "Model is 14B (2x7B MoE) — too large for CI golden data generation."
+ci_skip_reason: "Model is 14B (2x7B MoE) — too large for CI."
 notes: "FlexOLMo 2x7B post-norm MoE. FP accumulation differences vs HF batched dispatch (atol=0.15)."