feat: add LFM2-MoE support by ykhrustalev · Pull Request #2575 · microsoft/onnxruntime-genai

Yuri Khrustalev (ykhrustalev) · 2026-09-16T15:22:38Z

Problem
LFM2-MoE models (LFM2.5-8B-A1B, LFM2-24B-A2B) cannot be exported. They reuse the LFM2 conv/attention stack but route the later layers through sparse experts, and their router selects experts on sigmoid(logits) + expert_bias while mixing them with the unbiased sigmoid(logits) — routing the fused MoE/QMoE op cannot express, since it only does softmax over the top-k.

Solution

Adds LFM2MoEModel for Lfm2MoeForCausalLM, exported as model type lfm2_moe: dense MLP for the first num_dense_layers, fused MoE/QMoE with interleaved SwiGLU experts after that
Performs expert selection in the graph and feeds the op log(sigmoid(logits)) at the selected experts and a large negative sentinel elsewhere, so the op's softmax reproduces the model's normalized routing weights exactly on every EP
Keeps the router MatMul out of int4 quantization so rounding cannot flip an expert choice
Registers lfm2_moe in the runtime through ModelType::IsLFM2, reusing the LFM2 conv-state cache and its rewind restriction
Fixes make_moe_init writing moe_attrs keys nothing reads, which made every int4 MoE export emit the float MoE op instead of QMoE
Switches non-CUDA block-wise QMoE weights to the signed-scale MatMulNBits encoding already used for CUDA, so no block extreme is clipped

Testing

LFM2.5-8B-A1B verified against transformers: fp32 matches to 5e-5 (KL 0.0000) on CPU; int4 KL 0.096 / 89% top-1; CUDA fp16 KL 0.006 and int4 KL 0.10 on an A10
New test/python/builder/test_lfm2_moe.py pins the router graph, the dense/MoE layer split, expert repacking, and the QMoE op selection; new C++ case covers the lfm2_moe runtime path

Depends on microsoft/onnxruntime#32644 for usable CPU decode speed on these models.

LFM2-MoE (LFM2.5-8B-A1B, LFM2-24B-A2B) keeps the LFM2 conv/attention stack but routes the later layers through sparse experts. The router selects on sigmoid(logits)+expert_bias and mixes on the unbiased sigmoid, which the fused MoE/QMoE op cannot express, so the selection is done in the graph and the op is fed log(sigmoid) at the selected experts and a sentinel elsewhere. Also fixes make_moe_init writing moe_attrs keys nobody reads, which made every int4 MoE export emit the float MoE op instead of QMoE.

azure-pipelines · 2026-09-16T15:22:48Z

Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run.

Drop the redundant sorted=True on the router TopK (ONNX default), remove comments that restate the code, and state in the README that the signed-scale QMoE encoding was validated on CPU only. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Copilot

🟡 Changes recommended

Quantized-checkpoint export has a blocking failure, and packed expert biases can be lost.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds LFM2-MoE export and runtime support, including graph-based routing and QMoE quantization.

Changes:

Adds dense/MoE builder support, expert repacking, and routed graph selection.
Registers lfm2_moe with LFM2 runtime and cache behavior.
Fixes QMoE initialization and signed blockwise quantization.
Adds builder, runtime, quantization, and documentation coverage.

File summaries

File	Reviewed changes
`test/python/builder/test_qmoe_weights.py`	Tests signed blockwise QMoE quantization.
`test/python/builder/test_lfm2_moe.py`	Tests routing, layer selection, expert packing, and QMoE generation.
`test/cpp/lfm2_cache_layer_types_test.cpp`	Tests `lfm2_moe` runtime handling.
`src/python/py/models/README.md`	Documents LFM2-MoE and quantization behavior.
`src/python/py/models/loaders/quark.py`	Updates quantization references.
`src/python/py/models/builders/lfm2.py`	Implements the LFM2-MoE graph and routing. Findings: critical (2 votes) quantized checkpoints expose `layer.mlp`, causing the current `layer.feed_forward` access to fail; moderate (1 vote) packed expert biases are overwritten with zeros; nit (3 votes) `ROUTER_SENTINEL` is class-scoped; nit (3 votes) the routing method should be split into reusable stages; nit (1 vote) configuration values should be centralized in `moe_attrs`.
`src/python/py/models/builders/base.py`	Fixes QMoE configuration and blockwise quantization.
`src/python/py/models/builders/__init__.py`	Exports the new builder class.
`src/python/py/models/builder.py`	Registers builder dispatch for LFM2-MoE.
`src/models/model.cpp`	Routes `lfm2_moe` through the LFM2 runtime.
`src/models/model_type.h`	Registers the new model type.
`src/generator/generators.cpp`	Applies LFM2 rewind restrictions.

Review details

Suppressed comments (2)

src/python/py/models/builders/lfm2.py:255

The packed-expert IR carries gate_up_bias and down_bias (and the Quark loader populates them), but this unconditional initialization replaces both with zeros. Any pre-quantized checkpoint with nonzero expert projection biases will silently lose them; preserve the packed bias fields in that branch and only synthesize zeros for biasless dense LFM2 experts.

        num_experts = self.moe_attrs["num_experts"]
        self.make_initializer(
            torch.zeros(num_experts, 2 * self.moe_intermediate_size),
            f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias",
            to=self.io_dtype,
        )
        self.make_initializer(
            torch.zeros(num_experts, self.hidden_size),
            f"model.layers.{layer_id}.moe.experts.down_proj.bias",
            to=self.io_dtype,

src/python/py/models/builders/lfm2.py:216

These four new configuration values control the MoE graph, but are stored as ad-hoc model attributes instead of in an _attrs dictionary. Please centralize them in moe_attrs with comments and read them from there, so the model's scenario-specific state remains consistent with the builder design and does not spread across more fields.

        self.num_dense_layers = config.num_dense_layers
        self.moe_intermediate_size = config.moe_intermediate_size
        self.use_expert_bias = config.use_expert_bias
        self.routed_scaling_factor = config.routed_scaling_factor

Files reviewed: 12/12 changed files
Comments generated: 3
Review effort level: Lite

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Resolve the feed-forward module through an accessor so the quantized-checkpoint IR (which names it mlp) works alongside the Hugging Face naming, forward packed expert biases instead of always emitting zeros, move the router sentinel into moe_attrs, and split the router into score, selection and mask stages.

Yuri Khrustalev (ykhrustalev) · 2026-09-16T20:24:17Z

Thanks — all five addressed in fix: address LFM2-MoE review findings.

layer.feed_forward vs layer.mlp: the feed-forward module is now resolved through get_feed_forward_module, which accepts either name, and the w1/w3/w2 aliasing only runs when those attributes exist (the quantized IR already uses gate_proj/up_proj/down_proj). Note the quantized-checkpoint IR still has no operator_norm, conv or ffn_norm, so pre-quantized checkpoints remain unsupported for the whole LFM2 family, as on main — this change just removes the LFM2-MoE-specific obstacle.
Packed expert biases: make_moe_expert_initializers never writes biases, so the zeros were the only source rather than an overwrite. Packed biases are now forwarded when a checkpoint supplies them, with zeros as the fallback (LFM2-MoE experts have no bias).
ROUTER_SENTINEL: moved into moe_attrs["router_sentinel"].
Router method size: split into make_moe_router_scores, make_moe_router_selection and make_moe_router_mask, with make_moe_router composing them.

Re-verified after the refactor: exported fp32 graph still matches transformers to 3e-7 on the dummy model and the int4 graph still emits QMoE; 730 builder tests pass, including new cases for the loader-independent accessor and the bias passthrough.

Tianlei Wu (tianleiwu)

Review Summary

The LFM2-MoE feature itself is solid. The log(sigmoid) + sentinel trick to express LFM2's sigmoid router through the softmax-only fused MoE/QMoE op is mathematically exact, and I verified it line by line against transformers' Lfm2MoeSparseMoeBlock.route_tokens_to_experts: selection on sigmoid + expert_bias, mixing on the unbiased sigmoid, renormalized over the selected experts. Moving routed_scaling_factor from the routing weights to a Mul on the op output is algebraically equivalent (expert mixing is linear in the weights) and cheaper. Splitting make_moe_router into four small emitters addresses the earlier "one large method" feedback well, and the new 358-line test file pins the exact node sequence, wiring, attributes and sentinel representability rather than just asserting "it builds".

My concerns are almost entirely about the two cross-cutting MoE changes bundled with the new model, not about LFM2 itself.

Verification performed on `5ee16167`

Check	Result
`pytest test/python/builder`	729 passed, 4 skipped
Router math vs `transformers` `Lfm2MoeSparseMoeBlock`	matches
`gate_up_proj` interleave vs `Qwen2MoeExperts` (`chunk(2, dim=-1)`)	correct
`Lfm2MoeConfig` attribute availability	all builder-read fields present
`moe_attrs` readers/writers after rename	consistent, no stale keys in `src/`
CPU QMoE dequant (`contrib_ops/cpu/moe/moe_quantization_cpu.cc`)	`(q - 8) * scale` — signed scale is safe
WebGPU QMoE (`contrib_ops/webgpu/moe/qmoe.cc`)	routes through `ApplyMatMulNBits` — MatMulNBits layout is exactly what it wants
Origin of the `moe_attrs` key bug	introduced by #2386 (`git log -L 1030,1040:src/python/py/models/builders/base.py`)

Top concerns (details inline)

The make_moe_init key fix is broader than the description says and should probably be its own PR — see the comment on base.py.
The signed-scale QMoE switch changes exported bytes for every non-CUDA MoE model, and TRT-RTX is unverified.
int4 MoE export on the DML EP now raises where it previously (incorrectly) succeeded.

Findings outside the diff

CudaQuantizer.symmetric_blockwise_quantize (src/python/py/models/quantization/cuda_quantizer.py:623) now has zero callers. The only surviving references are the two loaders/quark.py docstrings this PR rewrites to point at it. Please delete it, or add a note on why it is retained.
The non-CUDA blockwise branch does not validate block_size, while the CUDA branch rejects anything outside (32, 64, 128). Now that both paths share CudaQuantizer.matmulnbits_blockwise_quantize, the check could be hoisted.

Minor notes

The LFM2MoEModel docstring names the router Lfm2MoeTopKRouter; upstream the class is Lfm2MoeSparseMoeBlock (with route_tokens_to_experts). Worth correcting so the reference is greppable.
HF normalizes with routing_weights / (sum + 1e-6); the exported graph normalizes without the epsilon. The delta is negligible, but a one-line comment would document the known parity gap.
moe_attrs["router_sentinel"] and moe_attrs["router_probs"] are not documented in the moe_attrs declaration block (base.py:413), which the model-builder guidelines ask for. router_probs is also not a setting — it is a per-layer graph value name threaded from make_moe_router to make_moe_subgraph. Consider returning it like Qwen3MoEModel.make_moe_subgraph does with its combine output.

Stale threads

The three open threads from the automated reviewer were filed against e720558c and look fully addressed by de40a22f / 5ee16167 (get_feed_forward_module, moe_attrs["router_sentinel"], and the four-way router split).

…diate_size_init hook Restore the strict block_auto_adjust_ff_dim read for dense LFM2 so a config missing the key raises instead of silently inverting transformers' default of True. LFM2MoEModel overrides the hook to use intermediate_size as is, since Lfm2MoeConfig has no such field.

…nd QMoE encoding scope Router graph: - Clamp the gathered sigmoid scores before Log so a score that flushed to zero cannot fall below the sentinel and let the MoE op's internal top-k pick a different expert than the graph selected. - Gather the selected experts before Log so Log runs over top_k entries instead of num_experts. - Emit Log, TopK, GatherElements and ScatterElements through new make_<op> wrappers on Model instead of raw make_node calls. - Return the masked router scores from make_moe_router and thread them into make_moe_subgraph instead of stashing them in moe_attrs; document router_sentinel in the moe_attrs declaration block. - Move the router MatMul quantization exclusion out of the score emitter into make_moe_preprocessing so graph emission has no side effects on the checkpoint IR. - Name the upstream router (Lfm2MoeSparseMoeBlock.route_tokens_to_experts) and note the missing 1e-6 normalization epsilon. Reuse: - Add Model.make_moe_expert_names and Model.make_interleaved_swiglu_moe_preprocessing; LFM2-MoE and Qwen3.5-MoE now share the gate/up interleave, zero expert biases and expert initializer names, and make_moe_expert_initializers uses the same helper. - Drop the LFM2-MoE packed-expert bias passthrough and its tests: no loader produces packed LFM2-MoE experts yet. - Hoist the TRT-RTX swiglu_limit workaround out of the LFM2-MoE and Qwen3.5-MoE constructors into make_moe_init. QMoE encoding: - Keep TRT-RTX on the original symmetric blockwise encoding until the signed-scale MatMulNBits grid has been measured on that EP; CPU and WebGPU stay on the signed-scale grid. This restores _symmetric_blockwise_quantize with a single caller. - Hoist the block_size 32/64/128 check so CPU and WebGPU validate it too. README: split the block-wise QMoE paragraph into encoding contract, ORT_USE_MLAS_Q4_GEMM_MOE note, and the onnxruntime#32644 dependency, and attribute the decode numbers to LFM2.5-8B-A1B int4 on a 30-core x86 host. Verified: tiny random LFM2-MoE exported fp32 and int4 on CPU and run in ONNX Runtime 1.30 against transformers 5.17; fp32 logits match to 2e-7 and int4 KL is 1.7e-4. test/python/builder: 738 passed, 3 skipped.

Yuri Khrustalev (ykhrustalev) · 2026-09-16T23:01:30Z

Addressed the summary-level items in f763e12 as well: symmetric_blockwise_quantize is no longer dead (TRT-RTX uses it), block-size validation is shared, the docstring names Lfm2MoeSparseMoeBlock.route_tokens_to_experts, the missing 1e-6 normalization epsilon is noted in make_moe_router_mask, and router_sentinel is documented in the moe_attrs block while router_probs is returned instead of stored. The make_moe_init key fix lives in #2581; this branch keeps its copy until that merges.

…class CodeQL flagged the LFM2-MoE override for requiring five positional arguments while Model.make_moe_subgraph requires four. Declare router_probs as an optional parameter on both, document it on the base method, and raise in the LFM2-MoE override when it is missing, since that builder cannot address its router output by name.

Tianlei Wu (tianleiwu)

Reviewed current head 161f6032. Most earlier concerns are addressed. The complete Python builder suite passes: 738 passed, 4 skipped; the focused LFM2-MoE/QMoE subset passes 87 tests, 3 skipped. git diff --check also passes.

One correctness issue remains: the omitted HF routing-normalization epsilon is not uniformly a rounding-level difference. Executing a real builder-generated FP32 MoE graph against a tiny Transformers LFM2-MoE model demonstrates a material output-magnitude error for small selected scores and a nonzero output when the reference returns zero. Numerical evidence and the proposed correction are on the existing router thread, rather than duplicated as a new inline thread. Please preserve that normalization factor and cover it with an executed-graph parity test. The earlier review summary's characterization of this gap as negligible was too broad.

The two new inline notes are nonblocking model-builder convention suggestions. I accept separating the shared MoE key fix into #2581; that PR remains open, so this branch still needs the planned rebase after it lands. The prior fixed/accepted discussions have been resolved.

Local execution used CPU ORT 1.28.0 and Transformers 5.15.0 with controlled inputs, not full pretrained-checkpoint evaluation. Native tests were not rebuilt at this head, and CUDA/WebGPU/TRT-RTX execution was not revalidated.

Hugging Face mixes the selected experts with sigmoid_i / (sum_selected + 1e-6) while the fused MoE/QMoE op normalizes to sigmoid_i / sum_selected. For tokens whose selected scores are all tiny the difference is not rounding level: with four selected logits of -16 the HF weights sum to 0.31, and when the scores flush to zero HF returns exactly zero while the op still emits a mixture. - make_moe_router now also returns a per-token output scale sum_selected / (sum_selected + 1e-6), summed before the clamp so flushed scores give a zero factor, with routed_scaling_factor folded in. It is reshaped to the MoE output's leading dims and multiplied into the op output in make_moe_subgraph. - Emit the router sigmoid as 1 / (1 + exp(-x)). ONNX Runtime's CPU Sigmoid kernel has an absolute error of about 6e-8, returning 5.96e-8 for x = -16 instead of 1.125e-7 and 0 below x = -17; since the scores are logged and summed that is a large relative error. The explicit form matches torch.sigmoid to fp32 rounding. Adds make_exp, make_neg and make_reciprocal wrappers. - Move num_dense_layers, use_expert_bias and routed_scaling_factor into moe_attrs, documented in the base declaration block, and move make_moe_router_shape onto Model. - Add an executed-graph test: the emitted MoE layer runs in ORT and is compared with HF's routing math for ordinary scores, the 0.31-mass case and the flushed-to-zero case, with and without routed_scaling_factor. Verified against a tiny random Lfm2MoeForCausalLM in transformers 5.17 on ORT 1.30: fp32 logits match to 1.8e-7, int4 KL 1.7e-4. test/python/builder: 743 passed, 3 skipped.

…MoE executed graph The executed-graph test only exercised ordinary, tiny and fully flushed router scores. Add a partly flushed token, which is the only case that catches a clamp floor leaking weight to flushed experts, and a token whose selected sigmoids sum to exactly HF's 1e-6 epsilon, where HF keeps half the routing mass. Replace the loose 0.2 < ratio < 1.6 check with exact ratios against the epsilon-free normalization, so the fixture fails if it drifts out of the regime where the output scale matters. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…M kernels (#32644) ### Description **Problem** On CPU, QMoE runs block-wise 4-bit and 8-bit experts slowly. Before every matrix multiply it converts the expert weights back to fp32, and it gives each expert its own thread. During decode only `top_k` experts run, so most threads sit idle. On an 8B-A1B model that means about 3 tokens/s. The existing MLAS Q4 fast path doesn't solve this: it only runs on AVX-512, is off by default, and re-quantizes the weights. **Solution** - **Pack weights once at load time.** Block-wise 4-bit and 8-bit experts are packed for the same MLAS kernels that `MatMulNBits` uses. The weights aren't re-quantized, and no fp32 copy is kept. Zero points work too, because QMoE stores them the same way `MatMulNBits` does. - **Use the threads better.** When fewer experts are active than there are threads, experts run one after another and each matrix multiply gets the whole thread pool. - **New `accuracy_level` attribute on `QMoE`, with the same meaning as in `MatMulNBits`:** - 4-bit: the default keeps fp32 activations. `accuracy_level=4` switches to int8 activations, which is faster but slightly less accurate. - 8-bit: MLAS has no fp32 kernel, so int8 activations are used at every level. - **Environment override:** `ORT_QMOE_CPU_QNBIT_GEMM` takes precedence over the attribute. - `fp32` forces fp32 activations, so 8-bit experts go back to the old path. - `int8` forces int8 activations. - `0` turns the new path off. - Any other value is an error. - **Shared pre-packed weights across sessions.** The packed buffer is tagged with how it was packed. A buffer packed differently is rejected instead of being used. - **Fallback warning.** If a block-wise node can't use the new path, a warning with the reason is logged once, because the old path is about 10x slower for decode. - **Unchanged:** row-wise scales, 2-bit weights, and other execution providers. **Testing** - New `MoETest.QMoETest_CPU_Int4_BlockWise_*` and `QMoETest_CPU_Int8_BlockWise_*` tests compare against an fp32 reference. They cover decode, prefill, bias with block size 64, fp16, `accuracy_level=4`, zero points, each environment variable setting, and shared pre-packed weights across two sessions. All `MoETest.*` and `MatMulNBits*` tests pass in `onnxruntime_provider_test` on Linux x64 (AVX-512). - **Model benchmark:** - **Model:** LFM2.5-8B-A1B (32 experts, top-4, int4 with block size 32), exported with microsoft/onnxruntime-genai#2575. - **Machine:** Xeon 8358, 30 threads, same build for every row. - **Decode:** genai greedy generation. - **Prefill:** a single `InferenceSession.run`. - **Accuracy:** KL divergence against the Hugging Face fp32 logits on an 81-token prompt. | ORT_QMOE_CPU_QNBIT_GEMM | decode tok/s | prefill 256 tok/s | prefill 16 tok/s | KL vs HF fp32 | top-1 agree | |-----------------------------|-------------:|------------------:|-----------------:|--------------:|------------:| | `0` (old path) | 3.2 | 147 | 27 | 0.096 | 0.901 | | unset (fp32, the default) | 55.2 | 346 | 148 | 0.098 | 0.889 | | `int8` (`accuracy_level=4`) | 60.7 | 364 | 103 | 0.106 | 0.889 | An earlier int8-only run on an M3 Ultra gave: - **Decode:** 4 → 125 tok/s. - **Prefill (256 tokens):** 190 → 855 tok/s. - **KL divergence:** 0.112 → 0.101. The fp32 default hasn't been measured on Apple Silicon yet. There it uses the NEON fp32 kernels, because KleidiAI only provides int8. ### Motivation and Context This makes block-wise quantized MoE models fast enough for CPU decode by reusing the `MatMulNBits` kernels instead of converting weights back to fp32 on every call. fp32 activations stay the default where MLAS supports them because int8 loses a little accuracy. Models can opt in to int8 with `accuracy_level=4`, as with `MatMulNBits`. This pairs with microsoft/onnxruntime-genai#2575, which exports LFM2-MoE models in this format. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

Tianlei Wu (tianleiwu)

Re-reviewed at 519a6ad6.

The blocking issue from my previous round is resolved: make_moe_router_selected_scores / make_moe_router_output_scale now apply sum_selected / (sum_selected + 1e-6) with routed_scaling_factor folded in, the sum is taken before the clamp so a flushed-to-zero token yields exactly zero, and test_lfm2_moe_executed_graph_matches_hf_routing executes the emitted graph in ORT against the HF routing math for the ordinary, tiny-mass, fully-flushed, partly-flushed and exactly-epsilon cases. The spelled-out Reciprocal(1 + Exp(Neg(x))) sigmoid is the right call given the CPU Sigmoid kernel's absolute-error profile. I have resolved that thread and the two convention threads.

What is left is all in the CPU/WebGPU quantization switch, which is the part of this PR with the widest blast radius (it changes the bytes of every non-CUDA MoE export, not just LFM2-MoE). Two of them I would fix before merge:

qweight last dim now keeps the block padding. _matmulnbits_blockwise_quantize returns [N, num_blocks * ceil(block_size/pack)]; _symmetric_blockwise_quantize trimmed it to [N, ceil(K/pack)]. ORT's moe_helper.h requires the unpadded shape, so any MoE model whose hidden_size or moe_intermediate_size is not a multiple of the block size now produces a model the QMoE op rejects at session creation. Inline comment with details.
qmoe_block_size 16 and 256 now fail on CPU/WebGPU, but README.md:586 and builder.py:702 still document the 32/64/128 restriction as CUDA-only. The builder instructions require option constraints to be consistent across both files.

One observation worth capturing rather than fixing: the switch also flips the INT8 byte encoding from two's complement to q + 128. ORT CPU uses default_zp_8bit = 128.0f, so the new encoding is the correct one and the previous INT8 CPU/WebGPU export was off by 128 on every weight — a second real bug fix hiding in this PR. Only INT4 is measured and tested today; please add INT8 coverage and mention it in the description so it is not lost.

Two nitpicks are inline (an unreachable ValueError guard, and make_moe_router_shape hardcoding the non-paged row dim). Also, test_lfm2_moe.py re-implements the builder-module import shim already in test_qmoe_weights.py, and _moe_attrs() hardcodes a snapshot of the moe_attrs defaults that will drift from base.py — worth deriving from the real class.

Nothing else blocking. The dedup into make_interleaved_swiglu_moe_preprocessing / make_moe_expert_names, holding TRT-RTX on the old encoding, and excluding the router MatMul from int4 all look right, and the op_type regression fix is valuable well beyond this model.

…dress review findings - Trim the block padding off the CPU/WebGPU/CUDA-raw MatMulNBits QMoE storage so qweight stays [E, N, ceil(K/pack)], which is what the QMoE op validates. Exports whose hidden_size or moe_intermediate_size is not a multiple of qmoe_block_size failed at session creation with a shape mismatch. - Document that the 32/64/128 block-size restriction applies to CPU, CUDA and WebGPU (TRT-RTX still accepts 16/256) in README.md and builder.py. - Derive the MoE router row dim from use_paged_attention (num_tokens vs batch_size * sequence_length), use the helper in Qwen3.5-MoE too, and pin the Qwen router Reshape shape for both token layouts. - Drop the unreachable router-output guard in LFM2MoEModel.make_moe_subgraph. - Move the moe_attrs defaults into a make_moe_attrs_init(config) hook so tests build them from the real class instead of a hardcoded snapshot. - Share the builder-module import shim between test_lfm2_moe.py and test_qmoe_weights.py (_builder_test_utils.py); add INT8 offset-128 encoding and padding-trim coverage.

Tianlei Wu (tianleiwu)

Re-reviewed at 93aa6da0. The four open concerns from the previous round are addressed, and I have resolved those threads. No remaining actionable findings.

The final padding fix retains unpadded CPU/raw QMoE weights, rejects odd raw INT4 dimensions, and explicitly rejects partial blocks on WebGPU, whose MatMulNBits kernels require whole-block strides. The new executed CPU tests cover partial blocks in both projections for INT4 and INT8. README and CLI help now agree on these constraints. The paged router shape and unreachable guard comments are also addressed.

Validation on this head: 768 Python builder tests passed, 4 skipped, including the new executed partial-block cases; git diff --check passed. An additional CPU probe of the retained trimming implementation with ORT 1.30.0 agreed with an independently dequantized expert reference to 3.73e-8 max absolute error for aligned and partial-block INT4/INT8 cases. Native tests were not rebuilt, and CUDA, WebGPU, TRT-RTX, and full pretrained-checkpoint inference were not run in this round.

The accepted split/rebase arrangement for the shared MoE key fix in #2581 still applies; that PR remains open. This approval supersedes my earlier changes-requested review.

Tianlei Wu (tianleiwu) · 2026-09-18T04:03:24Z

/azp run macOS arm64 Build

azure-pipelines · 2026-09-18T04:03:34Z

Azure Pipelines: Successfully started running 1 pipeline(s).

**Problem** LFM2 worked here as a text-only model. The LFM2-VL models, which also read pictures, did not work at all. Nothing turned a picture into something the model could read, and nothing made room for it in the prompt. **Solution** - Adds picture handling for LFM2-VL: shrink the picture to a size the model likes, cut it into small squares, and pass those along with a note saying which squares are real picture and how they were laid out. - Swaps each `<image>` marker in the prompt for the right number of empty slots, so the picture and the text line up. The prompt has to carry exactly one `<image>` per picture, in picture order; any other count is refused up front, the same rule the Hugging Face processor applies, rather than guessed at. One prompt at a time: a one-item list works, longer lists are refused instead of silently losing the text. - Registers the two new model names, so the loader knows how to run them: one for the picture model, one for its text half. - Adds a setting for how many squares each picture gets padded to. - Adds a tutorial and a ready-made config file. Known gaps: a very large picture is shrunk rather than cut into tiles, so it loses some detail, and a picture more than 200 times wider than it is tall (or the reverse) is refused by the resize step. **Testing** - New tests check the counting: how many squares a picture turns into, and how many prompt slots it needs. - A tiny LFM2-VL model is checked in under `test/models/lfm2-vl`, so CI runs the whole pipeline: the picture processor, the hand-off into the decoder, and generation compared against a plain ONNX Runtime replay, with and without a picture and across prefill chunk boundaries. It also pins the refusals above, the fp16 picture-input path, and the pixel values against the Hugging Face processor when `transformers` and `torchvision` are installed (added to the CPU test requirements, except on Windows ARM64 where no `torchvision` wheel exists). - Ran all six published models on Linux CPU, LFM2-VL 450M, 1.6B, 3B and LFM2.5-VL 450M, 1.6B, 3B, and compared against the original PyTorch version. What we hand the model is the same, except pixel values, which differ by at most 1 out of 255. - One picture, two pictures, and no picture all give sensible answers. The smallest model writes the same caption as PyTorch, word for word. **Also in this PR** - The Windows ARM64 job now force-reinstalls the freshly built wheel. The self-hosted runner kept an older install with the same version number, so pip skipped the new one and the Python tests ran against a stale build. - `IsLFM2`, the model-type list and the rewind guard are also touched by #2575 (LFM2-MoE); whichever lands second needs a small rebase so `IsLFM2` covers `lfm2`, `lfm2_moe` and `lfm2_vl_text`. --------- Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Tianlei Wu (tianleiwu) · 2026-09-18T04:23:56Z

Yuri Khrustalev (@ykhrustalev), please resolve conflicts

LFM2-VL (microsoft#2571) landed on main and touches the same LFM2 files, so every conflict is two model types being added in the same place: - model_type.h: keep both "lfm2_moe" and "lfm2_vl_text" in the LLM list (now 33) and in IsLFM2, which both need the conv state cache. - generators.cpp: take main's refactor; its ModelType::IsLFM2 call already covers lfm2_moe, so the rewind guard needs no further change. - builder.py / README.md: keep both architectures and both bullets. - lfm2_cache_layer_types_test.cpp: both sides added the same test for a different model type. Share the body in an ExpectLfm2Runtime helper and keep main's stronger rewind assertions for both.

Tianlei Wu (tianleiwu) · 2026-09-18T18:31:09Z

/azp run macOS arm64 Build

azure-pipelines · 2026-09-18T18:31:20Z

Azure Pipelines: Successfully started running 1 pipeline(s).

Copilot AI lite review requested due to automatic review settings September 16, 2026 15:22

Yuri Khrustalev (ykhrustalev) requested a review from a team as a code owner September 16, 2026 15:22

Copilot started reviewing on behalf of Yuri Khrustalev (ykhrustalev) September 16, 2026 15:23

View session

Yuri Khrustalev (ykhrustalev) mentioned this pull request Sep 16, 2026

perf: run block-wise int4/int8 QMoE CPU experts on the MLAS QNBit GEMM kernels microsoft/onnxruntime#32644

Merged

Copilot AI reviewed Sep 16, 2026

View reviewed changes

Comment thread src/python/py/models/builders/lfm2.py Outdated

Comment thread src/python/py/models/builders/lfm2.py Outdated

Comment thread src/python/py/models/builders/lfm2.py Outdated

Yuri Khrustalev (ykhrustalev) force-pushed the ykhrustalev/add-lfm2-moe-support branch from de40a22 to 774a6a2 Compare September 16, 2026 20:21

Yuri Khrustalev (ykhrustalev) force-pushed the ykhrustalev/add-lfm2-moe-support branch from 774a6a2 to 5ee1616 Compare September 16, 2026 20:21

Tianlei Wu (tianleiwu) reviewed Sep 16, 2026

View reviewed changes

Yuri Khrustalev (ykhrustalev) mentioned this pull request Sep 16, 2026

fix: restore the moe_attrs keys the MoE emitters read #2581

Closed

Yuri Khrustalev (ykhrustalev) added 2 commits September 16, 2026 18:16

github-advanced-security Bot found potential problems Sep 16, 2026

View reviewed changes

Comment thread src/python/py/models/builders/lfm2.py Fixed

Tianlei Wu (tianleiwu) requested changes Sep 16, 2026

View reviewed changes

Comment thread src/python/py/models/builders/lfm2.py Outdated

Comment thread src/python/py/models/builders/lfm2.py Outdated

Yuri Khrustalev (ykhrustalev) and others added 2 commits September 16, 2026 19:42

Tianlei Wu (tianleiwu) reviewed Sep 17, 2026

View reviewed changes

Comment thread src/python/py/models/builders/base.py

Comment thread src/python/py/models/builders/base.py

Comment thread src/python/py/models/builders/base.py Outdated

Comment thread src/python/py/models/builders/lfm2.py Outdated

Yuri Khrustalev (ykhrustalev) added 2 commits September 17, 2026 14:08

fix: reject unsupported QMoE weight dimensions

93aa6da

Tianlei Wu (tianleiwu) previously approved these changes Sep 17, 2026

View reviewed changes

Tianlei Wu (tianleiwu) enabled auto-merge (squash) September 17, 2026 19:08

Yuri Khrustalev (ykhrustalev) mentioned this pull request Sep 17, 2026

feat: add LFM2-VL / LFM2.5-VL vision support #2571

Merged

auto-merge was automatically disabled September 18, 2026 04:28
Head branch was pushed to by a user without write access

Yuri Khrustalev (ykhrustalev) dismissed Tianlei Wu (tianleiwu)’s stale review via f1c22ae September 18, 2026 04:28

Yuri Khrustalev (ykhrustalev) mentioned this pull request Sep 18, 2026

feat: add LFM2.5-8B-A1B recipes microsoft/olive-recipes#626

Open

Tianlei Wu (tianleiwu) approved these changes Sep 18, 2026

View reviewed changes

Tianlei Wu (tianleiwu) enabled auto-merge (squash) September 18, 2026 18:31

Tianlei Wu (tianleiwu) merged commit 1b1b171 into microsoft:main Sep 18, 2026
16 of 17 checks passed

Yuri Khrustalev (ykhrustalev) mentioned this pull request Sep 21, 2026

feat: add LFM2-Audio / LFM2.5-Audio speech input and output support #2601

Merged

Yuri Khrustalev (ykhrustalev) mentioned this pull request Oct 15, 2026

docs: refresh the LFM2.5-8B-A1B setup for released genai and ONNX Runtime microsoft/olive-recipes#667

Draft

Conversation

Yuri Khrustalev (ykhrustalev) commented Sep 16, 2026

Uh oh!

azure-pipelines Bot commented Sep 16, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Suppressed comments (2)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Yuri Khrustalev (ykhrustalev) commented Sep 16, 2026

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Choose a reason for hiding this comment

Review Summary

Verification performed on 5ee16167

Top concerns (details inline)

Findings outside the diff

Minor notes

Stale threads

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Yuri Khrustalev (ykhrustalev) commented Sep 16, 2026

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Tianlei Wu (tianleiwu) left a comment

Choose a reason for hiding this comment

Uh oh!

Tianlei Wu (tianleiwu) commented Sep 18, 2026

Uh oh!

azure-pipelines Bot commented Sep 18, 2026

Uh oh!

Tianlei Wu (tianleiwu) commented Sep 18, 2026

Uh oh!

Tianlei Wu (tianleiwu) commented Sep 18, 2026

Uh oh!

azure-pipelines Bot commented Sep 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Verification performed on `5ee16167`