Group MLA q_a_proj/kv_a_proj_with_mqa projections in auto_quantize

TRT-LLM fuses the MLA low-rank input projections (DeepSeek/GLM lineage) into
a single fused_qkv_a_proj_with_mqa GEMM, so auto_quantize must assign both
shards one shared quantization format; without this rule any scoring method
can emit checkpoints the runtime cannot fuse.

Signed-off-by: Joshua Hill <joshua.hill@baseten.co>

Author: joshua-hill

modelopt/torch/quantization/algorithms.py

@@ -595,6 +595,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)$")
+_MLA_A_PROJ_RE = re.compile(r"^(.*?)\.(?:q_a_proj|kv_a_proj_with_mqa)$")
 
 
 def _linear_attn_qkvz_group_key(_model, name: str) -> str | None:
@@ -607,6 +608,11 @@ def _linear_attn_ba_group_key(_model, name: str) -> str | None:
     return f"{m.group(1)}/ba" if m else None
 
 
+def _mla_a_proj_group_key(_model, name: str) -> str | None:
+    m = _MLA_A_PROJ_RE.match(name)
+    return f"{m.group(1)}/qkv_a" if m else None
+
+
 def _module_search_space_signature(module_search_spaces) -> tuple:
     """Return a checkpoint-stable description of module-specific candidate spaces."""
     return tuple(
@@ -649,6 +655,14 @@ class _AutoQuantizeBaseSearcher(BaseSearcher, ABC):
         r"^(.*?)\.(gate_proj|up_proj)$",  # gate_proj, up_proj for llama like models
         r"^(.*?)\.(\d+\.(w1|w2|w3))$",  # mixtral experts
         r"^(.*?)\.((w1_linear|w2_linear|w3_linear)\.\d+)$",  # dbrx experts
+        # MLA low-rank input projections (DeepSeek/GLM lineage): TRT-LLM fuses them into
+        # fused_qkv_a_proj_with_mqa, so the shards must share one quantization format.
+        # A callable (not a regex) because a regex rule keys on match.group(1) -- the
+        # attention path -- which is the key the q_proj/k_proj/v_proj rule above already
+        # returns. On MLA built without a q LoRA rank (q_lora_rank=None, e.g.
+        # DeepSeek-V2-Lite) ``q_proj`` and ``kv_a_proj_with_mqa`` are siblings, so a
+        # shared key would merge the unfused q_proj into this group.
+        _mla_a_proj_group_key,
         # Qwen3.5/3.6 hybrid linear_attn: vLLM fuses (in_proj_qkv, in_proj_z)
         # into ``in_proj_qkvz`` and (in_proj_a, in_proj_b) into ``in_proj_ba`` and
         # requires fused shards to share quant_algo. Two callables (not one

tests/unit/torch/quantization/test_autoquant.py

@@ -1537,3 +1537,84 @@ def test_get_auto_quantize_config_emits_fused_expert_quantizer_names(with_persis
     assert f"{module_name}.down_proj_weight_quantizer" in quantizer_names
     assert f"{module_name}.weight_quantizer" not in quantizer_names
 
+
+def test_mla_projections_share_one_group():
+    """TRT-LLM fuses q_a_proj + kv_a_proj_with_mqa, so they must share one quant format."""
+
+    class _MLAAttention(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.q_a_proj = torch.nn.Linear(32, 32)
+            self.kv_a_proj_with_mqa = torch.nn.Linear(32, 32)
+            self.o_proj = torch.nn.Linear(32, 32)
+
+        def forward(self, x):
+            return self.o_proj(self.q_a_proj(x) + self.kv_a_proj_with_mqa(x))
+
+    class _MLABlock(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.self_attn = _MLAAttention()
+
+        def forward(self, x):
+            return self.self_attn(x)
+
+        def get_input(self):
+            return torch.randn(1, 4, 32)
+
+    model = _MLABlock()
+    mtq.auto_quantize(
+        model,
+        constraints={"effective_bits": 8.0},
+        quantization_formats=[mtq.INT8_DEFAULT_CFG],
+        data_loader=[model.get_input() for _ in range(2)],
+        forward_step=lambda model, batch: model(batch),
+        loss_func=lambda output, data: output.sum(),
+        num_calib_steps=2,
+        num_score_steps=2,
+        method="gradient",
+    )
+    hparam = model.self_attn.q_a_proj.get_hparam("quant_recipe")
+    assert model.self_attn.kv_a_proj_with_mqa.get_hparam("quant_recipe") == hparam
+    assert model.self_attn.o_proj.get_hparam("quant_recipe") != hparam
+
+
+def test_mla_group_does_not_absorb_unfused_q_proj():
+    """With q_lora_rank=None there is no q_a_proj; q_proj is NOT fused with kv_a_proj."""
+
+    class _MLAAttention(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.q_proj = torch.nn.Linear(32, 32)
+            self.kv_a_proj_with_mqa = torch.nn.Linear(32, 32)
+            self.o_proj = torch.nn.Linear(32, 32)
+
+        def forward(self, x):
+            return self.o_proj(self.q_proj(x) + self.kv_a_proj_with_mqa(x))
+
+    class _MLABlock(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.self_attn = _MLAAttention()
+
+        def forward(self, x):
+            return self.self_attn(x)
+
+        def get_input(self):
+            return torch.randn(1, 4, 32)
+
+    model = _MLABlock()
+    mtq.auto_quantize(
+        model,
+        constraints={"effective_bits": 8.0},
+        quantization_formats=[mtq.INT8_DEFAULT_CFG],
+        data_loader=[model.get_input() for _ in range(2)],
+        forward_step=lambda model, batch: model(batch),
+        loss_func=lambda output, data: output.sum(),
+        num_calib_steps=2,
+        num_score_steps=2,
+        method="gradient",
+    )
+    q_hparam = model.self_attn.q_proj.get_hparam("quant_recipe")
+    assert model.self_attn.kv_a_proj_with_mqa.get_hparam("quant_recipe") != q_hparam
+    assert model.self_attn.o_proj.get_hparam("quant_recipe") != q_hparam