Add gemma4 OpenVINO NPU support - #746
Open
xiaoyu-work wants to merge 3 commits into
Open
xiaoyu-work wants to merge 3 commits into
xiaoyu-work wants to merge 3 commits into
Conversation
Emit an OpenVINO-compatible rank-4 attention graph, preserve opset 23 Attention semantics, and use canonical Gather extraction for per-layer inputs so NPUW can compile repeated decoder blocks. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d7d02d41-4f2e-4bdd-ae7d-0d22fbc6d3bd
Drop checkpoint K/V projection and norm tensors for layers that borrow a shared KV cache, including packed quantized sidecars that have no target module. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d7d02d41-4f2e-4bdd-ae7d-0d22fbc6d3bd
Remove one-off requires_* capability flags and branch directly on the OpenVINO execution provider where the Gemma4 and opset workarounds are applied. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d7d02d41-4f2e-4bdd-ae7d-0d22fbc6d3bd
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Critical OpenVINO RoPE and dtype-boundary issues remain, and runtime parity coverage is missing.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 3
Open (4)
What changed in this PR
Adds OpenVINO/NPU support for Gemma4 through rank-4 attention graphs, float32 boundaries, shared-KV preprocessing, and opset lowering.
Changes:
- Adds OpenVINO decoder and embedding I/O handling.
- Implements OpenVINO attention with RoPE and KV broadcasting.
- Adds graph, preprocessing, and opset-lowering tests.
| File | Summary |
|---|---|
src/mobius/tasks/_gemma4.py |
Adds OpenVINO I/O handling; text-only and multimodal boundaries remain incomplete. |
src/mobius/models/gemma4.py |
Adds OpenVINO attention and shared-KV handling; RoPE interleave compatibility and runtime parity require changes. |
src/mobius/models/gemma4_test.py |
Adds topology and preprocessing tests; runtime prefill/decode coverage is still needed. |
src/mobius/_builder.py |
Adds OpenVINO opset lowering. |
src/mobius/_builder_test.py |
Adds opset-lowering test coverage. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Comment on lines
+1164
to
+1167
| first, second = op.Split(states, [head_dim // 2, head_dim // 2], axis=-1, _outputs=2) | ||
| rotated = op.Concat(op.Neg(second), first, axis=-1) | ||
| left = op.Mul(states, cos) | ||
| right = op.Mul(rotated, sin) |
| op = builder.op | ||
| caps = ep_capabilities() | ||
| is_openvino = caps.name == "openvino" | ||
| decoder_io_dtype = ir.DataType.FLOAT if is_openvino else config.dtype |
Comment on lines
+827
to
+835
| cast_outputs = ep_capabilities().name == "openvino" | ||
|
|
||
| def component_output(value: ir.Value) -> ir.Value: | ||
| if not cast_outputs or value.dtype == ir.DataType.FLOAT: | ||
| return value | ||
| output = op.Cast(value, to=ir.DataType.FLOAT) | ||
| output.type = ir.DataType.FLOAT | ||
| output.shape = value.shape | ||
| return output |
Comment on lines
+477
to
+482
| assert counts["Attention"] == config.num_hidden_layers | ||
| assert counts["RotaryEmbedding"] == 0 | ||
| assert counts["Softmax"] == 0 | ||
| assert counts["CumSum"] == 1 | ||
| assert inputs_embeds.dtype == ir.DataType.FLOAT | ||
| assert logits.dtype == ir.DataType.FLOAT |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add gemma4 OpenVINO NPU support