Fix q/k/v quantizer wiring in register_attention_for_kv_quant

CHANGELOG.rst

@@ -52,6 +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. 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.

modelopt/torch/quantization/plugins/attention.py

@@ -52,7 +52,9 @@ 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: 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",)
 
     def is_bmm(node):
@@ -72,6 +74,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 +173,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 +194,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:

tests/unit/torch/quantization/plugins/test_attention_quant.py

@@ -52,6 +52,54 @@ 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
+
+
+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},
@@ -115,6 +163,37 @@ 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)
+
+
+@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()