Skip to content

Add Gemma4 GenAI text generation example with benchmarking - #249

Merged
justinchuby merged 7 commits into
mainfrom
gemma4-genai-example
May 6, 2026
Merged

Add Gemma4 GenAI text generation example with benchmarking#249
justinchuby merged 7 commits into
mainfrom
gemma4-genai-example

Conversation

@justinchuby

Copy link
Copy Markdown
Member

Summary

Add examples/gemma4_genai.py — a streamlined example for Gemma4 text generation via ORT GenAI with performance benchmarking.

Features

  • Build + export: Uses mobius build --runtime ort-genai to create ONNX model with genai_config.json
  • Multi-step generation: Prefill + autoregressive decode with streaming output
  • Performance metrics: Tokens/sec, time-to-first-token (TTFT), decode throughput
  • Device selection: --device cpu or --device cuda
  • Precision: --dtype f32/f16/bf16
  • Model sizes: --model google/gemma-4-e2b-it (default) or any Gemma4 variant
  • Pre-built models: --model-dir to skip export, --save-to for export-only

Usage

# CPU (default)
python examples/gemma4_genai.py

# CUDA with FP16
python examples/gemma4_genai.py --device cuda --dtype f16

# Different model
python examples/gemma4_genai.py --model google/gemma-4-e4b-it --max-new-tokens 200

Relation to existing examples

This is a simplified, benchmark-focused version of gemma4_ort_genai.py. Key differences:

  • Focused on text-only generation (no VLM/image support)
  • Adds performance reporting (tok/s, TTFT)
  • Uses mobius CLI for export instead of inline config generation
  • Simpler code (~250 lines vs ~550 lines)

New example script for Gemma4 text generation via ORT GenAI:
- Builds ONNX model with mobius CLI (--runtime ort-genai)
- Runs multi-step generation with streaming output
- Reports tokens/sec, time-to-first-token, decode throughput
- Supports --device cpu/cuda, --dtype f32/f16/bf16
- Supports --model for different Gemma4 sizes
- Supports --save-to for export-only and --model-dir for pre-built

Usage:
  python examples/gemma4_genai.py --device cpu
  python examples/gemma4_genai.py --device cuda --dtype f16
Signed-off-by: Justin Chu <justinchu@microsoft.com>
@github-actions

github-actions Bot commented May 5, 2026

Copy link
Copy Markdown

Performance Comparison

Comparing c1daeaa7291641

Model Metric Baseline Current Delta
bert (feature-extraction) model_size_bytes 359 KB 359 KB +0.0%
bert (feature-extraction) num_nodes 60 60 +0.0%
falcon model_size_bytes 364 KB 364 KB +0.0%
falcon num_nodes 66 66 +0.0%
gemma2 model_size_bytes 428 KB 428 KB +0.0%
gemma2 num_nodes 107 107 +0.0%
gpt2 model_size_bytes 388 KB 388 KB +0.0%
gpt2 num_nodes 53 53 +0.0%
llama model_size_bytes 425 KB 425 KB +0.0%
llama num_nodes 61 61 +0.0%
llama (static-cache) model_size_bytes 425 KB 425 KB +0.0%
llama (static-cache) num_nodes 58 58 +0.0%
mamba (ssm-text-generation) model_size_bytes 296 KB 296 KB +0.0%
mamba (ssm-text-generation) num_nodes 98 98 +0.0%
phi3 model_size_bytes 421 KB 421 KB +0.0%
phi3 num_nodes 59 59 +0.0%
phi3 (static-cache) model_size_bytes 421 KB 421 KB +0.0%
phi3 (static-cache) num_nodes 56 56 +0.0%
qwen2 model_size_bytes 425 KB 425 KB +0.0%
qwen2 num_nodes 61 61 +0.0%
qwen2 (static-cache) model_size_bytes 425 KB 425 KB +0.0%
qwen2 (static-cache) num_nodes 58 58 +0.0%
qwen3_5_moe (hybrid-text-generation) model_size_bytes 506 KB 506 KB +0.0%
qwen3_5_moe (hybrid-text-generation) num_nodes 275 275 +0.0%
qwen3_5_text (hybrid-text-generation) model_size_bytes 458 KB 458 KB +0.0%
qwen3_5_text (hybrid-text-generation) num_nodes 129 129 +0.0%
qwen3_5_vl (hybrid-qwen-vl) model_size_bytes 977 KB 977 KB +0.0%
qwen3_5_vl (hybrid-qwen-vl) num_nodes 408 408 +0.0%
t5 (seq2seq) model_size_bytes 836 KB 836 KB +0.0%
t5 (seq2seq) num_nodes 166 166 +0.0%
whisper (speech-to-text) model_size_bytes 1008 KB 1008 KB +0.0%
whisper (speech-to-text) num_nodes 128 128 +0.0%

No performance regressions.

@codecov

codecov Bot commented May 5, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby
justinchuby force-pushed the gemma4-genai-example branch from 54054d9 to 92dfce8 Compare May 5, 2026 19:42
@justinchuby

Copy link
Copy Markdown
Member Author

@copilot please fix lint

Copilot AI commented May 5, 2026

Copy link
Copy Markdown
Contributor

@copilot please fix lint

Fixed in 36d0f2a — reformatted the list literal in build_and_export and the ternary expression for decode_tps per ruff-format.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby
justinchuby force-pushed the gemma4-genai-example branch from 36d0f2a to d01c8a1 Compare May 5, 2026 21:29
Gemma4 instruction-tuned models require chat-formatted prompts.
Raw text prompts produce degenerate repetitive output. Add
format_chat_prompt() that wraps user messages in the Gemma chat
template (<start_of_turn>user/model format).

Verified: with chat template, ORT GenAI output is token-for-token
identical to HuggingFace PyTorch (9/9 and 64/64 tokens matched).

Signed-off-by: Justin Chu <justinchu@microsoft.com>
Gemma4 e2b/e4b are multimodal — support text + image + audio via
GenAI's MultiModalProcessor. When --image or --audio is provided,
the script uses the VLM generation path with set_inputs().

- Text-only: uses chat template, appends tokens
- Multimodal: uses HF processor for template with media placeholders,
  loads media via og.Images/og.Audios, processes with MultiModalProcessor
- Auto-detects VLM vs text-only build based on --image/--audio flags

Signed-off-by: Justin Chu <justinchu@microsoft.com>
Gemma4 tokenizer has add_bos_token=False, so <bos> must be included
explicitly in the chat template. Without it, generation produces
garbled output because the model expects BOS as the first token.

The HF apply_chat_template() includes <bos> automatically, but our
manual format_chat_prompt() was missing it.

Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby
justinchuby enabled auto-merge (squash) May 6, 2026 03:54
@justinchuby
justinchuby disabled auto-merge May 6, 2026 03:55
GenAI's tokenizer reads the chat template from tokenizer_config.json
and applies it correctly (including <bos>). Remove the manual
format_chat_prompt() function and use:

  messages = json.dumps([{'role': 'user', 'content': prompt}])
  chat_prompt = tokenizer.apply_chat_template(messages)

This is more robust — it uses the model's own template definition
rather than a hardcoded format string that can drift.

Signed-off-by: Justin Chu <justinchu@microsoft.com>
@justinchuby
justinchuby merged commit 6d3b3e2 into main May 6, 2026
21 of 22 checks passed
@justinchuby
justinchuby deleted the gemma4-genai-example branch May 6, 2026 14:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants