From 85efdb8be5f05bbe1b0deb81ceec13b7ffd5416e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 24 May 2026 12:21:06 +0000 Subject: [PATCH 1/3] Initial plan From a24a424fa32960d9d8fc4d21988786c685d57962 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 24 May 2026 12:31:26 +0000 Subject: [PATCH 2/3] Import Qwen3.5-MoE support from onnxruntime-genai PR #2146 + fast unit test Agent-Logs-Url: https://github.com/xadupre/mbext/sessions/0d946b1a-b3d9-4d9a-bba8-57a66a71d021 Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> --- modelbuilder/builder.py | 4 + modelbuilder/builders/qwen.py | 213 +++++++++++++++++++++++++- tests/fast/test_random_qwen3_5_moe.py | 141 +++++++++++++++++ 3 files changed, 354 insertions(+), 4 deletions(-) create mode 100644 tests/fast/test_random_qwen3_5_moe.py diff --git a/modelbuilder/builder.py b/modelbuilder/builder.py index 66c1522..afb11bb 100644 --- a/modelbuilder/builder.py +++ b/modelbuilder/builder.py @@ -389,6 +389,10 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid from .builders.qwen import Qwen35CausalLMModel onnx_model = Qwen35CausalLMModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) + elif config.architectures[0] == "Qwen3_5MoeForConditionalGeneration": + from .builders.qwen import Qwen35MoeTextModel + + onnx_model = Qwen35MoeTextModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options) elif config.architectures[0] == "Qwen3VLForConditionalGeneration": text_config = config.text_config for key in text_config: diff --git a/modelbuilder/builders/qwen.py b/modelbuilder/builders/qwen.py index 4eb44a7..a09982c 100644 --- a/modelbuilder/builders/qwen.py +++ b/modelbuilder/builders/qwen.py @@ -1895,6 +1895,11 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options) + # Set the genai-config model_type here (rather than inside + # ``make_genai_config``) so subclasses (e.g. ``Qwen35MoeTextModel``) + # can override it cleanly. + self.model_type = "Qwen3_5ForConditionalGeneration" + # OffsetRMSNorm: Qwen3.5 uses (1 + weight) * RMSNorm(x). # Pre-bake the +1 into the weight initializer so the base class's # SkipSimplifiedLayerNormalization can be used directly. @@ -2802,10 +2807,11 @@ def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir): hf_config.save_pretrained(out_dir) # Temporarily restore the KV cache template keys and adjust attributes - # so the base class generates the right entries. - saved = {"num_layers": self.num_layers, "model_type": self.model_type} + # so the base class generates the right entries. ``self.model_type`` + # is already set in ``__init__`` (and may have been overridden by a + # subclass such as ``Qwen35MoeTextModel``). + saved = {"num_layers": self.num_layers} self.num_layers = len(self.layer_types) - self.model_type = "Qwen3_5ForConditionalGeneration" self.input_names["past_key_values.key"] = "past_key_values.%d.key" self.input_names["past_key_values.value"] = "past_key_values.%d.value" self.output_names["present.key"] = "present.%d.key" @@ -2815,7 +2821,6 @@ def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir): # Restore self.num_layers = saved["num_layers"] - self.model_type = saved["model_type"] del self.input_names["past_key_values.key"] del self.input_names["past_key_values.value"] del self.output_names["present.key"] @@ -2849,3 +2854,203 @@ def load_weights(self, input_path): return Qwen3_5ForCausalLM.from_pretrained( self.model_name_or_path, cache_dir=self.cache_dir, token=self.hf_token, trust_remote_code=self.hf_remote ) + + +class Qwen35MoeTextModel(Qwen35TextModel): + """Qwen3.5 MoE hybrid model builder. + + Extends :class:`Qwen35TextModel` with Mixture-of-Experts MLP layers. + Each decoder layer replaces the dense MLP with: + + - A router that selects top-k experts from ``num_experts`` candidates. + - Packed routed expert weights (``gate_up_proj`` and ``down_proj``). + - A shared expert (always-active) with its own sigmoid gating signal. + + The attention side (GatedDeltaNet linear + gated full attention) is + inherited unchanged from the parent class. + + Imported from `microsoft/onnxruntime-genai PR #2146 + `_. + """ + + def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): + # Map Qwen3.5-MoE config attributes to what the base class expects. + if hasattr(config, "text_config"): + tc = config.text_config + # Base class reads ``num_local_experts``; MoE config uses ``num_experts``. + if hasattr(tc, "num_experts") and not hasattr(tc, "num_local_experts"): + tc.num_local_experts = tc.num_experts + # Base class reads ``intermediate_size``; MoE has ``moe_intermediate_size``. + if not hasattr(tc, "intermediate_size") and hasattr(tc, "moe_intermediate_size"): + tc.intermediate_size = tc.moe_intermediate_size + + super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options) + + self.model_type = "Qwen3_5_MoeForConditionalGeneration" + + # MoE attributes specific to Qwen3.5-MoE. + self.moe_attrs["activation_type"] = "swiglu" + self.moe_attrs["swiglu_fusion"] = 1 + self.moe_attrs["normalize_routing_weights"] = True + + self.moe_intermediate_size = getattr(config, "moe_intermediate_size", 512) + self.shared_expert_intermediate_size = getattr(config, "shared_expert_intermediate_size", self.moe_intermediate_size) + + # MoE layers use MoE/QMoE ops instead of individual MatMul nodes, + # so remove any /mlp/ MatMul overrides that don't apply. + algo_config = self.quant_attrs["int4"].get("algo_config") + if algo_config is not None and hasattr(algo_config, "customized_weight_config"): + keys_to_remove = [k for k in algo_config.customized_weight_config if "/mlp/" in k] + for k in keys_to_remove: + del algo_config.customized_weight_config[k] + + def load_weights(self, input_path): + from transformers import Qwen3_5MoeForConditionalGeneration + + print("Loading Qwen3_5MoeForConditionalGeneration model...") + return Qwen3_5MoeForConditionalGeneration.from_pretrained( + self.model_name_or_path, cache_dir=self.cache_dir, token=self.hf_token, trust_remote_code=self.hf_remote + ) + + def make_layer(self, layer_id, layer): + """Override to use MoE instead of dense MLP.""" + attn_module = layer.linear_attn if self.layer_types[layer_id] == "linear_attention" else layer.self_attn + self.make_layernorm( + layer_id, + layer.input_layernorm, + skip=not self.layernorm_attrs["first_layernorm"], + simple=self.layernorm_attrs["simple"], + location="input", + ) + self.make_attention(layer_id, attn_module, root_input=self.layernorm_attrs["output_0"]) + self.make_layernorm( + layer_id, layer.post_attention_layernorm, skip=True, simple=self.layernorm_attrs["simple"], location="post_attention" + ) + self.make_moe(layer_id, layer.mlp, root_input=self.layernorm_attrs["output_0"]) + + self.layernorm_attrs["first_layernorm"] = False + if layer_id == self.num_layers - 1: + self.layernorm_attrs["last_layernorm"] = True + + def make_moe(self, layer_id, mlp, root_input): + """Build MoE + shared expert subgraph for one decoder layer.""" + basename = f"/model/layers.{layer_id}/moe" + op_type = self.moe_attrs["op_type"] + moe_weight_type = f"{'q' if op_type == 'QMoE' else ''}weight" + + # --- Router (bias-free gate) --- + router_basename = f"{basename}/router/MatMul" + router_matmul_name = self.make_matmul(mlp.gate, router_basename, root_input) + router_reshape_name = f"{basename}/router/Reshape" + self.make_reshape( + router_reshape_name, + [f"{router_matmul_name}/output_0", f"/model/constants/INT64/{[-1, self.moe_attrs['num_experts']]}"], + dtype=self.io_dtype, + shape=["batch_size * sequence_length", self.moe_attrs["num_experts"]], + ) + + # --- Routed expert weights --- + gate_up_proj_weight = f"model.layers.{layer_id}.moe.experts.gate_up_proj.{moe_weight_type}" + gate_up_proj_scales = f"model.layers.{layer_id}.moe.experts.gate_up_proj.scales" + gate_up_proj_bias = f"model.layers.{layer_id}.moe.experts.gate_up_proj.bias" + down_proj_weight = f"model.layers.{layer_id}.moe.experts.down_proj.{moe_weight_type}" + down_proj_scales = f"model.layers.{layer_id}.moe.experts.down_proj.scales" + down_proj_bias = f"model.layers.{layer_id}.moe.experts.down_proj.bias" + + # Repack HF concatenated [gate|up] to ORT interleaved [g0,u0,g1,u1,...] for swiglu_fusion=1. + raw_gate_up = mlp.experts.gate_up_proj + half = raw_gate_up.shape[1] // 2 + interleaved = torch.stack([raw_gate_up[:, :half, :], raw_gate_up[:, half:, :]], dim=2).reshape_as(raw_gate_up) + + if op_type == "MoE": + self.make_initializer(interleaved, gate_up_proj_weight, to=self.io_dtype) + self.make_initializer(mlp.experts.down_proj, down_proj_weight, to=self.io_dtype) + else: + gate_up_qw_list, gate_up_sc_list = [], [] + down_qw_list, down_sc_list = [], [] + for i in range(self.moe_attrs["num_experts"]): + qw1, sc1 = self.make_qmoe_weights(interleaved[i]) + gate_up_qw_list.append(qw1) + gate_up_sc_list.append(sc1) + qw2, sc2 = self.make_qmoe_weights(mlp.experts.down_proj[i]) + down_qw_list.append(qw2) + down_sc_list.append(sc2) + self.make_initializer(torch.stack(gate_up_qw_list, dim=0).to(torch.uint8), gate_up_proj_weight) + self.make_initializer(torch.stack(down_qw_list, dim=0).to(torch.uint8), down_proj_weight) + self.make_initializer(torch.stack(gate_up_sc_list, dim=0), gate_up_proj_scales, to=self.io_dtype) + self.make_initializer(torch.stack(down_sc_list, dim=0), down_proj_scales, to=self.io_dtype) + + num_e = self.moe_attrs["num_experts"] + self.make_initializer(torch.zeros(num_e, 2 * self.moe_intermediate_size), gate_up_proj_bias, to=self.io_dtype) + self.make_initializer(torch.zeros(num_e, self.hidden_size), down_proj_bias, to=self.io_dtype) + + # --- MoE/QMoE op --- + moe_name = f"{basename}/{op_type}" + self.make_moe_op( + moe_name, + root_input=root_input, + router_probs=f"{router_reshape_name}/output_0", + weight1=gate_up_proj_weight, + scales1=gate_up_proj_scales if op_type == "QMoE" else "", + bias1=gate_up_proj_bias, + weight2=down_proj_weight, + scales2=down_proj_scales if op_type == "QMoE" else "", + bias2=down_proj_bias, + ) + + # --- Shared expert --- + shared_output = self.make_shared_expert(layer_id, mlp.shared_expert, mlp.shared_expert_gate, root_input) + combine_name = f"{basename}/Add" + self.make_add( + combine_name, + [f"{moe_name}/output_0", shared_output], + dtype=self.io_dtype, + shape=["batch_size", "sequence_length", self.hidden_size], + ) + self.layernorm_attrs["skip_input"] = f"{combine_name}/output_0" + + def make_shared_expert(self, layer_id, shared_expert, shared_expert_gate, root_input): + """Build the shared expert SiLU-MLP with sigmoid gating.""" + basename = f"/model/layers.{layer_id}/shared_expert" + + gate_matmul = self.make_matmul(shared_expert.gate_proj, f"{basename}/gate_proj/MatMul", root_input) + up_matmul = self.make_matmul(shared_expert.up_proj, f"{basename}/up_proj/MatMul", root_input) + + silu_sigmoid_name = f"{basename}/gate_proj/Sigmoid" + self.make_sigmoid( + silu_sigmoid_name, + f"{gate_matmul}/output_0", + self.io_dtype, + shape=["batch_size", "sequence_length", self.shared_expert_intermediate_size], + ) + + silu_mul_name = f"{basename}/gate_proj/Mul" + self.make_mul( + silu_mul_name, + [f"{gate_matmul}/output_0", f"{silu_sigmoid_name}/output_0"], + dtype=self.io_dtype, + shape=["batch_size", "sequence_length", self.shared_expert_intermediate_size], + ) + + gate_up_mul_name = f"{basename}/Mul" + self.make_mul( + gate_up_mul_name, + [f"{silu_mul_name}/output_0", f"{up_matmul}/output_0"], + dtype=self.io_dtype, + shape=["batch_size", "sequence_length", self.shared_expert_intermediate_size], + ) + + down_matmul = self.make_matmul(shared_expert.down_proj, f"{basename}/down_proj/MatMul", f"{gate_up_mul_name}/output_0") + + gate_matmul_name = self.make_matmul(shared_expert_gate, f"{basename}_gate/MatMul", root_input) + gate_sigmoid_name = f"{basename}_gate/Sigmoid" + self.make_sigmoid(gate_sigmoid_name, f"{gate_matmul_name}/output_0", self.io_dtype, shape=["batch_size", "sequence_length", 1]) + + gated_mul_name = f"{basename}/GatedMul" + self.make_mul( + gated_mul_name, + [f"{down_matmul}/output_0", f"{gate_sigmoid_name}/output_0"], + dtype=self.io_dtype, + shape=["batch_size", "sequence_length", self.hidden_size], + ) + return f"{gated_mul_name}/output_0" diff --git a/tests/fast/test_random_qwen3_5_moe.py b/tests/fast/test_random_qwen3_5_moe.py new file mode 100644 index 0000000..6cafe2a --- /dev/null +++ b/tests/fast/test_random_qwen3_5_moe.py @@ -0,0 +1,141 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- +"""Fast unit tests for the Qwen3.5-MoE ONNX builder. + +These tests cover the changes imported from `microsoft/onnxruntime-genai +PR #2146 `_, +which add support for the ``Qwen3_5MoeForConditionalGeneration`` HF +architecture (256 routed experts + 1 shared expert with SwiGLU). +""" + +import os +import unittest + +from modelbuilder.ext_test_case import ExtTestCase, hide_stdout, requires_transformers + +QWEN3_5_MOE_MODEL_NAME = "Qwen/Qwen3.5-MoE" + + +def _make_qwen3_5_moe_config(layer_types, num_hidden_layers=None): + """Return a minimal ``Qwen3_5MoeConfig`` suitable for offline unit tests. + + Parameters + ---------- + layer_types: + List of layer type strings, e.g. ``["full_attention", "linear_attention"]``. + The number of layers is inferred from this list unless + ``num_hidden_layers`` is provided. + num_hidden_layers: + Explicit override; defaults to ``len(layer_types)``. + """ + from transformers.models.qwen3_5_moe.configuration_qwen3_5_moe import Qwen3_5MoeConfig, Qwen3_5MoeTextConfig + + if num_hidden_layers is None: + num_hidden_layers = len(layer_types) + + # partial_rotary_factor=0.25, head_dim=64 -> rdim=16, rdim_half=8. + # mrope_section=[2, 3, 3]: height positions at stride-1 offsets within rdim_half. + rope_cfg = {"type": "mrope", "rope_type": "default", "mrope_section": [2, 3, 3], "rope_theta": 10000.0, "partial_rotary_factor": 0.25} + + text_config = Qwen3_5MoeTextConfig( + hidden_size=128, + num_hidden_layers=num_hidden_layers, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=64, + max_position_embeddings=256, + vocab_size=32000, + rms_norm_eps=1e-6, + layer_types=layer_types, + linear_num_key_heads=2, + linear_num_value_heads=2, + linear_key_head_dim=16, + linear_value_head_dim=16, + linear_conv_kernel_dim=4, + # MoE-specific settings (kept small for CI). + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=64, + shared_expert_intermediate_size=64, + ) + text_config.rope_scaling = rope_cfg + text_config.rope_parameters = rope_cfg + + config = Qwen3_5MoeConfig(text_config=text_config, bos_token_id=1, eos_token_id=2) + config.architectures = ["Qwen3_5MoeForConditionalGeneration"] + return config + + +class TestRandomQwen3_5Moe(ExtTestCase): + def _build_model(self, config, precision, provider): + """Create a random-weight HF MoE model and build its ONNX export. + + Returns the output directory containing ``model.onnx``. + """ + import torch + from transformers import AutoModelForImageTextToText + + from modelbuilder.builder import create_model + + basename = f"test_qwen3_5_moe_{precision}_{provider}_" + "_".join(config.text_config.layer_types) + model_dir_full = self.get_model_dir(basename) + output_dir, cache_dir = self.get_dirs(basename) + + torch.manual_seed(42) + model = AutoModelForImageTextToText.from_config(config) + model.eval() + model.save_pretrained(model_dir_full) + + tokenizer = self.make_word_level_tokenizer() + tokenizer.save_pretrained(model_dir_full) + + create_model( + model_name=QWEN3_5_MOE_MODEL_NAME, + input_path=model_dir_full, + output_dir=output_dir, + precision=precision, + execution_provider=provider, + cache_dir=cache_dir, + ) + return output_dir + + @requires_transformers("5") + @hide_stdout() + def test_qwen3_5_moe_fp32_cpu_full_attention_build(self): + """Build a Qwen3.5-MoE decoder with only ``full_attention`` layers. + + Verifies that ``Qwen35MoeTextModel`` registers correctly, emits MoE + ops in the ONNX graph, and produces a model with the expected + ``model_type`` in ``genai_config.json``. + """ + import json + + import onnx + + config = _make_qwen3_5_moe_config(["full_attention", "full_attention"]) + output_dir = self._build_model(config, "fp32", "cpu") + + text_onnx_path = os.path.join(output_dir, "model.onnx") + self.assertExists(text_onnx_path) + + onnx_model = onnx.load(text_onnx_path) + self.assertIsNotNone(onnx_model) + + # Confirm an MoE op (MoE or QMoE) appears in the graph along with the + # shared-expert sigmoid gate (com.microsoft Sigmoid via make_sigmoid). + op_types = {node.op_type for node in onnx_model.graph.node} + self.assertTrue(("MoE" in op_types) or ("QMoE" in op_types), f"Expected an MoE/QMoE op in the graph, found: {sorted(op_types)}") + + # genai_config.json should expose the MoE-specific model_type. + genai_config_path = os.path.join(output_dir, "genai_config.json") + self.assertExists(genai_config_path) + with open(genai_config_path, encoding="utf-8") as f: + genai_config = json.load(f) + self.assertEqual(genai_config["model"]["type"], "qwen3_5_moe") + + +if __name__ == "__main__": + unittest.main(verbosity=2) From 377b8b68c1ffcd7241de9e8a06ab66695d2f5594 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 24 May 2026 12:42:46 +0000 Subject: [PATCH 3/3] Add ORT-vs-PT discrepancy + first-token check for Qwen3.5-MoE Agent-Logs-Url: https://github.com/xadupre/mbext/sessions/b964cd1f-c911-4ab1-b273-972fb0386e8f Co-authored-by: xadupre <22452781+xadupre@users.noreply.github.com> --- tests/fast/test_random_qwen3_5_moe.py | 130 ++++++++++++++++++++++++++ 1 file changed, 130 insertions(+) diff --git a/tests/fast/test_random_qwen3_5_moe.py b/tests/fast/test_random_qwen3_5_moe.py index 6cafe2a..d616eba 100644 --- a/tests/fast/test_random_qwen3_5_moe.py +++ b/tests/fast/test_random_qwen3_5_moe.py @@ -14,6 +14,8 @@ import os import unittest +import numpy as np + from modelbuilder.ext_test_case import ExtTestCase, hide_stdout, requires_transformers QWEN3_5_MOE_MODEL_NAME = "Qwen/Qwen3.5-MoE" @@ -136,6 +138,134 @@ def test_qwen3_5_moe_fp32_cpu_full_attention_build(self): genai_config = json.load(f) self.assertEqual(genai_config["model"]["type"], "qwen3_5_moe") + # ------------------------------------------------------------------ # + # Discrepancy: HF PyTorch vs ONNX Runtime CPU # + # ------------------------------------------------------------------ # + + def _prefill_feed(self, model, config, precision, batch_size=1, seq_len=5): + """Return ``(inputs_embeds_pt, position_ids_3d, onnx_feed)`` for a + Qwen3.5-MoE prefill step with all-``full_attention`` layers. + """ + import torch + + text_cfg = config.text_config + np_dtype = self.get_input_np_dtype(precision) + + torch.manual_seed(0) + input_ids = torch.randint(3, text_cfg.vocab_size, (batch_size, seq_len)) + with torch.no_grad(): + inputs_embeds_pt = model.model.language_model.embed_tokens(input_ids) + + pos = np.arange(seq_len, dtype=np.int64) + position_ids_3d = np.stack([pos] * 3, axis=0)[:, None, :] # [3, 1, S] + position_ids_3d = np.broadcast_to(position_ids_3d, (3, batch_size, seq_len)).copy() + + feed = { + "inputs_embeds": inputs_embeds_pt.numpy().astype(np_dtype), + "attention_mask": np.ones((batch_size, seq_len), dtype=np.int64), + "position_ids": position_ids_3d, + } + for i, lt in enumerate(text_cfg.layer_types): + if lt == "full_attention": + feed[f"past_key_values.{i}.key"] = np.zeros( + (batch_size, text_cfg.num_key_value_heads, 0, text_cfg.head_dim), dtype=np_dtype + ) + feed[f"past_key_values.{i}.value"] = np.zeros( + (batch_size, text_cfg.num_key_value_heads, 0, text_cfg.head_dim), dtype=np_dtype + ) + return inputs_embeds_pt, position_ids_3d, feed + + @requires_transformers("5") + @hide_stdout() + def test_qwen3_5_moe_fp32_cpu_discrepancy_full_attention(self): + """Compare ONNX Runtime prefill logits with HF PyTorch forward. + + Builds a tiny random-weight ``Qwen3_5MoeForConditionalGeneration`` + decoder, runs both ``onnxruntime`` CPU and the HF model's + ``language_model`` forward (with ``inputs_embeds`` and the 3-D + mRoPE ``position_ids``), and asserts: + + - The per-element maximum absolute difference of the logits is + below ``1e-3`` (fp32, dominated by the ``com.microsoft:MoE`` + kernel and ``GroupQueryAttention``). + - The greedy first-token prediction (``argmax`` of the last-row + logits) agrees between PyTorch and ONNX Runtime — this is the + "first token difference" that ORT-GenAI greedy generation + relies on. + """ + import torch + from transformers import AutoModelForImageTextToText + + from modelbuilder.builder import create_model + + precision, provider = "fp32", "cpu" + config = _make_qwen3_5_moe_config(["full_attention", "full_attention"]) + + basename = f"test_qwen3_5_moe_disc_{precision}_{provider}" + model_dir_full = self.get_model_dir(basename) + output_dir, cache_dir = self.get_dirs(basename) + + torch.manual_seed(42) + model = AutoModelForImageTextToText.from_config(config) + model.eval() + model.save_pretrained(model_dir_full) + tokenizer = self.make_word_level_tokenizer() + tokenizer.save_pretrained(model_dir_full) + + create_model( + model_name=QWEN3_5_MOE_MODEL_NAME, + input_path=model_dir_full, + output_dir=output_dir, + precision=precision, + execution_provider=provider, + cache_dir=cache_dir, + ) + + onnx_path = os.path.join(output_dir, "model.onnx") + self.assertExists(onnx_path) + + sess = self._check_with_ort(onnx_path, cpu=True) + onnx_input_names = {i.name for i in sess.get_inputs()} + + inputs_embeds_pt, position_ids_3d, feed = self._prefill_feed(model, config, precision) + feed = {k: v for k, v in feed.items() if k in onnx_input_names} + ort_outputs = sess.run(None, feed) + ort_logits = ort_outputs[0] + self.assertEqual(ort_logits.shape, (1, 5, config.text_config.vocab_size)) + + # HF forward using the same inputs_embeds + 3-D position_ids. + with torch.no_grad(): + pt_out = model.model.language_model( + inputs_embeds=inputs_embeds_pt, + position_ids=torch.from_numpy(position_ids_3d), + attention_mask=torch.ones(1, inputs_embeds_pt.shape[1], dtype=torch.long), + use_cache=False, + ) + pt_logits = model.lm_head(pt_out.last_hidden_state).numpy() + + # Discrepancy (uses the same helper as other random-weight tests). + disc = self.get_numpy_discrepancy(pt_logits, ort_logits) + self.log_results( + { + "step": "prefill", + "precision": precision, + "model_id": QWEN3_5_MOE_MODEL_NAME, + "experiment": "forward", + "provider": provider, + "test": basename, + "input_type": "text", + "kind": "fast", + **disc, + } + ) + np.testing.assert_allclose(pt_logits, ort_logits, atol=1e-3, rtol=1e-3) + + # First-token agreement: the greedy next token (argmax of the last + # row of logits) is what ORT-GenAI's greedy generation step consumes. + pt_first = int(np.argmax(pt_logits[0, -1, :])) + ort_first = int(np.argmax(ort_logits[0, -1, :])) + self.assertEqual(pt_first, ort_first) + if __name__ == "__main__": unittest.main(verbosity=2)