Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions scripts/generate_dashboard.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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, [])
Expand All @@ -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
Expand Down Expand Up @@ -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),
Expand Down
38 changes: 33 additions & 5 deletions scripts/generate_golden.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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,
}


Expand Down
129 changes: 127 additions & 2 deletions src/mobius/_testing/generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
)

Comment on lines +234 to +253

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

KV cache initialization in OnnxSeq2SeqGenerator hard-codes (batch_size, config.num_key_value_heads, 0, config.head_dim) for every past_key_values.* input. However Seq2SeqTask creates caches using num_attention_heads (and dtype from the model), and naming is split into .self. and .cross.. To avoid head-count mismatches and future shape variations, initialize caches from dec_session.get_input_shape(name) (like _make_empty_kv_cache() does in tests/e2e_golden_test.py) and only override the sequence dimension (0 for self; encoder seq len for cross if needed).

Suggested change
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,
)
past_kv: dict[str, np.ndarray] = {}
for name in self.dec_session.input_names:
if not name.startswith("past_key_values."):
continue
# Build the cache tensor from the decoder input contract so the
# helper follows the model's declared head count and layout.
cache_shape = list(self.dec_session.get_input_shape(name))
if len(cache_shape) >= 1:
cache_shape[0] = batch_size
if len(cache_shape) >= 3:
if ".cross." in name:
# Cross-attention cache is keyed by encoder time steps.
cache_shape[2] = src_seq_len
else:
# Self-attention cache starts empty and grows each step.
cache_shape[2] = 0
past_kv[name] = np.zeros(tuple(cache_shape), dtype=np.float32)

Copilot uses AI. Check for mistakes.
# 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,
}
Comment on lines +258 to +264

Copilot AI Apr 22, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

OnnxSeq2SeqGenerator.generate() feeds the decoder attention_mask with shape [batch, src_seq_len] and never grows it across decode steps. In Seq2SeqTask, the decoder input attention_mask is defined as [batch, past_seq_len + dec_seq_len], i.e. it should reflect the decoder time dimension (like OnnxGenerator does for causal LM). This can break models that use the decoder mask (and can also fail shape checks depending on runtime). Build attention_mask from the current decoder length (past_seq_len + cur_dec_ids.shape[1]) and update past_seq_len each step.

Copilot uses AI. Check for mistakes.

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,
Expand Down
6 changes: 6 additions & 0 deletions src/mobius/_testing/golden.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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,
)
Expand Down
59 changes: 51 additions & 8 deletions src/mobius/_testing/torch_reference.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)}")


# ---------------------------------------------------------------------------
Expand Down
3 changes: 2 additions & 1 deletion testdata/cases/causal-lm/apertus-8b.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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)."
1 change: 0 additions & 1 deletion testdata/cases/causal-lm/arcee-afm-4.5b.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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."
2 changes: 1 addition & 1 deletion testdata/cases/causal-lm/cohere-r7b.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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."
2 changes: 1 addition & 1 deletion testdata/cases/causal-lm/cohere2-r7b.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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."
3 changes: 2 additions & 1 deletion testdata/cases/causal-lm/ernie4_5-21b-moe.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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."
2 changes: 1 addition & 1 deletion testdata/cases/causal-lm/flex-olmo-2x7b.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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)."
2 changes: 1 addition & 1 deletion testdata/cases/causal-lm/gemma-4-e4b.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,11 @@ inputs:
prompts:
- "Here is my poem:"

skip_reason: "Weight shape mismatch in audio_encoder.encoder.output_proj."
level: "L4+L5"

generation:
max_new_tokens: 20
do_sample: false

skip_reason: "Gated repo (google/gemma-4-E4B-it requires authentication)."
notes: "Gemma 4 E4B Any-to-Any (text-only path). Dual RoPE, KV sharing, double-wide MLP, sliding+global attention, logit softcapping, audio encoder."
Loading
Loading