[qwen] Drop fp32 LayerNorm casts in Qwen3.5 text model on ORT >= 1.26

modelbuilder/builders/qwen.py

@@ -1905,15 +1905,30 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         # SkipSimplifiedLayerNormalization can be used directly.
         self.layernorm_attrs["add_offset"] = 1
 
-        # HF Qwen3_5RMSNorm always computes in float32 regardless of model
-        # dtype.  Force the builder to cast inputs to fp32 before LayerNorm
-        # and cast back after, matching HF behaviour and preventing precision
-        # loss that compounds across 36+ layers in fp16/bf16 builds.
-        self.layernorm_attrs["cast"]["use_fp32"] = True
-        self.layernorm_attrs["cast"]["root_input"] = True
-        self.layernorm_attrs["cast"]["skip_input"] = True
-        self.layernorm_attrs["cast"]["output_0"] = True
-        self.layernorm_attrs["cast"]["output_3"] = True
+        # Keep RMSNorm IO in the model's native dtype (fp16/bf16) instead of
+        # casting around every LayerNorm.  Casting inputs/outputs to fp32 adds
+        # ~216 Cast nodes in a 24-layer build (108 to-fp32 + 108 to-fp16) and
+        # measurably hurts generation throughput on GPU/iGPU backends, while
+        # output quality remains coherent and structurally identical without
+        # the casts.  See microsoft/onnxruntime-genai#2101.
+        #
+        # This relies on the fp16 ``SkipSimplifiedLayerNormalization`` kernel
+        # that landed in ORT 1.26.  On older ORT releases the native kernel
+        # accumulates in fp16 and loses precision across the 36+ Qwen3.5
+        # layers, so we keep the explicit fp32 cast wrapping in that case.
+        try:
+            import onnxruntime as _ort
+
+            _ort_ver = tuple(int(x) for x in _ort.__version__.split(".")[:2])
+        except (ImportError, ValueError):
+            # Unknown/dev builds are assumed to be recent enough.
+            _ort_ver = (99, 99)
+        if _ort_ver < (1, 26):
+            self.layernorm_attrs["cast"]["use_fp32"] = True
+            self.layernorm_attrs["cast"]["root_input"] = True
+            self.layernorm_attrs["cast"]["skip_input"] = True
+            self.layernorm_attrs["cast"]["output_0"] = True
+            self.layernorm_attrs["cast"]["output_3"] = True
 
         # 3D position_ids for mRoPE: [3, batch_size, sequence_length]
         self.input_shapes["position_ids"] = [3, "batch_size", "sequence_length"]

tests/fast/test_random_qwen3_5.py

@@ -273,6 +273,65 @@ def test_qwen3_5_fp16_cpu_hybrid_build(self):
         self.assertIsNotNone(outputs[0])
         self.assertEqual(outputs[0].shape, (1, 5, 32000))
 
+    @requires_transformers("5")
+    @hide_stdout()
+    def test_qwen3_5_fp16_no_layernorm_fp32_casts(self):
+        """Verify that the fp16 build does not wrap LayerNorm IO with fp32 Casts.
+
+        ``Qwen3_5TextModel`` previously forced ``use_fp32`` on every RMSNorm,
+        which inserted ~216 Cast nodes (to-fp32 + to-fp16) around the LayerNorm
+        ops in a 24-layer build and measurably hurt generation throughput.  See
+        microsoft/onnxruntime-genai#2101.  When the runtime is ORT >= 1.26 the
+        builder keeps LayerNorm IO in the model's native dtype: no Cast-to-fp32
+        should appear on the inputs of the ``SkipSimplifiedLayerNormalization``
+        ops, and no Cast-from-fp32 should appear on their outputs.  On older
+        ORT releases the fp32 cast wrapping is still emitted (the native fp16
+        kernel loses precision across the 36+ Qwen3.5 layers), so this test is
+        skipped there.
+        """
+        import onnx
+        import onnxruntime as ort
+
+        ort_ver = tuple(int(x) for x in ort.__version__.split(".")[:2])
+        if ort_ver < (1, 26):
+            self.skipTest(f"requires onnxruntime >= 1.26, got {ort.__version__}")
+
+        config = _make_qwen3_5_config(["full_attention", "full_attention"])
+        _, output_dir = self._build_and_save_model(config, "fp16", "cpu")
+
+        text_onnx_path = os.path.join(output_dir, "model.onnx")
+        self.assertExists(text_onnx_path)
+
+        onnx_model = onnx.load(text_onnx_path)
+        nodes_by_output = {out: node for node in onnx_model.graph.node for out in node.output}
+
+        FLOAT = onnx.TensorProto.FLOAT
+        offending = []
+        for node in onnx_model.graph.node:
+            if node.op_type != "SkipSimplifiedLayerNormalization":
+                continue
+            # No upstream Cast-to-fp32 feeding the LayerNorm inputs.
+            for inp in node.input:
+                producer = nodes_by_output.get(inp)
+                if producer is None or producer.op_type != "Cast":
+                    continue
+                to_attr = next((a.i for a in producer.attribute if a.name == "to"), None)
+                if to_attr == FLOAT:
+                    offending.append((node.name, "input", producer.name))
+            # No downstream Cast consuming the LayerNorm outputs from fp32.
+            for out in node.output:
+                if not out:
+                    continue
+                for consumer in onnx_model.graph.node:
+                    if consumer.op_type != "Cast" or out not in consumer.input:
+                        continue
+                    # Cast from fp32 means the LayerNorm output was fp32.
+                    val_info = next((v for v in list(onnx_model.graph.value_info) + list(onnx_model.graph.output) if v.name == out), None)
+                    if val_info is not None and val_info.type.tensor_type.elem_type == FLOAT:
+                        offending.append((node.name, "output", consumer.name))
+
+        self.assertEqual(offending, [], f"Unexpected fp32 Cast nodes around LayerNorm: {offending}")
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)