Gemma 4 E2B: GGUF INT4 direct conversion + QNN HTP lowering (#407)

## Summary

Enables converting the **unsloth/gemma-4-E2B-it GGUF (Q4_K_M)** directly
into a working INT4 ONNX model that targets the **QNN EP / Hexagon HTP**
— reusing the GGUF's existing INT4 weights instead of running expensive
Olive quantization. Pairs with olive-recipes#432 (adds a GGUF-direct
`config_gguf.json`).

## What's here

**1. Gemma 4 GGUF INT4 conversion fixes** (make `mobius build-gguf
--keep-quantized` work for E2B):
- Collapse per-layer `feed_forward_length` array → scalar
`intermediate_size` + `use_double_wide_mlp`.
- Map the top-level per-layer-input tensors + strip the
`layer_scalar.weight` artefact.
- Thread the INT4 `MatMulNBits` linear class through Gemma4
(`--keep-quantized` was silently producing a **float** model) + tie the
float LM head.
- Fix KV-shared standard-Attention shape inference (opset-24 `Attention`
present outputs are un-inferred → `o_proj` collapsed → model
unloadable).
- Fix the Gemma GGUF activation default (`silu` → `gelu_pytorch_tanh`;
the wrong default produced garbage output).

**2. QNN HTP lowering** (two EP-capability-gated transforms; both
default-on, off for `qnn`):
- `supports_attention=False` → `DecomposeAttentionPass` rewrites the
fused opset-24 `Attention` op into SDPA primitives
(Reshape/Transpose/MatMul/Softmax/Add, Tile for GQA). An `onnx_ir`
InPlacePass so it can rewire the 3 outputs (present-KV are graph outputs
on some layers, dead on shared-KV layers — mixed arity the pattern
rewriter can't express).
- `supports_matmul_nbits=False` → `com.microsoft::MatMulNBits` lowered
to blocked-uint4 `DequantizeLinear` + `MatMul` (QDQ), implemented as a
**model-local ir.Function** expanded by the existing InlinePass
mechanism (like SkipLayerNormalization / PackedMultiHeadAttention).
Supported EPs keep the native contrib op + kernel.

## Verification

- Dequantized GGUF weights match the gated `google/gemma-4-E2B-it`
checkpoint (corr ≥ 0.98).
- Both lowerings are numerically identical to the originals via ORT
(maxdiff 0.0 for prefill/decode/GQA/mask; ~5e-7 with softcap due to fp
rounding).
- The INT4 model generates correct text on ORT ("capital of France →
Paris", "capital of Japan → Tokyo", valid haiku).
- A full `qnn` build lowers to **0 Attention / 35 Softmax** and **0
MatMulNBits / 205 DequantizeLinear**, all 30 present-KV graph outputs
preserved; CPU/CUDA builds unchanged (still emit GQA + MatMulNBits).
- Empirically confirmed on a Snapdragon X box (ORT 1.27, QNN SDK
2.48.40): HTP runs the primitive ops
(MatMul/Softmax/RMSNormalization/Gelu) but **not** fused
Attention/MatMulNBits — exactly what this lowering removes. (Plugin QNN
EP is engaged via `add_provider_for_devices`, and HTP requires fully
static shapes.)

## Remaining / follow-up

- Unit tests for `DecomposeAttentionPass` and the `MatMulNBits` function
body (in progress).
- End-to-end HTP EPContext compile validation (olive-recipes
`config_gguf.json` path) not yet hardware-validated.

🤖 Generated with GitHub Copilot CLI

---------

Signed-off-by: Justin Chu <justinchu@microsoft.com>
Signed-off-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

Author: 3 people

src/mobius/__main__.py

@@ -428,22 +428,22 @@ def _cmd_build_gguf(args: argparse.Namespace) -> None:
     output_dir = args.output or os.path.splitext(gguf_path)[0] + "_onnx"
     os.makedirs(output_dir, exist_ok=True)
 
-    if mmproj_path is not None:
-        print(f"Multimodal mode: fusing vision/audio encoder from mmproj {mmproj_path}...")
-        pkg = build_from_gguf(
-            gguf_path,
-            mmproj=mmproj_path,
-            dtype=args.dtype,
-            execution_provider=args.execution_provider,
-            keep_quantized=args.keep_quantized,
-        )
-    else:
-        pkg = build_from_gguf(
-            gguf_path,
-            dtype=args.dtype,
-            keep_quantized=args.keep_quantized,
-            execution_provider=args.execution_provider,
-        )
+    if args.max_seq_len is not None and not args.static_cache:
+        raise SystemExit("Error: --max-seq-len can only be used with --static-cache.")
+    if args.max_seq_len is not None and args.max_seq_len <= 0:
+        raise SystemExit("Error: --max-seq-len must be a positive integer.")
+    if mmproj_path is not None and args.static_cache:
+        raise SystemExit("Error: --static-cache cannot be used with --mmproj.")
+
+    pkg = build_from_gguf(
+        gguf_path,
+        mmproj=mmproj_path,
+        dtype=args.dtype,
+        keep_quantized=args.keep_quantized,
+        execution_provider=args.execution_provider,
+        static_cache=args.static_cache,
+        max_seq_len=args.max_seq_len,
+    )
 
     pkg.save(
         output_dir,
@@ -770,6 +770,26 @@ def main(argv: list[str] | None = None) -> None:
             "Either way the quantized model runs directly in the target runtime."
         ),
     )
+    gguf_parser.add_argument(
+        "--static-cache",
+        action="store_true",
+        help=(
+            "Use a static KV cache (pre-allocated fixed-width buffers written "
+            "in place via TensorScatter) instead of the dynamic concat-grow "
+            "cache. Produces a fully static-shaped graph as required by "
+            "fixed-shape runtimes such as the QNN HTP backend."
+        ),
+    )
+    gguf_parser.add_argument(
+        "--max-seq-len",
+        type=int,
+        default=None,
+        metavar="N",
+        help=(
+            "Maximum sequence length for static cache buffers. Only used with "
+            "--static-cache. Defaults to max_position_embeddings from config."
+        ),
+    )
     gguf_parser.set_defaults(func=_cmd_build_gguf)
 
     # --- list ---

src/mobius/_execution_providers.py

@@ -71,6 +71,33 @@ class EpCapabilities:
             (query/key norm over the head dimension) to rank-3 and back via
             HtpRank4RMSNorm.  ``True`` leaves it unchanged.  Set ``False``
             only for the QNN HTP, which miscomputes rank-4 RMSNormalization.
+        supports_attention: ``False`` decomposes the fused opset-24
+            ``Attention`` op into scaled-dot-product primitives
+            (Reshape/Transpose/MatMul/Softmax/Add, plus Expand for GQA) via
+            DecomposeAttention.  ``True`` leaves the fused op unchanged.  Set
+            ``False`` only for runtimes without an ``Attention`` kernel (QNN
+            HTP), where the fused op would otherwise be forced onto CPU.
+        supports_rotary_embedding: ``False`` decomposes the opset-24
+            ``RotaryEmbedding`` op into rotate-half primitives (Reshape/Slice/
+            Mul/Sub/Add/Concat) via DecomposeRotaryEmbedding.  ``True`` leaves
+            the fused op unchanged.  Set ``False`` only for runtimes without a
+            ``RotaryEmbedding`` kernel (QNN HTP), where it is forced onto CPU.
+        supports_tensor_scatter: ``False`` rewrites the static-cache
+            ``TensorScatter`` in-place KV write into ``ScatterND`` (batch=1) via
+            TensorScatterToScatterND.  ``True`` leaves ``TensorScatter``
+            unchanged.  Set ``False`` only for runtimes without a
+            ``TensorScatter`` kernel (QNN HTP), where it is forced onto CPU.
+        supports_range: ``False`` replaces the ONNX ``Range`` op in
+            :func:`~mobius.components.create_static_cache_attention_bias` with a
+            precomputed ``Constant(arange)`` + ``Slice``.  ``True`` (the
+            default) leaves ``Range`` unchanged.  Set ``False`` only when
+            static cache is used on runtimes that lack a ``Range`` kernel (QNN
+            HTP), where the ``Range`` node would otherwise be forced onto CPU.
+        supports_matmul_nbits: ``False`` converts ``com.microsoft::MatMulNBits``
+            (blockwise-INT4 weight) into a standard ``DequantizeLinear`` +
+            ``MatMul`` (QDQ) pair via MatMulNBitsToQDQ.  ``True`` leaves the
+            contrib op unchanged.  Set ``False`` only for runtimes that lack a
+            ``MatMulNBits`` kernel but can consume QDQ weights (QNN HTP).
         default_int4_accuracy_level: Default accuracy level for INT4
             quantization (0 = highest accuracy, 4 = fastest).
         provider_options: Default ORT GenAI provider options dict for this EP.
@@ -121,6 +148,11 @@ class EpCapabilities:
     supports_fused_moe: bool = True
     supports_packed_multi_head_attention: bool = False
     supports_rank4_rmsnorm: bool = True
+    supports_attention: bool = True
+    supports_matmul_nbits: bool = True
+    supports_rotary_embedding: bool = True
+    supports_tensor_scatter: bool = True
+    supports_range: bool = True
     default_int4_accuracy_level: int = 0
     provider_options: dict[str, str] = dataclasses.field(default_factory=dict)
     enable_graph_capture: bool = False
@@ -344,6 +376,11 @@ def _register_builtins() -> None:
             },
             supports_past_present_share_buffer=False,  # standard-Attention KV concat
             supports_rank4_rmsnorm=False,  # HTP miscomputes rank-4 RMSNorm (q/k norm)
+            supports_attention=False,  # no HTP Attention kernel — decompose to SDPA
+            supports_matmul_nbits=False,  # no HTP MatMulNBits kernel — convert to QDQ
+            supports_rotary_embedding=False,  # no HTP RotaryEmbedding — rotate-half
+            supports_tensor_scatter=False,  # no HTP TensorScatter — ScatterND (batch=1)
+            supports_range=False,  # no HTP Range — Constant(arange)+Slice in static bias
         ),
         # onnx-standard: ONNX-only runtime — emits zero custom-domain ops.
         # All com.microsoft ops (SkipLayerNorm, PackedMHA) are expanded via

src/mobius/_optimizations.py

@@ -61,6 +61,8 @@
 )
 from mobius.functions import register_function_bodies
 from mobius.rewrite_rules import (
+    decompose_attention_pass,
+    decompose_rope_rules,
     gelu_fusion_rules,
     group_query_attention_rules,
     htp_rank4_rmsnorm_rules,
@@ -69,6 +71,7 @@
     skip_layer_norm_rules,
     skip_norm_rules,
     static_empty_kv_rules,
+    tensor_scatter_to_scatternd_rules,
     unpack_qkv_rules,
 )
 
@@ -318,6 +321,29 @@ def _get_optimization_passes(
     if not caps.supports_rank4_rmsnorm:
         lower.append(("HtpRank4RMSNorm", list(htp_rank4_rmsnorm_rules())))
 
+    # Decompose the opset-24 RotaryEmbedding op into rotate-half primitives for
+    # EPs without a RotaryEmbedding kernel (QNN HTP), where the op falls to CPU.
+    if not caps.supports_rotary_embedding:
+        lower.append(("DecomposeRotaryEmbedding", list(decompose_rope_rules())))
+
+    # Rewrite the static-cache TensorScatter KV write into ScatterND for EPs
+    # without a TensorScatter kernel (QNN HTP), where it falls to CPU.
+    if not caps.supports_tensor_scatter:
+        lower.append(("TensorScatterToScatterND", list(tensor_scatter_to_scatternd_rules())))
+
+    # Decompose the fused opset-24 Attention op into SDPA primitives for EPs
+    # without an Attention kernel (QNN HTP), where the fused op falls to CPU.
+    # An IR pass (not a rewrite rule) so it can rewire the op's 3 outputs, whose
+    # KV-cache present tensors are graph outputs on some layers and dead on
+    # shared-KV layers. Runs after the batched rewrite rules (lower_ir_passes).
+    if not caps.supports_attention:
+        lower.append(("DecomposeAttention", decompose_attention_pass()))
+
+    # NOTE: MatMulNBits → QDQ lowering is handled by the InlinePass mechanism
+    # (a registered com.microsoft::MatMulNBits function body), not a rewrite
+    # rule — see optimize_model()'s _should_inline and mobius.functions. It is
+    # gated by ``supports_matmul_nbits`` there.
+
     # --- Graph-capture rewrite ---
     # Supports graph capture for shared-KV layer models on WebGPU EP
     # (currently Gemma4). Pattern-based: only fires on models that emit the
@@ -419,6 +445,12 @@ def _should_inline(func: ir.Function) -> bool:
             return not caps.supports_skip_layer_norm
         if func.domain == "com.microsoft" and func.name == "PackedMultiHeadAttention":
             return not caps.supports_packed_multi_head_attention
+        # MatMulNBits (blockwise-INT4) → QDQ (DequantizeLinear + MatMul) for EPs
+        # without a MatMulNBits kernel (QNN HTP). Supported EPs (CPU/CUDA/…) keep
+        # the compact contrib op and its native kernel; the function body stays
+        # registered but uninlined (kernels take precedence over local functions).
+        if func.domain == "com.microsoft" and func.name == "MatMulNBits":
+            return not caps.supports_matmul_nbits
         return False
 
     inline_pass = common_passes.InlinePass(criteria=_should_inline)

src/mobius/components/_common.py

@@ -335,23 +335,45 @@ def create_static_cache_attention_bias(
         Additive bias of shape ``(B, 1, S_q, max_seq_len)``: ``0`` where the
         query may attend the slot, ``dtype.min`` where masked.
     """
-    zero_scalar = op.Constant(value_int=0)
-    one_scalar = op.Constant(value_int=1)
-
-    # Q absolute positions: write_indices[b] + arange(S_q) -> (B, S_q, 1).
-    seq_scalar = op.Squeeze(seq_len, op.Constant(value_ints=[0]))
-    q_offsets = op.Range(zero_scalar, seq_scalar, one_scalar)  # (S_q,)
+    # KV positions: dense cache slot ids 0 .. max_seq_len - 1 -> (max_seq_len,).
+    # For EPs that support Range (CPU, CUDA, etc.) emit a single dynamic Range
+    # node.  For EPs without a Range kernel (QNN HTP, supports_range=False),
+    # emit a precomputed Constant — max_seq_len is always a concrete int in
+    # static-cache mode — and derive q_offsets via Slice, which both Constant
+    # and Slice run on HTP without CPU fallback.
+    from mobius._build_context import ep_capabilities
+
+    if ep_capabilities().supports_range:
+        zero_scalar = op.Constant(value_int=0)
+        one_scalar = op.Constant(value_int=1)
+        kv_slots = op.Range(
+            zero_scalar, op.Constant(value_int=max_seq_len), one_scalar
+        )  # (max_seq_len,)
+
+        # Q absolute positions: write_indices[b] + arange(S_q) -> (B, S_q, 1).
+        seq_scalar = op.Squeeze(seq_len, op.Constant(value_ints=[0]))
+        q_offsets = op.Range(zero_scalar, seq_scalar, one_scalar)  # (S_q,)
+    else:
+        # QNN HTP has no Range kernel — use a precomputed constant for the
+        # full kv_slots array and Slice it to S_q for q_offsets; both ops
+        # run on HTP.
+        kv_slots = op.Constant(
+            value=ir.tensor(np.arange(max_seq_len, dtype=np.int64))
+        )  # (max_seq_len,)
+
+        # arange(S_q) = kv_slots[:S_q] via Slice.
+        q_offsets = op.Slice(
+            kv_slots,
+            op.Constant(value_ints=[0]),  # start
+            seq_len,  # end == [S_q]
+            op.Constant(value_ints=[0]),  # axis
+        )  # (S_q,)
     q_abs_2d = op.Add(
         op.Unsqueeze(write_indices, [1]),  # (B, 1)
         op.Unsqueeze(q_offsets, [0]),  # (1, S_q)
     )  # (B, S_q)
     q_abs = op.Unsqueeze(q_abs_2d, [2])  # (B, S_q, 1)
 
-    # KV positions: dense cache slot ids 0 .. max_seq_len - 1 -> (max_seq_len,).
-    kv_slots = op.Range(
-        zero_scalar, op.Constant(value_int=max_seq_len), one_scalar
-    )  # (max_seq_len,)
-
     # Causal: a query at absolute position q may attend slot k iff q >= k.
     # Broadcasting (B, S_q, 1) >= (max_seq_len,) -> (B, S_q, max_seq_len).
     full_mask = op.GreaterOrEqual(q_abs, kv_slots)

src/mobius/functions/__init__.py

@@ -33,6 +33,9 @@
 from mobius.functions.linear_attention import (
     linear_attention,
 )
+from mobius.functions.matmul_nbits import (
+    matmul_nbits,
+)
 from mobius.functions.packed_multi_head_attention import (
     packed_multi_head_attention,
 )
@@ -53,6 +56,7 @@
 # (kernel_size, channels, num_heads are baked into the function body) and are
 # registered per-model by ``tasks._base._register_linear_attention_functions``.
 _FUNCTION_BUILDERS: dict[ir.OperatorIdentifier, Callable[[], ir.Function]] = {
+    (_DOMAIN, "MatMulNBits", ""): matmul_nbits,
     (_DOMAIN, "PackedMultiHeadAttention", ""): packed_multi_head_attention,
     (_DOMAIN, "SkipLayerNormalization", ""): skip_layer_normalization,
     (_DOMAIN, "SkipSimplifiedLayerNormalization", ""): skip_simplified_layer_normalization,
@@ -100,6 +104,7 @@ def register_function_bodies(model: ir.Model) -> None:
     "causal_conv_nd_with_state",
     "get_function",
     "linear_attention",
+    "matmul_nbits",
     "packed_multi_head_attention",
     "register_function_bodies",
     "skip_layer_normalization",