From aa0fd353be117f07956476f4757cb96fa150b209 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Fri, 24 Apr 2026 23:40:47 +0000 Subject: [PATCH 1/4] Add k_quant_linear mixed-precision quantization for hybrid attention models --- src/python/py/models/builders/qwen.py | 71 ++++++++++++++------------- 1 file changed, 36 insertions(+), 35 deletions(-) diff --git a/src/python/py/models/builders/qwen.py b/src/python/py/models/builders/qwen.py index 0e228bb000..f99f9e72f8 100644 --- a/src/python/py/models/builders/qwen.py +++ b/src/python/py/models/builders/qwen.py @@ -947,6 +947,20 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): if "partial_rotary_factor" in config.rope_scaling: config.partial_rotary_factor = config.rope_scaling["partial_rotary_factor"] + # Parse layer types before super().__init__() because + # make_int4_algo_config() is called from the base class init + # and needs self.layer_types to identify linear attention layers. + num_layers = getattr(getattr(config, "text_config", config), "num_hidden_layers", 0) + if hasattr(config, "layer_types") and config.layer_types is not None: + self.layer_types = list(config.layer_types) + elif hasattr(config, "full_attention_interval") and config.full_attention_interval is not None: + interval = config.full_attention_interval + self.layer_types = [ + "full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(num_layers) + ] + else: + self.layer_types = ["full_attention"] * num_layers + super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options) # OffsetRMSNorm: Qwen3.5 uses (1 + weight) * RMSNorm(x). @@ -984,17 +998,6 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): # Pre-compute cos/sin cache tables and interleaving masks for mRoPE self._make_rotary_caches() - # Parse layer types - if hasattr(config, "layer_types") and config.layer_types is not None: - self.layer_types = list(config.layer_types) - elif hasattr(config, "full_attention_interval") and config.full_attention_interval is not None: - interval = config.full_attention_interval - self.layer_types = [ - "full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(self.num_layers) - ] - else: - self.layer_types = ["full_attention"] * self.num_layers - # Store linear attention config self.linear_key_head_dim = getattr(config, "linear_key_head_dim", 128) self.linear_value_head_dim = getattr(config, "linear_value_head_dim", 128) @@ -1013,33 +1016,31 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): # Disable fused RoPE in attention op - we apply mRoPE manually self.attention_attrs["use_rope_in_attn"] = False - # Mixed-precision quantization for linear attention layers. - # Baseline: whole model INT4. Override linear attention layer nodes - # to INT8 for better accuracy with modest size increase. - # - # Linear attention recurrence accumulates errors across the full sequence, - # unlike softmax attention which normalizes per-step. - int8_nodes = {} - for i, lt in enumerate(self.layer_types): - if lt == "linear_attention": - # All linear attention projections: INT8 - for proj in ("in_proj_a", "in_proj_b", "in_proj_qkv", "in_proj_z", "out_proj"): - int8_nodes[f"/model/layers.{i}/linear_attn/{proj}/MatMul"] = {"bits": 8} - # MLP projections in linear attention layers: INT8 - for proj in ("gate_proj", "up_proj", "down_proj"): - int8_nodes[f"/model/layers.{i}/mlp/{proj}/MatMul"] = {"bits": 8} - - if int8_nodes: - algo_config = self.quant_attrs["int4"].get("algo_config") - if algo_config is not None and hasattr(algo_config, "customized_weight_config"): - algo_config.customized_weight_config.update(int8_nodes) - else: - algo_config = RTNWeightOnlyQuantConfig(customized_weight_config=int8_nodes) - self.quant_attrs["int4"]["algo_config"] = algo_config - # Replace standard KV cache I/O with hybrid cache I/O self._setup_hybrid_cache_io() + def make_int4_algo_config(self, quant_method: str): + """Extend base int4_algo_config with ``k_quant_linear``. + + Promotes all linear attention projections and their MLPs from INT4 to + INT8, since linear attention recurrence accumulates quantization errors + across the full sequence (no softmax normalization). + """ + if quant_method != "k_quant_linear": + return super().make_int4_algo_config(quant_method) + + int8_nodes = { + f"/model/layers.{i}/{section}/{proj}/MatMul": {"bits": 8} + for i, lt in enumerate(self.layer_types) + if lt == "linear_attention" + for section, projs in ( + ("linear_attn", ("in_proj_a", "in_proj_b", "in_proj_qkv", "in_proj_z", "out_proj")), + ("mlp", ("gate_proj", "up_proj", "down_proj")), + ) + for proj in projs + } + return RTNWeightOnlyQuantConfig(customized_weight_config=int8_nodes) + def _setup_hybrid_cache_io(self): """Set up hybrid cache I/O: KV cache for attention layers, conv_state + recurrent_state for linear attention layers.""" From 5adbe2e3fd5a3fdb008be8c8f45573055ded948e Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Sat, 25 Apr 2026 02:27:56 +0000 Subject: [PATCH 2/4] address PR comments --- src/python/py/models/builders/base.py | 2 +- src/python/py/models/builders/qwen.py | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index 6fdea95334..709c80561c 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -387,7 +387,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): if hasattr(config, "tie_word_embeddings") and config.tie_word_embeddings is not None else False, ) - self.int8_lm_head = extra_options.get("int4_algo_config", "default") in {"k_quant_mixed", "k_quant_last", "rtn_last"} + self.int8_lm_head = extra_options.get("int4_algo_config", "default") in {"k_quant_mixed", "k_quant_last", "k_quant_linear", "rtn_last"} # shared_embeddings conflicts with exclude_embeds and exclude_lm_head if self.shared_embeddings and (self.exclude_embeds or self.exclude_lm_head): diff --git a/src/python/py/models/builders/qwen.py b/src/python/py/models/builders/qwen.py index f99f9e72f8..fdc795731d 100644 --- a/src/python/py/models/builders/qwen.py +++ b/src/python/py/models/builders/qwen.py @@ -8,7 +8,6 @@ import numpy as np import onnx_ir as ir import torch -from onnxruntime.quantization.matmul_nbits_quantizer import RTNWeightOnlyQuantConfig from transformers import ( AutoConfig, Qwen2_5_VLForConditionalGeneration, @@ -950,7 +949,9 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): # Parse layer types before super().__init__() because # make_int4_algo_config() is called from the base class init # and needs self.layer_types to identify linear attention layers. - num_layers = getattr(getattr(config, "text_config", config), "num_hidden_layers", 0) + # Mirror base class logic: prefer extra_options["num_hidden_layers"] when present. + text_config = getattr(config, "text_config", config) + num_layers = extra_options.get("num_hidden_layers", getattr(text_config, "num_hidden_layers", 0)) if hasattr(config, "layer_types") and config.layer_types is not None: self.layer_types = list(config.layer_types) elif hasattr(config, "full_attention_interval") and config.full_attention_interval is not None: @@ -1039,7 +1040,10 @@ def make_int4_algo_config(self, quant_method: str): ) for proj in projs } - return RTNWeightOnlyQuantConfig(customized_weight_config=int8_nodes) + int4_algo_config = super().make_int4_algo_config("k_quant") + existing_weight_config = getattr(int4_algo_config, "customized_weight_config", None) or {} + int4_algo_config.customized_weight_config = {**existing_weight_config, **int8_nodes} + return int4_algo_config def _setup_hybrid_cache_io(self): """Set up hybrid cache I/O: KV cache for attention layers, From 37d619cbd0e05cb024e36a2a964ada3c2a37827e Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 28 Apr 2026 06:51:38 +0000 Subject: [PATCH 3/4] Update base.py with k_quant_linear --- src/python/py/models/builders/base.py | 13 ++++++++++++- src/python/py/models/builders/qwen.py | 25 ------------------------- 2 files changed, 12 insertions(+), 26 deletions(-) diff --git a/src/python/py/models/builders/base.py b/src/python/py/models/builders/base.py index 709c80561c..6f457e8731 100644 --- a/src/python/py/models/builders/base.py +++ b/src/python/py/models/builders/base.py @@ -709,7 +709,7 @@ def make_int4_algo_config(self, quant_method: str): customized_weight_config["/lm_head/MatMul"] = {"bits": 8} int4_algo_config = RTNWeightOnlyQuantConfig(customized_weight_config=customized_weight_config) - elif quant_method in {"k_quant", "k_quant_mixed", "k_quant_last"}: + elif quant_method in {"k_quant", "k_quant_mixed", "k_quant_last", "k_quant_linear"}: if quant_method != "k_quant": customized_weight_config["/lm_head/MatMul"] = {"bits": 8} @@ -729,6 +729,17 @@ def make_int4_algo_config(self, quant_method: str): customized_weight_config["/model/layers." + str(i) + "/attn/v_proj/MatMul"] = {"bits": 8} customized_weight_config["/model/layers." + str(i) + "/mlp/down_proj/MatMul"] = {"bits": 8} + if quant_method == "k_quant_linear" and hasattr(self, "layer_types"): + # Promote linear attention projections and their MLPs to INT8. + # Linear attention recurrence accumulates quantization errors across + # the full sequence (no softmax normalization). + for i, lt in enumerate(self.layer_types): + if lt == "linear_attention": + for proj in ("in_proj_a", "in_proj_b", "in_proj_qkv", "in_proj_z", "out_proj"): + customized_weight_config[f"/model/layers.{i}/linear_attn/{proj}/MatMul"] = {"bits": 8} + for proj in ("gate_proj", "up_proj", "down_proj"): + customized_weight_config[f"/model/layers.{i}/mlp/{proj}/MatMul"] = {"bits": 8} + customized_weight_config["/lm_head/MatMul"] = {"bits": 8} int4_algo_config = KQuantWeightOnlyQuantConfig(customized_weight_config=customized_weight_config) diff --git a/src/python/py/models/builders/qwen.py b/src/python/py/models/builders/qwen.py index fdc795731d..4e483afe57 100644 --- a/src/python/py/models/builders/qwen.py +++ b/src/python/py/models/builders/qwen.py @@ -1020,31 +1020,6 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options): # Replace standard KV cache I/O with hybrid cache I/O self._setup_hybrid_cache_io() - def make_int4_algo_config(self, quant_method: str): - """Extend base int4_algo_config with ``k_quant_linear``. - - Promotes all linear attention projections and their MLPs from INT4 to - INT8, since linear attention recurrence accumulates quantization errors - across the full sequence (no softmax normalization). - """ - if quant_method != "k_quant_linear": - return super().make_int4_algo_config(quant_method) - - int8_nodes = { - f"/model/layers.{i}/{section}/{proj}/MatMul": {"bits": 8} - for i, lt in enumerate(self.layer_types) - if lt == "linear_attention" - for section, projs in ( - ("linear_attn", ("in_proj_a", "in_proj_b", "in_proj_qkv", "in_proj_z", "out_proj")), - ("mlp", ("gate_proj", "up_proj", "down_proj")), - ) - for proj in projs - } - int4_algo_config = super().make_int4_algo_config("k_quant") - existing_weight_config = getattr(int4_algo_config, "customized_weight_config", None) or {} - int4_algo_config.customized_weight_config = {**existing_weight_config, **int8_nodes} - return int4_algo_config - def _setup_hybrid_cache_io(self): """Set up hybrid cache I/O: KV cache for attention layers, conv_state + recurrent_state for linear attention layers.""" From 09acb576112f481d65741229a5047af60083e806 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 28 Apr 2026 18:43:54 +0000 Subject: [PATCH 4/4] Add k_quant_linear in supported options --- src/python/py/models/builder.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py index 0cc75acdd3..52c4726faf 100644 --- a/src/python/py/models/builder.py +++ b/src/python/py/models/builder.py @@ -401,13 +401,14 @@ def get_args(): Use this option when you want to exclude certain nodes from being quantized. Separate the node names with a ',' when passing them here (e.g. int4_nodes_to_exclude=/lm_head/MatMul,/model/embed_tokens/Gather) int4_algo_config = Method for int4 quantization. Default is 'default'. - Currently supported options are: 'default', 'rtn', 'rtn_last', 'k_quant', 'k_quant_mixed', 'k_quant_last'. + Currently supported options are: 'default', 'rtn', 'rtn_last', 'k_quant', 'k_quant_mixed', 'k_quant_last', 'k_quant_linear'. default = algo_config passed to MatMulNBitsQuantizer is None. Quantizer uses default RTN algorithm. All MatMuls are quantized as int4.(different node naming conventions to `rtn`) rtn = RTN algorithm for int4 quantization. rtn_last = RTN algorithm where only the last MatMul (/lm_head/MatMul) is quantized as int8. Other MatMuls are quantized as int4. k_quant = k_quant algorithm for int4 quantization. k_quant_mixed = k_quant algorithm with mixed precision (int4 + int8). k_quant_last = k_quant algorithm where only the last MatMul (/lm_head/MatMul) is quantized as int8. Other MatMuls are quantized as int4. + k_quant_linear = k_quant algorithm with linear attention layer projections and MLPs promoted to int8 (for hybrid attention models like Qwen3.5). shared_embeddings = Enable weight sharing between embedding and LM head layers. Default is false. Use this option to share weights and reduce model size by eliminating duplicate weights. For quantized models (INT4/UINT4): Shares quantized weights using GatherBlockQuantized. Only works with rtn and k_quant algorithms, and cannot be used if LM head is excluded.