Add merge_shared_input_projections option to the Qwen3.8 Flash builder

docs/ModelBuilderConfiguration.md

@@ -27,7 +27,7 @@ These fields are siblings of `type` in an Olive `ModelBuilder` pass. The builder
 | Field | Owner and purpose |
 | --- | --- |
 | `builder_config_version` | Structured builder schema version, proposed as `2`. Not written to the generated runtime config. |
-| `target_options` | Target text-decoder export policy: `quant_config` and `attention`. |
+| `target_options` | Target text-decoder export policy: `quant_config`, `attention`, and `optimizations`. |
 | `drafter_options` | Drafter selection, source, quantization, attention, `shared_weights`, and type-specific export settings. |
 | `speculative_options` | Graph requirements connecting the target and drafter. |
 | `runtime_config` | Inline runtime JSON fragment or a path/resource referencing one. |
@@ -97,7 +97,13 @@ Explicit selections must not silently fall back. Head counts and other checkpoin
 
 Pool size, scheduler limits, utilization targets, and prefill chunk size belong to `runtime_config`. They remain constrained by the exported graph and supported runtime behavior even though they are not graph-construction options.
 
-Compatible projection optimizations are automatic and are not configuration options. Numeric precision and layouts remain controlled by the existing quantization settings.
+### Graph Optimizations
+
+Compatible projection optimizations are automatic and are not configuration options, with one opt-in exception described below. Numeric precision and layouts remain controlled by the existing quantization settings.
+
+`target_options.optimizations.merge_shared_input_projections` (Qwen3.8 Flash Next only, boolean, default `false`) maps to the legacy `merge_shared_input_projections` extra option. It merges the projections that read the same activation into one `MatMul` over row-concatenated weights plus one `Split` per group: GatedDeltaNet QKV/Z/A/B, sparse-attention packed Q/K/V with the output gate and indexer Q/K, and the MoE router with the shared expert's gate/up. Each `Split` output keeps the name of the unmerged `MatMul` output, so the graph inputs, outputs, and state names are unchanged. The shared expert's one-wide output gate stays separate.
+
+The option requires dense, bias-free projections in an fp32, fp16, or bf16 export and rejects configurations it cannot merge rather than silently skipping them. The one exception is the sparse-attention group, which stays unmerged when its Q/K/V projections are not packed. It applies to the decoder only (the MTP head is exported without it), and other model builders ignore it. When both spellings are present, the structured field takes precedence with a warning.
 
 ## 4. Drafter Configuration
 
@@ -609,6 +615,7 @@ The same component ownership can later accommodate audio without changing the ta
 | `paged_chunk_size` | Runtime `search.chunk_size` |
 | `enable_cuda_graph` | Runtime decoder session/provider settings |
 | `use_device_allocator_for_initializers` | Preserved, default `false`; paged Qwen3.8 sessions enable it automatically |
+| `merge_shared_input_projections` | `target_options.optimizations.merge_shared_input_projections` |
 | Olive `search` | Runtime `search` |
 | Any other `extra_options` key | No canonical destination yet; see rule 8 |
 

src/python/py/models/README.md

@@ -122,7 +122,7 @@ python src/python/py/models/builder.py -i path_to_dense_checkpoint -o output -e
   --runtime_config '{"search":{"max_length":128}}'
 ```
 
-`target_options` routes `quant_config` and `attention` to the existing exporter. Target MLP gate/up, QKV, and eligible Q/K normalization fusion are automatic and have no configuration options. `quant_config.format` is the canonical graph-layout key; `runtime` remains a parsing alias. The target rejects an explicit checkpoint policy until its loaders implement both paths. Root CLI `precision` is optional when target weight type is explicit.
+`target_options` routes `quant_config`, `attention`, and `optimizations` to the existing exporter. Target MLP gate/up, QKV, and eligible Q/K normalization fusion are automatic and have no configuration options; the only `optimizations` field is the opt-in `merge_shared_input_projections` for the Qwen3.8 Flash Next decoder (see [Qwen3.8 Merged Shared-Input Projections](docs/model-specific-options.md#qwen38-merged-shared-input-projections)). `quant_config.format` is the canonical graph-layout key; `runtime` remains a parsing alias. The target rejects an explicit checkpoint policy until its loaders implement both paths. Root CLI `precision` is optional when target weight type is explicit.
 
 DFlash2 and DSpark selection requires a local checkpoint `path` and paged target attention. DSpark uses BF16 body I/O; DFlash2 defaults to BF16 but also supports FP16 through `drafter_options.quant_config.io_dtype`. Omitted target taps are inferred from checkpoint `target_layer_ids + 1` before construction. DFlash2 parses `auto`, `required`, and `off` sharing modes and validates adopted quantized-head node attributes. MTP/DSpark currently accept only `auto`. Structured drafter selections that conflict with legacy drafter paths are rejected.
 

src/python/py/models/builder.py

@@ -212,6 +212,7 @@ def check_extra_options(
         "enable_cuda_fpa_intb_gemm",
         "exclude_mtp",
         "use_device_allocator_for_initializers",
+        "merge_shared_input_projections",
     ]
 
     for key in bools:
@@ -1228,6 +1229,9 @@ def get_args():
                     Each per-layer entry is a scalar (per_tensor) or a length-(num_kv_heads * head_size) vector (per_channel).
                     An optional "qmax" records the divisor the file was calibrated with (128 for int8, 8 for int4, 448 for fp8);
                     the builder then rescales to the requested scheme, so one file can serve several bit widths.
+                merge_shared_input_projections = Qwen3.8 Flash Next only. Merge projections that read the same activation
+                    (linear-attention QKV/Z/A/B, sparse-attention QKV/gate/indexer, and the router with the shared expert's gate/up)
+                    into one MatMul plus Split per group. Default is false. Requires fp32, fp16, or bf16 dense projections; the MTP head is unchanged.
                 use_webgpu_fp32 = Use FP32 I/O precision for WebGPU EP.
                     Use this option to enable GPUs that do not support FP16 on WebGPU (e.g. GTX 10xx).
                 use_cuda_bf16 = Use BF16 I/O precision in quantized ONNX models for CUDA EP.

src/python/py/models/builder_config.py

@@ -914,7 +914,16 @@ def flatten_target_options(
             flattened[legacy_key] = kv_cache[field_name]
 
     optimizations = options.get("optimizations", {})
-    check_fields(optimizations, set(), "target_options.optimizations")
+    check_fields(optimizations, {"merge_shared_input_projections"}, "target_options.optimizations")
+    if "merge_shared_input_projections" in optimizations:
+        if not isinstance(optimizations["merge_shared_input_projections"], bool):
+            raise ValueError("target_options.optimizations.merge_shared_input_projections must be a boolean")
+        warn_structured_override(
+            legacy_options,
+            "merge_shared_input_projections",
+            "target_options.optimizations.merge_shared_input_projections",
+        )
+        flattened["merge_shared_input_projections"] = optimizations["merge_shared_input_projections"]
 
     return flattened, quant_config, effective_precision
 

src/python/py/models/builders/qwen3_5/text.py

@@ -480,9 +480,7 @@ def make_linear_attention_input_proj(self, layer_id, attention, root_input):
         """
         basename = f"/model/layers.{layer_id}/linear_attn"
 
-        qkv_name, z_name = self.make_linear_attention_qkv_z_proj(layer_id, attention, root_input)
-
-        b_name, a_name = self.make_linear_attention_a_b_proj(layer_id, attention, root_input)
+        qkv_name, z_name, b_name, a_name = self.make_linear_attention_projections(layer_id, attention, root_input)
 
         conv_input = qkv_name
         if not self.use_paged_attention:
@@ -501,6 +499,12 @@ def make_linear_attention_input_proj(self, layer_id, attention, root_input):
 
         return z_name, b_name, a_name, conv_input, conv_weight_name
 
+    def make_linear_attention_projections(self, layer_id, attention, root_input):
+        """Build the QKV, Z, B and A projections, returning the names of their output values."""
+        qkv_name, z_name = self.make_linear_attention_qkv_z_proj(layer_id, attention, root_input)
+        b_name, a_name = self.make_linear_attention_a_b_proj(layer_id, attention, root_input)
+        return qkv_name, z_name, b_name, a_name
+
     def make_linear_attention_qkv_z_proj(self, layer_id, attention, root_input):
         basename = f"/model/layers.{layer_id}/linear_attn"
         qkv_name = self.make_matmul(attention.in_proj_qkv, f"{basename}/qkv_proj/MatMul", root_input)

src/python/py/models/builders/qwen3_8/model.py

@@ -100,6 +100,7 @@ def make_mtp_model(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_opti
         mtp_options["filename"] = "mtp.onnx"
         mtp_options.pop("include_hidden_states", None)
         mtp_options.pop("exclude_lm_head", None)
+        mtp_options.pop("merge_shared_input_projections", None)
         self.mtp = Qwen4ExpMTPTextModel(
             copy.deepcopy(config),
             self.mtp_attrs["io_dtype"],