Pack float and INT4 QKV projections with deferred weight binding

docs/ep_quickstart.md

@@ -100,7 +100,9 @@ print(sorted(ep_registry))
 # Inspect an EP's capabilities
 caps = get_ep("cuda")
 print(caps.gqa_dtypes)          # frozenset({FLOAT16, BFLOAT16})
-print(caps.qkv_pack_dtypes)     # frozenset({FLOAT, FLOAT16, BFLOAT16})
+print(caps.qkv_pack_dtypes)     # frozenset({FLOAT, FLOAT16, BFLOAT16}); legacy GQA-input packing
+print(caps.qkv_projection_pack_dtypes)  # frozenset(); experimental float projection packing is off
+print(caps.matmul_nbits_qkv_pack_dtypes)  # frozenset(); experimental INT4 packing is off
 print(caps.supports_fused_rope) # True
 print(caps.enable_graph_capture) # True
 print(caps.provider_options)     # {'enable_skip_layer_norm_strict_mode': '1'}
@@ -132,6 +134,15 @@ pkg = mobius.build("meta-llama/Llama-3.2-1B",
 Unrecognised EP names raise `ValueError` during build-time validation, before
 graph construction or optimization starts — safe-fail by default.
 
+Experimental float `PackQKVProjections` and INT4 projection packing require
+independent `qkv_projection_pack_dtypes` and `matmul_nbits_qkv_pack_dtypes`
+capability overrides, respectively; both are **default-off on every built-in EP**.
+The existing float GQA-input `PackQKV` stage retains its `qkv_pack_dtypes` gate.
+Neither INT4 nor CUDA FP16/BF16 float projection-packing numerical accuracy is
+qualified by this branch. Read the
+[numerical limitations and opt-in example](execution_providers.md#projection-packing-and-numerical-limitations)
+before enabling it; activation/component precision flags do not opt in.
+
 ---
 
 ## 5. Common EP configurations at a glance
@@ -141,7 +152,7 @@ graph construction or optimization starts — safe-fail by default.
 | Portable ONNX (maximum compatibility) | `"default"` | any | No EP-specific vendor fusions (e.g. no GQA/PackQKV) |
 | Strict standard ONNX (no custom ops at all) | `"onnx-standard"` | any | All `com.microsoft` ops expanded via InlinePass; safe for non-ORT runtimes |
 | ORT CPU inference | `"cpu"` | `"f32"` | GQA fusion for FP32 |
-| NVIDIA GPU | `"cuda"` | `"f16"` or `"bf16"` | GQA + SkipNorm + PackQKV |
+| NVIDIA GPU | `"cuda"` | `"f16"` or `"bf16"` | GQA + SkipNorm + existing float PackQKV plus new default float projection packing (also for masked decoders); experimental INT4 projection packing is off |
 | Windows GPU (DirectX) | `"dml"` | `"f16"` | RoPE lowered separately |
 | Browser / WebAssembly | `"webgpu"` | `"f16"` or `"f32"` | Shape ops replaced |
 | NVIDIA TensorRT-RTX | `"trt-rtx"` | `"f16"` or `"bf16"` | SkipLayerNorm expanded |

docs/execution_providers.md

@@ -100,6 +100,73 @@ ops to fused equivalents:
 
 These are applied during **Stage 2: Fusion** of the optimization pipeline.
 
+#### Projection packing and numerical limitations
+
+`PackQKVProjections` combines compatible projections into one GEMM and splits
+the result **before normalization, RoPE, attention, or caches**, without changing
+their contracts. It supports Q/K/V and Q/K with raw K reused for V.
+
+| Stage | Roles | Enablement |
+|---|---|---|
+| Existing float `PackQKV` GQA-input fusion | `decoder` | Unchanged |
+| Additional float `PackQKVProjections` | `decoder`, `masked-decoder` | **Experimental; disabled on every built-in EP**; explicit `qkv_projection_pack_dtypes` opt-in |
+| INT4 projection packing | `decoder`, `masked-decoder` | **Experimental; disabled on every built-in EP**, including CUDA FP16/BF16 |
+
+Experimental float and INT4 projection packing use independent process-local
+capabilities. Stock `--ep cuda` enables neither new stage; the existing
+`qkv_pack_dtypes`-gated GQA-input packing is unchanged. Set only the capabilities
+you intend to evaluate (the example opts into both):
+
+```python
+import dataclasses
+
+import onnx_ir as ir
+from mobius import get_ep, register_ep
+
+register_ep(
+    dataclasses.replace(
+        get_ep("cuda"),
+        qkv_projection_pack_dtypes=frozenset({ir.DataType.FLOAT16}),
+        matmul_nbits_qkv_pack_dtypes=frozenset({ir.DataType.FLOAT16}),
+    ),
+    overwrite=True,
+)
+```
+
+Both opt-ins are independent of legacy float packing, activation/storage dtype,
+component precision, and FP8 KV caches. Neither byte-exact payload tests nor
+CPU FLOAT execution tests qualify CUDA FP16/BF16 numerical parity. Packing never
+dequantizes, requantizes, or promotes weights.
+`_qkv_projection_packing.py` implements both float and INT4.
+
+Packing requires static, non-overridable, compatible parameters and opset >= 13;
+ambiguous or overlapping groups and weights with outside consumers are skipped.
+INT4 additionally requires a shared rank-3 activation, matching attributes except
+output width, canonical UINT8 blocked weights, same-dtype scales, and consistently
+absent or packed UINT8 zero points. Flat/row-major sidecars are supported;
+block size must be a power of two >= 16 and `accuracy_level` omitted or zero.
+Mixed bit widths, extra operands, and caller-owned function bodies are skipped.
+
+Parameter Concats retain checkpoint names until binding, then initializer
+Concats and matrix Transposes fold lazily without changing dtype. Missing leaves
+are checked even with `fold_constants=False`. Generic Reshape/Transpose layout
+folding is not part of this pipeline; those transformations remain explicit.
+Overridable inputs remain untouched.
+
+**The INT4 numerical gate remains failed and unwaived:** large-dimension CUDA
+prefill/cached-decode logits exceeded `rtol=atol=1e-2`. Changing GEMM width can
+change floating-point results despite byte-exact weights. MMLU/generation results
+do not clear that gate; float projection packing has no separate numerical
+qualification here. No production accuracy or speed improvement is claimed.
+
+For `ORT_FPA_INTB_GEMM=1`, keep exporter or ORT constant folding enabled so
+parameter Concats do not block runtime prepacking; behavior is runtime-dependent.
+Packed float weights retain adapter-restoration metadata and lazy backing-file
+ownership after pruning. `ModelPackage.save()` rejects known same-path, hardlink,
+or symlink aliases to backing files across model, shard, policy, adapter, metadata,
+and MTP destinations. This is **not full TOCTOU protection**; save to a separate
+output directory.
+
 ### Tier 3 — EP-specific constraints (required, correctness)
 
 Some EPs cannot execute certain ops. These are handled by **lowering rules**
@@ -126,7 +193,8 @@ Every EP is fully described by a single `EpCapabilities` entry in the
 class EpCapabilities:
     name: str
     gqa_dtypes: frozenset[ir.DataType]               # dtypes where GQA fusion fires
-    qkv_pack_dtypes: frozenset[ir.DataType]          # dtypes where PackQKV fusion fires
+    qkv_pack_dtypes: frozenset[ir.DataType]          # legacy GQA-input packing only
+    qkv_projection_pack_dtypes: frozenset[ir.DataType]  # experimental float; default empty
     supports_fused_rope: bool = True                 # False → SeparateRoPE + UnpackQKV
     supports_skip_layer_norm: bool = True            # False → InlinePass expansion
     supports_fused_matmul: bool = True               # False → Transpose + MatMul via InlinePass

docs/gemma4-runtime.md

@@ -0,0 +1,186 @@
+# Gemma4 text and multimodal runtime qualification
+
+This runtime work is draft, not production or equivalence approved. It does not
+qualify optimized text-only performance, ASR, or experimental fixed/paged-cache
+execution.
+
+## HuggingFace alignment and input contract
+
+- Global decoder layers use causal attention only; image-block IDs must not
+  unmask future keys in global layers.
+- Local decoder layers use
+  `valid_key AND (causal OR same_image) AND (key > query - window)`.
+  Same-image future keys can be visible, but an image block cannot bypass the
+  strict left window boundary. Distances use physical query/KV slots, including
+  past offsets; padding holes do not compress the window. Only current-chunk
+  image-block IDs are available;
+  these IDs are not persisted alongside the KV cache.
+- The separate vision encoder's `pixel_values` input deliberately uses ONNX
+  `FLOAT`, including with BF16/FP16 projection weights. Supply FP32 processor
+  patches in `[0, 1]`, not BF16 patches or pre-normalized `[-1, 1]` values.
+  Normalization computes `2 * (pixels - 0.5)` in FP32 before casting to the
+  projection's working dtype. Existing consumers must update their input dtype;
+  `examples/gemma4_multimodal.py` demonstrates the new ABI. This does not change
+  the encoder-free `gemma4_unified` vision ABI.
+- Checkpoint `router.scale` is preserved. The router first multiplies the
+  normalized working-dtype activation by the learned scale, then applies the
+  separate `hidden_size ** -0.5` scalar with FP32 opmath and rounds back to the
+  working dtype. Folding that scalar into a low-precision learned scale changes
+  rounding and is not equivalent.
+- Checkpoint-enabled vision standardization retains `std_bias`/`std_scale`
+  directly on the vision submodel. Pooling coefficients and feature accumulation
+  are FP32, including the `1/9` coefficients for a 3x3 kernel. The pooled mean
+  then rounds to working dtype before FP32 scaling/standardization and the
+  final cast back for projection.
+
+## Native runtime requirements and limits
+
+Ordinary `execution_provider="cuda"` does **not** assert custom GEGLU support.
+The default-off `gemma4_custom_geglu` flag explicitly opts into the experimental
+custom CUDA ABI: set `MOBIUS_GEMMA4_CUSTOM_GEGLU=1` before importing Mobius, or
+use `override_flags(gemma4_custom_geglu=True)` around construction, export and
+weight preprocessing. Environment values are captured when the module-level
+flags instance is created on first import; changing the environment afterward
+does not update that instance. Without opt-in, float tanh-GELU experts use the portable
+exact-activation fallback; quantized GEGLU experts fail at construction with an
+explicit requirement message. Opt-in does not enable native GEGLU on non-CUDA
+providers. Tanh-GELU must not be replaced by a sigmoid-based SwiGLU approximation.
+
+The portable fallback evaluates experts densely, but selects out unchosen expert
+outputs before weighting. This prevents an overflowing unselected expert from
+contaminating the result through `0 * Inf` or `0 * NaN`; nonfinite selected
+experts still propagate. This safeguard is not full-checkpoint numerical or
+performance qualification.
+
+Each emitted native `MoE`/`QMoE` operator persists
+`mobius.gemma4.required_runtime_feature="custom-cuda-geglu-v1"` in its ONNX
+metadata. This declares the experimental CUDA `activation_type="geglu"`,
+`swiglu_fusion=1` requirement; it is not a claim of released ORT support or a
+claim that the exporting machine used a particular runtime binary.
+
+The frozen pre-port experiments used an **experimental modified ORT 1.32 build**:
+source `942356fce6bdc78c016e4b902e49492fdf695453` with seven dirty paths; binary
+SHA-256 `3f6b7ff482f8fbd4d3c639e807ed2ec35fa240c88c205dbee7822d7ec497e988`.
+This is not an unmodified upstream release. Its local patchset and full
+experiment records are not included in this repository; these identity fields
+do not provide a reproducible runtime build. Released ORT support is unverified.
+
+Dynamic vision-overlay decoders declare the arbitrary-attention-mask contract.
+Both local and global layers retain ONNX `Attention`: reducing a mask to a valid
+token count would discard internal padding holes even in causal global layers.
+This preserves the full additive bias through CUDA optimization. The default
+VLM qualification path uses dynamic KV cache, graph capture off, and past/present
+buffer sharing off.
+
+Text-only export uses the existing `text_only=True` builder option (CLI
+`--features text-only`) to remove the vision overlay. Its causal/windowed
+attention can use the existing GQA path when supported by the execution provider.
+GQA sequence lengths describe prefix validity, not arbitrary interior mask holes;
+use a standard-ONNX attention path when those holes must be preserved.
+That eligibility is separate from projection packing, custom GEGLU opt-in, and
+runtime KV-sharing/graph-capture settings. No experimental static-cache or
+paged-attention changes are part of this modeling work.
+
+## Recorded pre-port experiments, not validation of this branch
+
+The corrected real-weight native BF16 and INT4 VLM runs each scored
+**166/200**; HF BF16 scored **167/200** and HF INT4 **162/200**. These results use
+a fixed quantization-selection AI2D subset, not an independent full-AI2D
+evaluation, and must not be described as such.
+
+Those runs used frozen corrected exports built from the original dirty source
+worktree, with source hashes recorded in the originating experiment records.
+They were **not** exports of this isolated origin/main-based branch. Newer main
+has different leaf-layout/configuration features; these recorded scores (and
+any separately recorded pre-port performance results) do not validate this
+branch. The frozen exports also predate this branch's corrected FP32 pooling
+and physical-slot window fixes. Publication requires the original source hashes,
+checkpoint/quantization configuration, export identity, and runtime version to
+remain attached to the experiment provenance.
+
+Strict full-vocabulary logit parity **failed 9/9**, and KV parity **failed
+540/540**, at `atol=0.01, rtol=0.01`. Task-level scores do not establish numerical
+equivalence or production readiness. The isolated
+patch's tiny-graph and weight-layout tests are separate evidence. Full-checkpoint
+re-export and revalidation on the ported branch remain **unverified**. No model
+rebuild, download, GPU run, or performance measurement is implied by this note.
+Lead review is required before publication.
+
+## Prior combined-branch evidence
+
+Offline, CPU-only validation against base
+`d69efe03a9f4c1ada744bc31dd01077a710b3c39`:
+
+- Related runtime, package, expert-layout, adapter, and packing tests:
+  **656 passed, 2 skipped**. These include `models/gemma4_test.py`,
+  `components/_common_test.py`, `_model_package_test.py`,
+  `functions/matmul_nbits_test.py`, and
+  `rewrite_rules/_qkv_projection_packing_test.py` under `src/mobius/`.
+- GQA rewrite tests: **32 passed**. Four remote-config cases were excluded from
+  that offline invocation after reproducing their failures on the exact clean
+  base; their tests were not changed.
+- `tests/build_graph/core_test.py -k gemma4`: **10 passed**.
+  `tests/integration/gemma4_test.py::test_gemma4_bidirectional_mask_parity`:
+  **11 passed** against the actual pinned HF mask entry points.
+
+Changed-file lint, formatting, compilation, and whitespace checks passed.
+Strict mypy with resolved editable ONNX Script imports has zero new normalized
+diagnostic instances relative to base, but neither revision typechecks cleanly.
+These checks do not replace the unresolved full-checkpoint numerical gates.
+The modeling/pass split and subsequent review fixes require their own validation;
+the historical counts above are not results for the reorganized revision.
+
+## Packed-checkpoint contract
+
+Native Gemma4 experts require an explicit decoder quantization config:
+`quant_method="olive"`, integer-affine `bits=4`, a power-of-two
+`group_size >= 16`, `sym`, and canonical `quantize_moe=true`.
+Metadata also accepts the legacy `moe` alias; `quantize_moe` takes precedence
+when both are supplied. Expert quantization opt-in and the custom-runtime flag
+are separate requirements.
+Compute dtype must be FP16/BF16 and expert activation must be tanh-GELU.
+Both expert input dimensions must be divisible by the group size. Preserve the
+authoritative HF source-module names and any component/per-projection override
+or exclusion plan; unsupported expert layouts fail closed, while ordinary
+projection plans remain handled by the shared component loader.
+
+For each layer with `E` experts, hidden size `H`, intermediate size `I`, and
+group size `G`, Olive sidecars under `...layers.L.experts` are:
+
+| Stem | `_qweight` (UINT8) | `_scales` (floating point) | `_qzeros` (UINT8, asymmetric only) |
+|---|---|---|---|
+| `gate_up_proj` | `[E, 2*I, H/2]` | `[E, 2*I, H/G]` | `[E, 2*I, ceil((H/G)/2)]` |
+| `down_proj` | `[E, H, I/2]` | `[E, H, I/G]` | `[E, H, ceil((I/G)/2)]` |
+
+Weights use contiguous gate-then-up rows before graph-level interleaving.
+The adapter preserves packed codes and folds `router.per_expert_scale` into
+down-projection block scales; low-precision folding changes rounding order
+and is not claimed bitwise exact. The learned `router.scale` remains unchanged.
+Incomplete, stale, or shape-incompatible expert sidecars are rejected.
+
+Float projection packing retains version-1
+`pkg.mobius.qkv_projection_sources` metadata on the `Split`: original initializer
+names and output widths. Keep it through serialization so PEFT attachment can
+restore independent projections after weight folding. INT4 projection packing
+requires homogeneous bits/block size/accuracy level, scale dtype, zero-point
+presence/layout, and a known shape-generic fallback body when one is registered.
+Float groups whose source initializer or logical weight view has outside
+consumers are conservatively left unpacked, preserving existing tied-weight
+adapter attachment rather than inventing multiple-group restoration semantics.
+
+The legacy `mobius.qkv_projection_sources` key is accepted when reading older
+exports; new exports write `pkg.mobius.qkv_projection_sources`.
+
+Restored projections remain lazy slices, not eager copies of a large packed
+checkpoint. Their backing-file ownership is registered with the package save
+alias guard. Saving over a backing file (same path, hard link or symlink) is
+rejected before any output is written, with the source unchanged. Use a distinct,
+non-aliasing output location; ordinary packages without such derived lazy
+dependencies retain their existing save behavior.
+
+Backing-file ownership is also retained on the in-memory model, so initializer
+folding cannot hide it by replacing the restored tensor wrappers. A regression
+refolds the lazy projections before saving, checks refusal before any external
+read or source mutation, then verifies an exact round-trip at a safe destination.
+This metadata is transient, not serialized provenance or a guarantee for arbitrary
+model-replacement transformations.

docs/model-catalog.md

@@ -199,6 +199,10 @@ Models combining Mamba/SSM layers with transformer attention layers.
 
 Models that process both images and text.
 
+For Gemma4 26B-A4B, see the [draft runtime qualification and FP32 pixel input
+contract](gemma4-runtime.md), including the failed strict parity gate and limits
+of the prior evaluation evidence.
+
 | Model Type | Module Class | Task | Example HuggingFace Model |
 |---|---|---|---|
 | `llava` | `LLaVAModel` | `vision-language` | `llava-hf/llava-1.5-7b-hf` |

examples/gemma4_multimodal.py

@@ -132,7 +132,7 @@ def prepare_vision_feeds(
     Args:
         processor: ``AutoProcessor`` loaded for the Gemma 4 model.
         image_path: Path to a local image file (JPEG, PNG, etc.).
-        np_dtype: Numpy dtype for pixel values (default float32).
+        np_dtype: Retained for caller compatibility; pixels always use float32.
 
     Returns:
         ``{"pixel_values": [B, N, 3*P^2], "pixel_position_ids": int64[B, N, 2]}``
@@ -143,7 +143,8 @@ def prepare_vision_feeds(
     # Gemma 4 processor requires a text argument alongside images.
     # Use a placeholder image token so the processor computes correct dims.
     processed = processor(images=img, text="<image>", return_tensors="np")
-    pixel_values = processed["pixel_values"].astype(np_dtype)
+    # The vision ABI is FLOAT even when decoder/projection weights use BF16.
+    pixel_values = processed["pixel_values"].astype(np.float32)
     # The HF processor returns "image_position_ids"; our ONNX vision model input is "pixel_position_ids"
     pixel_position_ids = processed["image_position_ids"].astype(np.int64)
     return {"pixel_values": pixel_values, "pixel_position_ids": pixel_position_ids}