From ecb54a232f0c6d78a131783a4e52ec954270f3f1 Mon Sep 17 00:00:00 2001 From: Javier de Jesus Date: Wed, 30 Sep 2026 22:38:12 +0000 Subject: [PATCH 1/3] Fix q/k/v quantizer wiring in register_attention_for_kv_quant The AST patcher collected the two matmul nodes with breadth-first ast.walk and patched them by list index, so which quantizer landed on which operand depended on AST nesting depth. A flat layout (scores = q @ k.T, probs = softmax(scores), probs @ v) got k_bmm_quantizer on v, v_bmm_quantizer on k and q_bmm_quantizer on the softmax output. Collect the nodes in source order, children before parents, and swap the two index pairs. The q/k score matmul is expected to be written before the matmul that consumes its result. In the modeling_.py modules of transformers 5.14.1 the generated code changes for 65 of the 72 accepted classes and is byte-identical for the other 7. A class that defines the score matmul in a helper placed after forward can go from correct to wrong, because the order is source position, not call order. Signed-off-by: Javier de Jesus --- CHANGELOG.rst | 1 + .../torch/quantization/plugins/attention.py | 16 +++-- .../plugins/test_attention_quant.py | 61 +++++++++++++++++++ 3 files changed, 73 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ff62b3cc0d4..f564c58d5f4 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -60,6 +60,7 @@ Changelog - Fix Megatron Top-K/Top-P distillation gradients being multiplied by the tensor-parallel size when tensor parallelism is enabled; reported loss values are unchanged. - Fix the Puzzletron memory objective pricing the KV cache with ``hidden_size // num_attention_heads`` instead of the model's ``head_dim``, which misestimated it for models such as Mistral-Small, GPT-OSS and Gemma-3 and for every candidate with a pruned hidden size. Recompute subblock statistics produced with earlier versions. - Fix ONNX INT4 quantization ending the Python process when the INT4 packing extension fails to build, for example without a C++ compiler. It now falls back to the Python implementation. +- Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. 0.48.0 (2026-10-xx) ^^^^^^^^^^^^^^^^^^^ diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py index 2113edea8a7..355d376179a 100644 --- a/modelopt/torch/quantization/plugins/attention.py +++ b/modelopt/torch/quantization/plugins/attention.py @@ -72,6 +72,12 @@ def is_sdpa(node): def is_bin_matmul(node): return isinstance(node, ast.BinOp) and isinstance(node.op, ast.MatMult) + def walk_post_order(node): + # Source order, children first; expects the q/k score matmul before the matmul consuming it + for child in ast.iter_child_nodes(node): + yield from walk_post_order(child) + yield node + def patch(node, quantizer_names, transpose=False): for index, quantizer_name in enumerate(quantizer_names): if quantizer_name is None: @@ -165,7 +171,7 @@ def patch_binop(node, quantizer_names, transpose=False): bmm_nodes = [] sdpa_nodes = [] bin_matmul_nodes = [] - for node in ast.walk(head): + for node in walk_post_order(head): if is_bmm(node): bmm_nodes.append(node) if is_sdpa(node): @@ -186,16 +192,16 @@ def patch_binop(node, quantizer_names, transpose=False): # after transpose, the quantization will be per-token, i.e., # self.k_bmm_quantizer(key_states.transpose(-1, -2).transpose(-1, -2)).transpose(-1, -2) # removing the additional transpose is doable but not trivial - patch(bmm_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) - patch(bmm_nodes[1], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch(bmm_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True) + patch(bmm_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 BMM/Matmul operators with quantizers") if len(bin_matmul_nodes) == 2: patch_binop( - bin_matmul_nodes[1], + bin_matmul_nodes[0], quantizer_names=("q_bmm_quantizer", "k_bmm_quantizer"), transpose=True, ) - patch_binop(bin_matmul_nodes[0], quantizer_names=(None, "v_bmm_quantizer")) + patch_binop(bin_matmul_nodes[1], quantizer_names=(None, "v_bmm_quantizer")) print("Patching 2 @ operators with quantizers") if len(sdpa_nodes) == 1: diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index 22a86971a6d..155d810ab74 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/tests/unit/torch/quantization/plugins/test_attention_quant.py @@ -52,6 +52,42 @@ def forward(self, hidden_states, **kwargs): return F.scaled_dot_product_attention(q, k, v), None +class FlatMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = torch.matmul(q, k.transpose(-2, -1)) + probs = torch.softmax(scores, dim=-1) + return torch.matmul(probs, v) + + +class FlatBMMAttention(nn.Module): + def forward(self, q, k, v): + scores = torch.bmm(q, k.transpose(-2, -1)) + probs = torch.softmax(scores, dim=-1) + return torch.bmm(probs, v) + + +class FlatBinMatmulAttention(nn.Module): + def forward(self, q, k, v): + scores = q @ k.transpose(-2, -1) + probs = torch.softmax(scores, dim=-1) + return probs @ v + + +class NestedMatmulAttention(nn.Module): + def forward(self, q, k, v): + return torch.matmul(torch.softmax(torch.matmul(q, k.transpose(-2, -1)), dim=-1), v) + + +class NestedBMMAttention(nn.Module): + def forward(self, q, k, v): + return torch.bmm(torch.softmax(torch.bmm(q, k.transpose(-2, -1)), dim=-1), v) + + +class NestedBinMatmulAttention(nn.Module): + def forward(self, q, k, v): + return torch.softmax(q @ k.transpose(-2, -1), dim=-1) @ v + + kv_cache_config = { "quant_cfg": [ {"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": 4}, "enable": True}, @@ -115,6 +151,31 @@ def test_kv_quant_hf(model_getter, attn_cls): mtq.unregister(attn_cls) +@pytest.mark.parametrize( + "attn_cls", + [ + FlatMatmulAttention, + FlatBMMAttention, + FlatBinMatmulAttention, + NestedMatmulAttention, + NestedBMMAttention, + NestedBinMatmulAttention, + ], +) +def test_kv_quant_operand_wiring(attn_cls): + """Each q/k/v_bmm_quantizer must calibrate on its own operand, whatever the layout.""" + q, k, v = torch.randn(1, 4, 8), 10 * torch.randn(1, 4, 8), 100 * torch.randn(1, 4, 8) + assert mtq.plugins.register_attention_for_kv_quant(attn_cls) + try: + model = attn_cls() + mtq.quantize(model, kv_cache_config, lambda model: model(q, k, v)) + for name, operand in zip("qkv", (q, k, v)): + quantizer = getattr(model, f"{name}_bmm_quantizer") + assert quantizer.amax == operand.abs().max(), f"{name}_bmm_quantizer saw another tensor" + finally: + mtq.unregister(attn_cls) + + def test_kv_quant_bert(): """Test KV cache quantization on BERT model with decorated attention.""" model_test = get_tiny_bert() From 91a7aaec2a89d400be86cb66d525c414df5d05d5 Mon Sep 17 00:00:00 2001 From: javierdejesusda Date: Mon, 26 Oct 2026 21:26:25 +0100 Subject: [PATCH 2/3] Stop matching baddbmm in register_attention_for_kv_quant patch() wires node.args[0] and node.args[1], but torch.baddbmm(input, batch1, batch2) takes the additive input first. A positional call put q_bmm_quantizer on the bias and k_bmm_quantizer on q, leaving k unquantized (GPTNeoXJapaneseAttention), and a keyword call raised IndexError (BloomAttention). Drop baddbmm from the matched ops so these classes fail the existing count guard and get the warning to register a custom quantized attention class. Correct baddbmm support needs operand offsets and keyword handling in patch(), which is left for separate work. Signed-off-by: javierdejesusda --- CHANGELOG.rst | 2 +- .../torch/quantization/plugins/attention.py | 3 ++- .../plugins/test_attention_quant.py | 18 ++++++++++++++++++ 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 35bd610f44f..192adf0573e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -52,7 +52,7 @@ Changelog - ``mtq.fold_weight`` now folds sequential weight quantizers (e.g. W4A8 INT4 + FP8 weights) into the weights; it used to leave them unfolded and enabled. - Fix ONNX quantization placing a per-tensor Q/DQ pair on a rank-0 (scalar) constant input, which TensorRT's ONNX parser rejects (``convertAxis`` asserts the quantization axis lies in ``[0, nbDims]`` but ``nbDims == 0``). Rank-0 constant operands are now kept unquantized in the int8, fp8, and autotune export paths; the existing single-element Mul/Div/Pad stripping stays gated under ``target_dla``. - Fix PTQ of native Transformers FP8 experts to apply source weight scales during calibration and export. Re-quantize affected checkpoints from the original FP8 source. -- Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. +- Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. Attention classes that compute their scores with ``torch.baddbmm`` (e.g. ``BloomAttention``, ``GPTNeoXJapaneseAttention``) are no longer patched, since their operands were always wired wrong (or raised ``IndexError``), and now get the existing warning to register a custom quantized attention class. - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. - Fix ONNX remote AutoQDQ safety benchmarks returning infinite latency by running the generated engine with ``trtexec_safe`` on the configured target. diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py index 355d376179a..24e4dc4e6b4 100644 --- a/modelopt/torch/quantization/plugins/attention.py +++ b/modelopt/torch/quantization/plugins/attention.py @@ -52,7 +52,8 @@ def register_attention_for_kv_quant(attention_cls: type) -> bool: model_module = inspect.getmodule(attention_cls) head = ast.parse(source_code) - bmm_ops = ("matmul", "bmm", "baddbmm") + # No baddbmm: its matrix operands are args[1] and args[2], and patch() wires args[0] and args[1] + bmm_ops = ("matmul", "bmm") sdpa_ops = ("scaled_dot_product_attention",) def is_bmm(node): diff --git a/tests/unit/torch/quantization/plugins/test_attention_quant.py b/tests/unit/torch/quantization/plugins/test_attention_quant.py index 155d810ab74..ff9e9404ea4 100644 --- a/tests/unit/torch/quantization/plugins/test_attention_quant.py +++ b/tests/unit/torch/quantization/plugins/test_attention_quant.py @@ -88,6 +88,18 @@ def forward(self, q, k, v): return torch.softmax(q @ k.transpose(-2, -1), dim=-1) @ v +class BAddBMMAttention(nn.Module): + def forward(self, q, k, v, bias): + scores = torch.baddbmm(bias, q, k.transpose(-2, -1)) + return torch.bmm(torch.softmax(scores, dim=-1), v) + + +class BAddBMMKeywordAttention(nn.Module): + def forward(self, q, k, v, bias): + scores = torch.baddbmm(input=bias, batch1=q, batch2=k.transpose(-2, -1)) + return torch.bmm(torch.softmax(scores, dim=-1), v) + + kv_cache_config = { "quant_cfg": [ {"quantizer_name": "*[kv]_bmm_quantizer", "cfg": {"num_bits": 4}, "enable": True}, @@ -176,6 +188,12 @@ def test_kv_quant_operand_wiring(attn_cls): mtq.unregister(attn_cls) +@pytest.mark.parametrize("attn_cls", [BAddBMMAttention, BAddBMMKeywordAttention]) +def test_kv_quant_rejects_baddbmm(attn_cls): + """baddbmm takes the additive input first, so its operands cannot be wired by position.""" + assert not mtq.plugins.register_attention_for_kv_quant(attn_cls) + + def test_kv_quant_bert(): """Test KV cache quantization on BERT model with decorated attention.""" model_test = get_tiny_bert() From f8ad6a902d42f601947e34a063deacd4bb5ddac9 Mon Sep 17 00:00:00 2001 From: javierdejesusda Date: Mon, 26 Oct 2026 21:29:32 +0100 Subject: [PATCH 3/3] Clarify why register_attention_for_kv_quant skips baddbmm The comment only described the torch.baddbmm(input, batch1, batch2) form. Bloom calls alibi.baddbmm(batch1=..., batch2=...) with keywords, so an args offset alone would not fix it. Signed-off-by: javierdejesusda --- modelopt/torch/quantization/plugins/attention.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/plugins/attention.py b/modelopt/torch/quantization/plugins/attention.py index 24e4dc4e6b4..e4843a3a2ec 100644 --- a/modelopt/torch/quantization/plugins/attention.py +++ b/modelopt/torch/quantization/plugins/attention.py @@ -52,7 +52,8 @@ def register_attention_for_kv_quant(attention_cls: type) -> bool: model_module = inspect.getmodule(attention_cls) head = ast.parse(source_code) - # No baddbmm: its matrix operands are args[1] and args[2], and patch() wires args[0] and args[1] + # No baddbmm: patch() wires args[0] and args[1], which miss its batch1 and batch2 operands + # (torch.baddbmm passes the additive input first, and Bloom passes them as keywords) bmm_ops = ("matmul", "bmm") sdpa_ops = ("scaled_dot_product_attention",)