[Tianlei Wu (tianleiwu)]

Description

Stacked on #2713 (base kvaishnavi/qwen-38-flash-create, so the diff is only this change). Adds an off-by-default merge_shared_input_projections option to the Qwen3.8 Flash Next builder that merges projections reading the same activation into one MatMul followed by Split. The graph interface (inputs, outputs, state tensors, tensor names consumed downstream) is unchanged, so no runtime change is needed.

Motivation and Context

Our research decode stack (a different export, built from the Mobius exporter and rewritten by a graph pass, not from this builder) measured these launch-count reductions at S=1 CUDA-graph replay, 2k context:

Rewrite	Measured	Exactness
GDN qk_l2_norm=1 + pre-cast A_log/dt_bias	-1.03 ms	tolerance-level (KL 0.0017); cast fold alone is bit-exact
Recurrent state kept in op layout + constant Add(w, 1) folded	-0.48 ms	bit-exact
97 merged shared-input projection groups	-1.19 ms	bit-identical

Before adding anything I checked what #2713's builder already emits for the decoder. The first two rows are already in #2713: GatedDeltaNet is emitted with qk_l2_norm=1, A_log/dt_bias are float32 initializers consumed by the op, the recurrent state is the op's V-major layout with no transposes (paged path), the +1 on the RMSNorm-style weights is folded at export, and no constant Cast/Add or router Softmax remains. This PR therefore does not touch them; it only adds a test that pins that behavior. What is left is the third row: #2713 already packs QKV+Z, A+B and the hyper-connection down/inject, but these groups still read one activation through separate MatMuls.

Summary of Changes

New option

merge_shared_input_projections=true (extra option) or target_options.optimizations.merge_shared_input_projections (structured config). Default false. Qwen3.8 Flash Next only; the MTP head is exported unchanged.

Layer	Merged (one MatMul + Split per layer)	Replaces
GatedDeltaNet (36 of 48 layers)	QKV, Z, A, B	packed QKV/Z + packed A/B (2 MatMuls)
Sparse attention (12 layers)	packed Q/K/V, output gate, indexer Q/K	3 MatMuls (needs the default fuse_qkv)
MoE (48 layers)	router, shared-expert gate, shared-expert up	2 MatMuls (router + fused gate/up)

That is 96 groups for the official 48-layer checkpoint (the research stack had 97). The shared expert's one-wide output gate stays separate because a width that is not a multiple of 8 would misalign the merged GEMM. Each Split output keeps the name the unmerged MatMul output had. The option requires a dense, bias-free fp32/fp16/bf16 export and raises a ValueError otherwise instead of silently skipping groups.

Files

File	Change
builders/qwen3_8.py	make_merged_projections, router/shared-expert and sparse-attention merge paths, GDN override, option validation, drop the option for the MTP head
builders/qwen3_5.py	Split make_linear_attention_input_proj so a subclass can supply all four projections (make_linear_attention_projections, default behavior unchanged)
builder.py, builder_config.py	Bool parsing, help text, target_options.optimizations field
README.md, docs/ModelBuilderConfiguration.md	Option documentation and legacy-name mapping
test/python/builder/_qwen38_tiny.py	Synthetic official-shaped NVFP4 checkpoint generator (any dimensions)
test/python/builder/test_qwen38_merged_projections.py, test_builder_config.py	New tests

Testing

Run on CUDA GPU 1 (H200); the local transformers lacked qwen4_exp, so the tests ran with a local shim that registers the config class (CI environments that already run test_qwen38_nvfp4.py do not need it).

• pytest test/python/builder/test_qwen38_merged_projections.py test/python/builder/test_builder_config.py: all pass (20 new merge tests + 3 new config tests). They export the synthetic checkpoint through parse_extra_options/create_model in dense and paged layouts and check: default graph already has qk_l2_norm=1, fp32 A_log/dt_bias, no constant Cast/Add, no state transposes in paged; merged graph changes only the expected MatMul/Split counts, keeps inputs/outputs and every produced tensor name, has no dangling inputs and the merged weights are the bitwise concatenation of the default ones; sub-graph outputs match the default export (CPU fp32 allclose, CUDA bf16 within one bf16 step); unsupported inputs (quantized, biased, unaligned, mixed quantization exclusion, int4 export) are rejected; MTP head does not receive the option. Mutating the Split output order makes the numerical tests fail.
• test_qwen38_contrib_ops.py, test_qkv_fusion.py, test_qwen_gated_add.py, test_decoder_state_groups.py, test_qwen_mtp_export.py: the same 14 tests fail with and without this change (stale tests in the Add Qwen-3.8 Flash Next to model builder #2713 branch); everything else passes. test_qwen38_nvfp4.py could not be imported locally (needs Qwen4ExpVisionModel).
• Also exported a 4-layer synthetic checkpoint with the official tensor shapes (hidden 2560, 512 experts, 48 value heads, ...) with and without the option, and ran the extracted projection sub-graphs for all four group types on CUDA bf16 at 1 to 2048 tokens: outputs are bitwise identical (this was run locally, not added to the test suite).
• ruff check and ruff format --check on the changed test files.

Indicative GEMM-level microbenchmark (PyTorch/cuBLAS CUDA graph, official shapes, bf16, M=1, merged GEMM plus Split copies versus separate GEMMs; not ORT and not this PR's graph): about 16.5 us saved per GatedDeltaNet group, 8.5 us per sparse-attention group, 11.5 us per MoE group, about 1.25 ms over 36/12/48 layers.

Not validated

• No end-to-end export of the 74 GB checkpoint and no decode: the exported graph needs the contrib ops from the ORT stack (not available in a stock ORT) and the Add build and runtime support for Qwen-3.8 Flash Next #2714 runtime. Decode speedup and token agreement of this PR's output are unmeasured; the research numbers above are from a different export.
• Not tested with -p int4/int8 (the option rejects them), fuse_qkv=false for the sparse-attention group (kept separate by design) or on non-CUDA EPs.
• Merging changes GEMM shapes, so exactness depends on cuBLAS heuristics; it was bit-identical for the shapes and token counts above on H200 only.

Checklist

• Tests added/updated
• No breaking changes (option is off by default; default export is unchanged)
• Documentation updated