diff --git a/test/quantization/wrapq/wrappers/qwen_vl/test_quant_text_attention.py b/test/quantization/wrapq/wrappers/qwen_vl/test_quant_text_attention.py index 5db5f4cc..93806a40 100644 --- a/test/quantization/wrapq/wrappers/qwen_vl/test_quant_text_attention.py +++ b/test/quantization/wrapq/wrappers/qwen_vl/test_quant_text_attention.py @@ -14,6 +14,7 @@ import inspect import unittest +from unittest.mock import patch import torch from tico.quantization.wrapq.dtypes import DType @@ -180,6 +181,38 @@ def test_batched_layout_matches_unrolled_layout_without_quantization(self): torch.allclose(attn_unrolled, attn_batched, rtol=1e-4, atol=1e-5) ) + def test_unrolled_layout_runs_attention_per_query_head(self): + """Verify that grouped-query attention does not broadcast a KV head.""" + qattn = QuantQwen3VLTextAttention(self.fp_attn) + + batch_size, seq_len = 1, 4 + x = torch.randn(batch_size, seq_len, self.hidden_size) + pos = self._rand_rope(batch_size, seq_len) + mask = torch.zeros(batch_size, 1, seq_len, seq_len) + + logits_shapes: list[tuple[int, ...]] = [] + apply_scale = qattn._apply_attention_scale_if_needed + + def record_logits_shape(logits: torch.Tensor) -> torch.Tensor: + logits_shapes.append(tuple(logits.shape)) + return apply_scale(logits) + + with patch.object( + qattn, + "_apply_attention_scale_if_needed", + side_effect=record_logits_shape, + ): + with torch.no_grad(): + out, attn_weights = qattn(x, pos, attention_mask=mask) + + expected_logits_shape = (batch_size, seq_len, seq_len) + self.assertEqual(logits_shapes, [expected_logits_shape] * self.num_heads) + self.assertEqual(out.shape, (batch_size, seq_len, self.hidden_size)) + self.assertEqual( + attn_weights.shape, + (batch_size, self.num_heads, seq_len, seq_len), + ) + def test_forward_diff(self): qattn = QuantQwen3VLTextAttention(self.fp_attn) qattn.enable_calibration() diff --git a/tico/quantization/wrapq/wrappers/qwen_vl/quant_text_attention.py b/tico/quantization/wrapq/wrappers/qwen_vl/quant_text_attention.py index 9fe22019..f8eb5ab8 100644 --- a/tico/quantization/wrapq/wrappers/qwen_vl/quant_text_attention.py +++ b/tico/quantization/wrapq/wrappers/qwen_vl/quant_text_attention.py @@ -44,8 +44,8 @@ class QuantQwen3VLTextAttention(QuantModuleBase): """ Quantized Qwen3-VL text attention wrapper with selectable attention layout. - The default `npu_export` profile preserves the existing per-KV-head - unrolled graph for NPU export. Set `PTQConfig.model_args["profile"]` to + The default `npu_export` profile uses a per-query-head unrolled graph for + NPU export. Set `PTQConfig.model_args["profile"]` to "reference_eval" or pass `model_args["attention"]` overrides to run a Hugging Face-like batched attention graph for experiments. """ @@ -603,53 +603,50 @@ def _forward_unrolled( batch_size: int, q_len: int, ): - """Run the NPU-export-friendly per-KV-head attention path.""" + """Run the NPU-export-friendly per-query-head attention path.""" attn_weights_parts: list[torch.Tensor] = [] attn_out_parts: list[torch.Tensor] = [] - n_kv = present_k.size(1) - kv_rep = self.kv_rep + # Remove the singleton head dimension so every attention operation has + # the same rank and shape on both sides. The mask builder guarantees + # that the head dimension is one for this decomposed path. + attn_mask = attention_mask.squeeze(1) - for i in range(n_kv): - # (B, 1, K, H) - k_i = present_k[:, i : i + 1, :, :] - v_i = present_v[:, i : i + 1, :, :] + for kv_i in range(self.num_kv_heads): + # (B, K, H) + k_i = present_k[:, kv_i, :, :] + v_i = present_v[:, kv_i, :, :] - # (B, G, S, H) where G=kv_rep - h0 = i * kv_rep - h1 = (i + 1) * kv_rep - q_i = q_rot[:, h0:h1, :, :] + for rep_i in range(self.kv_rep): + q_idx = kv_i * self.kv_rep + rep_i + q_i = q_rot[:, q_idx, :, :] # (B, S, H) - # logits: (B, G, S, K) - logits_i = q_i @ k_i.transpose(-2, -1) - logits_i = self._apply_attention_scale_if_needed(logits_i) - logits_i = self._fq(logits_i, self.obs_logits) + # logits: (B, S, K) + logits_i = q_i @ k_i.transpose(-2, -1) + logits_i = self._apply_attention_scale_if_needed(logits_i) + logits_i = self._fq(logits_i, self.obs_logits) - # mask add: broadcast on head axis (1 -> G) - logits_i = self._fq(logits_i + attention_mask, self.obs_mask_add) + logits_i = self._fq(logits_i + attn_mask, self.obs_mask_add) - # softmax - attn_i = torch.softmax(logits_i, dim=-1, dtype=torch.float32).to( - q_rot.dtype - ) - attn_i = self._fq(attn_i, self.obs_softmax) + attn_i = torch.softmax(logits_i, dim=-1, dtype=torch.float32).to( + q_i.dtype + ) + attn_i = self._fq(attn_i, self.obs_softmax) - # out: (B, G, S, H) - out_i = self._fq(attn_i @ v_i, self.obs_attn_out) + # out: (B, S, H) + out_i = self._fq(attn_i @ v_i, self.obs_attn_out) - attn_weights_parts.append(attn_i) - attn_out_parts.append(out_i) + attn_weights_parts.append(attn_i) + attn_out_parts.append(out_i) - # Concatenate heads back: (B, n_h, S, K) / (B, n_h, S, H) + # Stack heads back: (B, n_h, S, K) / (B, S, n_h, H) attn_weights = self._fq( - torch.cat(attn_weights_parts, dim=1), self.obs_attn_weights + torch.stack(attn_weights_parts, dim=1), self.obs_attn_weights ) - attn_out_h = self._fq(torch.cat(attn_out_parts, dim=1), self.obs_attn_out_h) + attn_out_h = self._fq(torch.stack(attn_out_parts, dim=2), self.obs_attn_out_h) # Attention output - attn_out = attn_out_h.transpose(1, 2).reshape( - batch_size, q_len, -1 - ) # (B, S, n_h * H) + attn_out = attn_out_h.reshape(batch_size, q_len, -1) # (B, S, n_h * H) # Final projection out = self.o_proj(attn_out)