fix(autoquant): score grouped QKV at attention output

modelopt/torch/quantization/algorithms.py

@@ -445,6 +445,7 @@ def attrs(self) -> list[str]:
 
 _LINEAR_ATTN_QKVZ_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_qkv|in_proj_z)$")
 _LINEAR_ATTN_BA_RE = re.compile(r"^(.*?\.linear_attn)\.(?:in_proj_a|in_proj_b)$")
+_ATTN_QKV_RULE = r"^(.*?)\.(q_proj|k_proj|v_proj)$"
 
 
 def _linear_attn_qkvz_group_key(_model, name: str) -> str | None:
@@ -471,7 +472,7 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
     method_name: str | None = None
 
     quant_grouping_rules = [
-        r"^(.*?)\.(q_proj|k_proj|v_proj)$",  # q_proj, k_proj, v_proj for llama like models
+        _ATTN_QKV_RULE,  # q_proj, k_proj, v_proj for llama like models
         # gate_proj, up_proj, down_proj for Qwen3 like MoE models
         r"^(.*?\.mlp\.experts)\.\d+\.(gate_proj|up_proj|down_proj)$",
         r"^(.*?\.mixer\.experts)\.\d+\.(up_proj|down_proj)$",  # NemotronH MoE experts
@@ -1025,6 +1026,7 @@ class AutoQuantizeGradientSearcher(_AutoQuantizeBaseSearcher):
     method_name = "gradient"
 
     score_module_rules = [
+        _ATTN_QKV_RULE,
         # Use MLP layer output for gate_proj, up_proj, down_proj for Qwen3 like MoE models (local and shared experts)
         r"^(.*?\.mlp)\.experts\.\d+\.(gate_proj|up_proj|down_proj)$",
         r"^(.*?\.mixer)\.experts\.\d+\.(up_proj|down_proj)$",  # NemotronH MoE experts
@@ -1105,12 +1107,12 @@ def forward_backward_step(model, data):
     @torch.enable_grad()
     def _estimate_auto_quantize_scores(self, is_param_grad_enabled):
         # TODO: remove the no-quant recipe
-        def auto_quantize_score_estimate_forward(module, input, *args, **kwargs):
+        def auto_quantize_score_estimate_forward(module, *args, **kwargs):
             for hparam in module._hparams_for_scoring:
                 if hparam.is_configurable:
                     hparam.active = QuantRecipe(quant_cfg=None)
 
-            output = module._forward_original(input, *args, **kwargs)
+            output = module._forward_original(*args, **kwargs)
 
             # If gradient checkpointing is enabled, gradient will not be enabled in the global forward pass.
             # With gradient checkpointing, gradients are computed in the local forward pass during backward pass
@@ -1128,7 +1130,7 @@ def auto_quantize_score_estimate_forward(module, input, *args, **kwargs):
                         if recipe == QuantRecipe(quant_cfg=None):
                             continue
                         hparam.active = recipe
-                        output_diff = module._forward_original(input, *args, **kwargs)
+                        output_diff = module._forward_original(*args, **kwargs)
 
                         if isinstance(output_diff, tuple):
                             output_diff = output_diff[0] - output[0]

tests/unit/torch/quantization/test_autoquant.py

@@ -56,6 +56,37 @@ def forward(self, x):
         return x
 
 
+class _ParallelAttention(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.q_proj = torch.nn.Linear(4, 4, bias=False)
+        self.k_proj = torch.nn.Linear(4, 4, bias=False)
+        self.v_proj = torch.nn.Linear(4, 4, bias=False)
+
+        weight = torch.tensor(
+            [
+                [0.11, -0.37, 0.53, 0.89],
+                [-0.29, 0.47, 0.71, -0.97],
+                [0.19, 0.41, -0.67, 0.83],
+                [0.31, -0.59, 0.73, -0.91],
+            ]
+        )
+        for projection in (self.q_proj, self.k_proj, self.v_proj):
+            projection.weight.data.copy_(weight)
+
+    def forward(self, x):
+        return self.q_proj(x) + self.k_proj(x) + self.v_proj(x)
+
+
+class _ParallelAttentionModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = _ParallelAttention()
+
+    def forward(self, x, input_as_keyword=False):
+        return self.attn(x=x) if input_as_keyword else self.attn(x)
+
+
 class TransformerBlock(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -419,6 +450,56 @@ def test_auto_quantize_disabled_layers_no_poison():
     assert QuantRecipe(mtq.INT4_BLOCKWISE_WEIGHT_ONLY_CFG) in hparam.choices
 
 
+@pytest.mark.parametrize("input_as_keyword", [False, True])
+def test_auto_quantize_qkv_score_uses_attention_output(input_as_keyword):
+    model = _ParallelAttentionModel()
+    data = torch.tensor([[[0.13, -0.31, 0.79, -0.43], [0.61, 0.17, -0.23, 0.97]]])
+    output_grad = torch.tensor([[[0.7, -1.1, 0.3, 1.3], [-0.5, 0.9, 1.7, -0.2]]])
+
+    model, search_history = mtq.auto_quantize(
+        model,
+        constraints={"effective_bits": 16.0},
+        quantization_formats=[mtq.INT8_DEFAULT_CFG],
+        data_loader=[data],
+        forward_step=lambda model, batch: model(batch, input_as_keyword=input_as_keyword),
+        loss_func=lambda output, batch: (output * output_grad).sum(),
+        num_calib_steps=1,
+        num_score_steps=1,
+    )
+
+    hparam = model.attn.q_proj.get_hparam("quant_recipe")
+    assert model.attn.k_proj.get_hparam("quant_recipe") is hparam
+    assert model.attn.v_proj.get_hparam("quant_recipe") is hparam
+    assert hparam.score_modules == [model.attn]
+
+    no_quant_recipe = QuantRecipe(None)
+    int8_recipe = QuantRecipe(mtq.INT8_DEFAULT_CFG)
+    hparam.active = no_quant_recipe
+    attention_output = model.attn(data)
+    projection_outputs = [projection(data) for projection in hparam.quant_modules]
+    hparam.active = int8_recipe
+    quantized_attention_output = model.attn(data)
+    quantized_projection_outputs = [projection(data) for projection in hparam.quant_modules]
+
+    expected_score = ((output_grad * (quantized_attention_output - attention_output)) ** 2).sum()
+    old_projection_score = sum(
+        ((output_grad * (quantized - output)) ** 2).sum()
+        for output, quantized in zip(projection_outputs, quantized_projection_outputs)
+    )
+    qkv_stats = next(
+        stats
+        for stats in search_history["candidate_stats"].values()
+        if set(stats["module_names"]) == {"attn.q_proj", "attn.k_proj", "attn.v_proj"}
+    )
+    actual_score = qkv_stats["scores"][qkv_stats["formats"].index(int8_recipe)]
+
+    assert actual_score == pytest.approx(expected_score.item())
+    assert actual_score != pytest.approx(old_projection_score.item())
+    assert qkv_stats["costs"][qkv_stats["formats"].index(no_quant_recipe)] == sum(
+        projection.weight.numel() for projection in hparam.quant_modules
+    )
+
+
 INT4INT8_AWQ_CFG = {
     "quant_cfg": [
         {"quantizer_name": "*", "enable": False},