From 84f9cb7df999d6b0cb220bc67d92b5e8172effd4 Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Thu, 7 May 2026 10:23:54 +0800 Subject: [PATCH 1/2] [builder] Allow packed QKV MatMul under QK-Norm via post-MatMul Split MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously `use_packed_matmul` was disabled whenever q_norm or k_norm was set, so Qwen3 (and any other QK-Norm architecture) emitted three separate q_proj / k_proj / v_proj MatMulNBits nodes per layer. Allow packed MatMul in this case and insert a single ONNX Split node right after the packed `qkv_proj/MatMul` to recover Q/K/V tensors that feed the existing q_norm/k_norm path. This keeps the per-head SimplifiedLayerNormalization semantics unchanged (math equivalent, quantization unchanged) while reducing 3 MatMulNBits per layer to 1. A single `Split` is preferred over 3 `Slice` nodes because Split reads the packed output once and writes 3 outputs in a single dispatch, avoiding re-reading the same packed tensor 3x each decode step. The packed-bias branch is also gated off when QK-Norm forces unpack, to avoid mismatched shape on a packed Add over a sliced Q tensor. Verified on Qwen3-1.7B (int4, accuracy_level=4) — generated text is byte-identical to the unpacked baseline, and on RTX 5080 (WebGPU EP) gen TPS improves +5.5% (121.6 -> 128.3) with no prefill regression. --- src/python/py/models/builders/base.py | 40 +++++++++++++++++++++++++-- 1 file changed, 37 insertions(+), 3 deletions(-) diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index d79d1531ed..d5db2b5d39 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -530,8 +530,6 @@ def make_attention_init(self): self.attention_attrs["use_packed_matmul"] = ( self.ep not in ["dml"] and not self.matmul_attrs["use_lora"] - and not self.attention_attrs["q_norm"] - and not self.attention_attrs["k_norm"] and not self.extra_options.get("disable_qkv_fusion", False) ) @@ -1061,6 +1059,13 @@ def make_slice(self, name, inputs, dtype, shape): self.make_node("Slice", inputs=inputs, outputs=[output], name=name) self.make_value(output, dtype, shape=shape) + def make_split(self, name, inputs, axis, output_shapes, dtype): + outputs = [f"{name}/output_{i}" for i in range(len(output_shapes))] + self.make_node("Split", inputs=inputs, outputs=outputs, name=name, axis=axis) + for out, shape in zip(outputs, output_shapes): + self.make_value(out, dtype, shape=shape) + return outputs + def make_mul(self, name, inputs, dtype, shape): output = f"{name}/output_0" self.make_node("Mul", inputs=inputs, outputs=[output], name=name) @@ -2952,6 +2957,30 @@ def make_attention_input_proj(self, layer_id, attention, root_input, **kwargs): attention.q_proj, attention.k_proj, attention.v_proj, qkv_matmul_basename, root_input ) self.attention_attrs["q_path"] = f"{qkv_matmul_name}/output_0" + + # When q_norm/k_norm are present, the packed-QKV path inside GQA cannot be used + # (norm runs per-head before attention). Split here so downstream sees Q/K/V separately. + if self.attention_attrs["q_norm"] and self.attention_attrs["k_norm"]: + q_size = self.num_attn_heads * self.head_size + kv_size = self.num_kv_heads * self.head_size + split_name = f"/model/layers.{layer_id}/attn/qkv_proj/Split" + split_outputs = self.make_split( + split_name, + inputs=[ + f"{qkv_matmul_name}/output_0", + f"/model/constants/INT64/[{q_size}, {kv_size}, {kv_size}]", + ], + axis=-1, + output_shapes=[ + ["batch_size", "sequence_length", q_size], + ["batch_size", "sequence_length", kv_size], + ["batch_size", "sequence_length", kv_size], + ], + dtype=self.io_dtype, + ) + self.attention_attrs["q_path"] = split_outputs[0] + self.attention_attrs["k_path"] = split_outputs[1] + self.attention_attrs["v_path"] = split_outputs[2] else: q_matmul_basename = f"/model/layers.{layer_id}/attn/q_proj/MatMul" q_matmul_name = self.make_matmul(attention.q_proj, q_matmul_basename, root_input) @@ -2980,7 +3009,12 @@ def make_attention_input_proj(self, layer_id, attention, root_input, **kwargs): else: # Make Add nodes (if bias exists) - if self.attention_attrs["use_packed_matmul"] and qkv_dtype_equal and any_bias_exists: + if ( + self.attention_attrs["use_packed_matmul"] + and qkv_dtype_equal + and any_bias_exists + and not (self.attention_attrs["q_norm"] and self.attention_attrs["k_norm"]) + ): # Combine 3 Adds into 1 packed Add qkv_add_name = f"/model/layers.{layer_id}/attn/qkv_proj/Add" self.make_packed_add( From d35c5ef2b51d005423dff61d8fef885046e0bfd1 Mon Sep 17 00:00:00 2001 From: Xiaofei Han Date: Tue, 19 May 2026 16:34:23 +0800 Subject: [PATCH 2/2] resolve comments --- src/python/py/models/builders/base.py | 128 +++++++++++++++++--------- 1 file changed, 83 insertions(+), 45 deletions(-) diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index d5db2b5d39..16871d93fc 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -518,6 +518,9 @@ def is_packed_attn_supported(self) -> bool: return (self.ep, self.io_dtype) in valid_packed_attn_configurations def make_attention_init(self): + self.q_size = self.num_attn_heads * self.head_size + self.kv_size = self.num_kv_heads * self.head_size + if self.is_gqa_supported(): # Change model settings for GroupQueryAttention self.attention_attrs["op_type"] = "GroupQueryAttention" @@ -1059,12 +1062,10 @@ def make_slice(self, name, inputs, dtype, shape): self.make_node("Slice", inputs=inputs, outputs=[output], name=name) self.make_value(output, dtype, shape=shape) - def make_split(self, name, inputs, axis, output_shapes, dtype): - outputs = [f"{name}/output_{i}" for i in range(len(output_shapes))] + def make_split(self, name, inputs, outputs, dtypes, shapes, axis=-1): self.make_node("Split", inputs=inputs, outputs=outputs, name=name, axis=axis) - for out, shape in zip(outputs, output_shapes): - self.make_value(out, dtype, shape=shape) - return outputs + for out, dt, shape in zip(outputs, dtypes, shapes): + self.make_value(out, dt, shape=shape) def make_mul(self, name, inputs, dtype, shape): output = f"{name}/output_0" @@ -2373,13 +2374,13 @@ def make_qk_norm(self, layer_id, attention): q_reshape_2_name = f"/model/layers.{layer_id}/attn/q_norm/Reshape_2" q_reshape_2_inputs = [ q_layernorm_output, - f"/model/constants/INT64/[0, -1, {self.num_attn_heads * self.head_size}]", + f"/model/constants/INT64/[0, -1, {self.q_size}]", ] self.make_reshape( q_reshape_2_name, q_reshape_2_inputs, dtype=self.io_dtype, - shape=["batch_size", "sequence_length", self.num_attn_heads * self.head_size], + shape=["batch_size", "sequence_length", self.q_size], ) # Reshape K MatMul from BxSxD to Bx(SxN)xH before LayerNorm @@ -2426,13 +2427,13 @@ def make_qk_norm(self, layer_id, attention): k_reshape_2_name = f"/model/layers.{layer_id}/attn/k_norm/Reshape_2" k_reshape_2_inputs = [ k_layernorm_output, - f"/model/constants/INT64/[0, -1, {self.num_kv_heads * self.head_size}]", + f"/model/constants/INT64/[0, -1, {self.kv_size}]", ] self.make_reshape( k_reshape_2_name, k_reshape_2_inputs, dtype=self.io_dtype, - shape=["batch_size", "sequence_length", self.num_kv_heads * self.head_size], + shape=["batch_size", "sequence_length", self.kv_size], ) # Update q_path and k_path now @@ -2671,13 +2672,13 @@ def make_repeat_kv(self, layer_id, root_input, past_kv, present_kv, **kwargs): reshape_4_name = f"{basename}/Reshape_4" reshape_4_inputs = [ f"{transpose_2_name}/output_0", - f"/model/constants/INT64/[0, 0, {self.num_attn_heads * self.head_size}]", + f"/model/constants/INT64/[0, 0, {self.q_size}]", ] self.make_reshape( reshape_4_name, reshape_4_inputs, dtype=self.io_dtype, - shape=["batch_size", "sequence_length", self.num_attn_heads * self.head_size], + shape=["batch_size", "sequence_length", self.q_size], ) input_to_attention = f"{reshape_4_name}/output_0" @@ -2922,6 +2923,42 @@ def make_attention(self, layer_id, attention, root_input, **kwargs): # O_MatMul # | # O_Add + # + # GroupQueryAttention with packed QKV (no Q/K norm) example: + # + # root_input + # | + # QKV_MatMul seqlens_k total_seq_len past_key past_value + # | | | | | + # QKV_Add (packed) +------------+-----------+----------+ + # | | + # Q_Rotary / K_Rotary (in-attn or external) | + # | | + # GroupQueryAttention----------------------------+ + # | + # O_MatMul + # | + # O_Add + # + # GroupQueryAttention with packed QKV + Q/K norm example: + # + # root_input + # | + # QKV_MatMul + # | + # QKV_Add (packed, only if bias exists) + # | + # Split -> Q, K, V + # / | \ + # Q_Norm K_Norm V seqlens_k total_seq_len past_key past_value + # | | | | | | | + # Q_Rotary K_Rotary V +------------+-----------+----------+ + # \ | / | + # GroupQueryAttention----------------------------+ + # | + # O_MatMul + # | + # O_Add self.make_attention_input_proj(layer_id, attention, root_input, **kwargs) self.make_attention_qk_subgraph(layer_id, attention, root_input, **kwargs) self.make_attention_output_proj(layer_id, attention, root_input, **kwargs) @@ -2957,30 +2994,6 @@ def make_attention_input_proj(self, layer_id, attention, root_input, **kwargs): attention.q_proj, attention.k_proj, attention.v_proj, qkv_matmul_basename, root_input ) self.attention_attrs["q_path"] = f"{qkv_matmul_name}/output_0" - - # When q_norm/k_norm are present, the packed-QKV path inside GQA cannot be used - # (norm runs per-head before attention). Split here so downstream sees Q/K/V separately. - if self.attention_attrs["q_norm"] and self.attention_attrs["k_norm"]: - q_size = self.num_attn_heads * self.head_size - kv_size = self.num_kv_heads * self.head_size - split_name = f"/model/layers.{layer_id}/attn/qkv_proj/Split" - split_outputs = self.make_split( - split_name, - inputs=[ - f"{qkv_matmul_name}/output_0", - f"/model/constants/INT64/[{q_size}, {kv_size}, {kv_size}]", - ], - axis=-1, - output_shapes=[ - ["batch_size", "sequence_length", q_size], - ["batch_size", "sequence_length", kv_size], - ["batch_size", "sequence_length", kv_size], - ], - dtype=self.io_dtype, - ) - self.attention_attrs["q_path"] = split_outputs[0] - self.attention_attrs["k_path"] = split_outputs[1] - self.attention_attrs["v_path"] = split_outputs[2] else: q_matmul_basename = f"/model/layers.{layer_id}/attn/q_proj/MatMul" q_matmul_name = self.make_matmul(attention.q_proj, q_matmul_basename, root_input) @@ -3009,12 +3022,7 @@ def make_attention_input_proj(self, layer_id, attention, root_input, **kwargs): else: # Make Add nodes (if bias exists) - if ( - self.attention_attrs["use_packed_matmul"] - and qkv_dtype_equal - and any_bias_exists - and not (self.attention_attrs["q_norm"] and self.attention_attrs["k_norm"]) - ): + if self.attention_attrs["use_packed_matmul"] and qkv_dtype_equal and any_bias_exists: # Combine 3 Adds into 1 packed Add qkv_add_name = f"/model/layers.{layer_id}/attn/qkv_proj/Add" self.make_packed_add( @@ -3039,6 +3047,36 @@ def make_attention_input_proj(self, layer_id, attention, root_input, **kwargs): self.make_add_bias(attention.v_proj.bias, v_add_name, root_input=self.attention_attrs["v_path"]) self.attention_attrs["v_path"] = f"{v_add_name}/output_0" + # When q_norm/k_norm are present, the packed-QKV path inside GQA cannot be used + # (norm runs per-head before attention). Split here so downstream sees Q/K/V separately. + # Placed after the (optional) packed Add so packed bias fusion is preserved. + if ( + self.attention_attrs["use_packed_matmul"] + and qkv_dtype_equal + and self.attention_attrs["q_norm"] + and self.attention_attrs["k_norm"] + ): + split_name = f"/model/layers.{layer_id}/attn/qkv_proj/Split" + split_outputs = [f"{split_name}/output_{i}" for i in range(3)] + self.make_split( + split_name, + inputs=[ + self.attention_attrs["q_path"], + f"/model/constants/INT64/[{self.q_size}, {self.kv_size}, {self.kv_size}]", + ], + outputs=split_outputs, + dtypes=[self.io_dtype] * 3, + shapes=[ + ["batch_size", "sequence_length", self.q_size], + ["batch_size", "sequence_length", self.kv_size], + ["batch_size", "sequence_length", self.kv_size], + ], + axis=-1, + ) + self.attention_attrs["q_path"] = split_outputs[0] + self.attention_attrs["k_path"] = split_outputs[1] + self.attention_attrs["v_path"] = split_outputs[2] + def make_attention_qk_subgraph(self, layer_id, attention, root_input, **kwargs): # Make Q/K SimplifiedLayerNorm nodes if self.attention_attrs["q_norm"] and self.attention_attrs["k_norm"]: @@ -3145,8 +3183,8 @@ def make_attention_unpacked(self, layer_id, attention, root_input, **kwargs): def make_attention_unpacked_lora(self, layer_id, attention, qkv_linear, root_input, **kwargs): from peft.tuners.lora.layer import LoraLayer - q_size = self.num_attn_heads * self.head_size - kv_size = self.num_kv_heads * self.head_size + q_size = self.q_size + kv_size = self.kv_size # Create Q/K/V base layers q_proj = torch.nn.Linear(in_features=q_size, out_features=q_size) @@ -3209,8 +3247,8 @@ def make_attention_unpacked_lora(self, layer_id, attention, qkv_linear, root_inp attention.v_proj.scaling = qkv_linear.scaling def make_attention_unpacked_regular(self, layer_id, attention, qkv_linear, root_input, **kwargs): - q_size = self.num_attn_heads * self.head_size - kv_size = self.num_kv_heads * self.head_size + q_size = self.q_size + kv_size = self.kv_size attention.q_proj = torch.nn.Linear(in_features=q_size, out_features=q_size) attention.q_proj.weight = torch.nn.Parameter(qkv_linear.weight[:q_size, :], requires_grad=False)