Bug for register_attention_for_kv_quant

[AIgiraffe]

When I tested the register_attention_for_kv_quant function of modelopt, I found the following issues：
a simple attention pytorch code is follows:

class CommonAttention(nn.Module):
def forward(self, qkv: tuple, extra_arg=None):
# NOTE: Add unused argument y with default value to test that replaced attention retain original defaults
q, k, v = qkv

    attn = q @ k
    attn = F.softmax(attn)

    attn = (attn @ v)
    return attn 

@classmethod
def get_input(cls, device: str = "cpu"):
    q = torch.randn(1, 4, 8, device=device)
    k = torch.randn(1, 4, 8, device=device)
    v = torch.randn(1, 4, 8, device=device)
    return (q, k, v),

After I applied register_attention_for_kv_quant to CommonAttention，CommonAttention is follows:

class _QuantCommonAttention(nn.Module):

def forward(self, qkv: tuple, extra_arg=None):
    q, k, v = qkv
    k = k.transpose(-2, -1)
    attn = q @ self.v_bmm_quantizer(k)
    attn = F.softmax(attn)
    attn = self.q_bmm_quantizer(attn) @ torch.transpose(self.k_bmm_quantizer(torch.transpose(v, -1, -2)), -1, -2)
    return attn

@classmethod
def get_input(cls, device: str='cpu'):
    q = torch.randn(1, 4, 8, device=device)
    k = torch.randn(1, 4, 8, device=device)
    v = torch.randn(1, 4, 8, device=device)
    return (q, k, v)

Is this correct?

[ChenhanYu]

Triage Analysis

Field	Value
Classification	Bug
Severity	High
Complexity	Moderate
Auto-fixable	No

Summary

The register_attention_for_kv_quant function incorrectly patches binary matmul operations (@) in attention layers. The function applies quantizers to incorrect operands and uses wrong transpose operations, resulting in semantically incorrect transformations of the attention computation.

Root Cause / Approach

The patch_binop function has two critical issues: (1) It applies the v_bmm_quantizer to the wrong operand in the first matmul (should quantize v, not k), and (2) It applies incorrect transpose logic that doesn't preserve the original computation semantics. The original code attn @ v becomes (quantized_attn) @ (transposed-quantized-transposed v), which changes the shape and meaning of the operation.

Suggested Fix

The patch_binop function needs to be corrected to: (1) For the first @ operation (q @ k.T), apply quantizers to q and k appropriately; (2) For the second @ operation (attn @ v), apply q_bmm_quantizer to attn and v_bmm_quantizer to v without incorrect transpose wrapping. The transpose=True parameter should only affect the key tensor in the attention pattern (q @ k.T), not the value tensor in the output computation. Review the intended quantization points: typically q and k are quantized in the attention computation, and v is quantized separately in the weighted sum.

Relevant Files

• modelopt/torch/quantization/plugins/attention.py
• tests/unit/torch/quantization/plugins/test_attention_quant.py

---

Auto-triaged by pensieve /magic-triage