fix(quantizer): add dimension guards for clear errors on GQA shape mismatch (#69)

Previously, passing a tensor whose last dimension did not match the
quantizer's configured dim produced an opaque `RuntimeError` from
PyTorch's reshape. This surfaced with Gemma 4 E4B on transformers 5.x,
where GQA expansion changed the KV tensor shape before cache update.

- Add `ValueError` guards to `quantize()`, `dequantize()`, and
`estimate_inner_product()` that check `shape[-1] == self.dim` with
actionable error messages
- Add `TestDimMismatchGuard` (4 tests) and `TestGQAShapes` (5 tests)
covering both MSE and Prod quantizers with expanded head counts
- Extract `TestHFTokenPassthrough` to `tests/test_verify_token.py` (test
maturity: split oversized test file by concern)

Test: `uv run pytest tests/test_quantizer.py tests/test_verify_token.py
-v`

Closes #67

Author: Alberto-Codes

src/turboquant_vllm/quantizer.py

@@ -121,7 +121,18 @@ def quantize(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         Returns:
             Tuple of (indices, norms) where indices is a long tensor of
             shape (..., dim) and norms is a float tensor of shape (..., 1).
+
+        Raises:
+            ValueError: If ``x.shape[-1] != self.dim``.
         """
+        if x.shape[-1] != self.dim:
+            msg = (
+                f"Input last dimension {x.shape[-1]} does not match "
+                f"quantizer dim {self.dim}. Ensure the compressor's "
+                f"head_dim matches the model's actual head dimension."
+            )
+            raise ValueError(msg)
+
         # Store original shape and flatten to 2D
         orig_shape = x.shape
         flat = x.reshape(-1, self.dim).float()
@@ -151,7 +162,18 @@ def dequantize(self, indices: torch.Tensor, norms: torch.Tensor) -> torch.Tensor
 
         Returns:
             Reconstructed float tensor of shape (..., dim).
+
+        Raises:
+            ValueError: If ``indices.shape[-1] != self.dim``.
         """
+        if indices.shape[-1] != self.dim:
+            msg = (
+                f"Indices last dimension {indices.shape[-1]} does not match "
+                f"quantizer dim {self.dim}. Check that the same quantizer "
+                f"instance is used for quantize and dequantize."
+            )
+            raise ValueError(msg)
+
         orig_shape = indices.shape
         flat_idx = indices.reshape(-1, self.dim)
         flat_norms = norms.reshape(-1, 1)
@@ -316,7 +338,18 @@ def estimate_inner_product(
         Returns:
             Inner product estimates, shape matching broadcast of query and key
             batch dimensions.
+
+        Raises:
+            ValueError: If ``query.shape[-1] != self.dim``.
         """
+        if query.shape[-1] != self.dim:
+            msg = (
+                f"Query last dimension {query.shape[-1]} does not match "
+                f"quantizer dim {self.dim}. Ensure the compressor's "
+                f"head_dim matches the model's actual head dimension."
+            )
+            raise ValueError(msg)
+
         # MSE component: <q, k_mse>
         k_mse = self.mse_quantizer.dequantize(indices, norms)
         mse_term = (query.float() * k_mse).sum(dim=-1, keepdim=True)

tests/test_quantizer.py

@@ -175,3 +175,98 @@ def test_zero_vector_inner_product(
             query, indices, norms, qjl_signs, res_norms
         )
         assert abs(estimated.item()) < 0.01
+
+
+@pytest.mark.unit
+class TestDimMismatchGuard:
+    """Validate clear errors when tensor last dim != quantizer dim (GH-67)."""
+
+    def test_quantize_wrong_last_dim_raises_valueerror(self) -> None:
+        """Feeding shape[-1] != self.dim should raise ValueError, not RuntimeError."""
+        q = TurboQuantMSE(128, bits=4)
+        x = torch.randn(1, 2, 10, 256)  # last dim 256 != quantizer dim 128
+        with pytest.raises(ValueError, match="last dimension"):
+            q.quantize(x)
+
+    def test_dequantize_wrong_last_dim_raises_valueerror(self) -> None:
+        """Dequantize with indices whose last dim != self.dim should raise ValueError."""
+        q = TurboQuantMSE(128, bits=4)
+        indices = torch.zeros(1, 2, 10, 256, dtype=torch.long)
+        norms = torch.ones(1, 2, 10, 1)
+        with pytest.raises(ValueError, match="last dimension"):
+            q.dequantize(indices, norms)
+
+    def test_prod_quantize_wrong_dim_raises_valueerror(self) -> None:
+        """TurboQuantProd should surface ValueError from inner MSE quantizer."""
+        q = TurboQuantProd(128, bits=4)
+        x = torch.randn(1, 2, 10, 256)
+        with pytest.raises(ValueError, match="last dimension"):
+            q.quantize(x)
+
+    def test_estimate_inner_product_wrong_query_dim_raises_valueerror(self) -> None:
+        """Query with shape[-1] != self.dim should raise ValueError."""
+        q = TurboQuantProd(128, bits=4)
+        keys = torch.randn(1, 2, 10, 128)
+        indices, norms, signs, res_norms = q.quantize(keys)
+        bad_query = torch.randn(1, 2, 10, 256)  # wrong last dim
+        with pytest.raises(ValueError, match="last dimension"):
+            q.estimate_inner_product(bad_query, indices, norms, signs, res_norms)
+
+
+@pytest.mark.unit
+class TestGQAShapes:
+    """Validate quantizer handles GQA-expanded tensor shapes (GH-67).
+
+    GQA models may have different numbers of KV heads vs attention heads.
+    The quantizer must handle any leading dimensions as long as the last
+    dimension matches self.dim.
+    """
+
+    def test_quantize_gqa_expanded_shape(self) -> None:
+        """GQA-expanded tensor (batch, expanded_heads, seq, dim) should work."""
+        q = TurboQuantMSE(DIM, bits=4)
+        # Simulate GQA expansion: 2 KV heads expanded to 8 attention heads
+        x = torch.randn(1, 8, 32, DIM)
+        indices, norms = q.quantize(x)
+        assert indices.shape == (1, 8, 32, DIM)
+        assert norms.shape == (1, 8, 32, 1)
+
+    def test_quantize_unexpanded_kv_shape(self) -> None:
+        """Unexpanded KV tensor (batch, kv_heads, seq, dim) should work."""
+        q = TurboQuantMSE(DIM, bits=4)
+        x = torch.randn(1, 2, 32, DIM)
+        indices, norms = q.quantize(x)
+        assert indices.shape == (1, 2, 32, DIM)
+        assert norms.shape == (1, 2, 32, 1)
+
+    def test_roundtrip_gqa_expanded_quality(self) -> None:
+        """Round-trip through GQA-expanded shape preserves cosine quality."""
+        q = TurboQuantMSE(DIM, bits=4)
+        x = torch.randn(1, 8, 32, DIM)
+        indices, norms = q.quantize(x)
+        reconstructed = q.dequantize(indices, norms)
+        assert reconstructed.shape == x.shape
+        cos = torch.nn.functional.cosine_similarity(
+            x.flatten().float(), reconstructed.flatten().float(), dim=0
+        ).item()
+        assert cos > 0.80, f"Round-trip cosine {cos:.4f} below 0.80 threshold"
+
+    def test_quantize_head_dim_256_gqa(self) -> None:
+        """Head dim 256 with GQA-expanded heads (Gemma 4 E4B scenario)."""
+        q = TurboQuantMSE(256, bits=4)
+        x = torch.randn(1, 8, 16, 256)  # 8 expanded heads, head_dim=256
+        indices, norms = q.quantize(x)
+        assert indices.shape == (1, 8, 16, 256)
+        assert norms.shape == (1, 8, 16, 1)
+        reconstructed = q.dequantize(indices, norms)
+        assert reconstructed.shape == x.shape
+
+    def test_prod_quantize_gqa_expanded(self) -> None:
+        """TurboQuantProd handles GQA-expanded shapes (key compressor path)."""
+        q = TurboQuantProd(DIM, bits=4)
+        x = torch.randn(1, 8, 32, DIM)
+        indices, norms, signs, res_norms = q.quantize(x)
+        assert indices.shape == (1, 8, 32, DIM)
+        assert norms.shape == (1, 8, 32, 1)
+        assert signs.shape == (1, 8, 32, DIM)
+        assert res_norms.shape == (1, 8, 32, 1)

tests/test_verify.py

@@ -479,35 +479,4 @@ def test_run_verification_molmo2_passes(self) -> None:
             torch.cuda.empty_cache()
 
 
-class TestHFTokenPassthrough:
-    """Verify that HF_TOKEN is forwarded to all from_pretrained calls."""
-
-    def test_token_passed_to_auto_config(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        """AutoConfig.from_pretrained receives token= from HF_TOKEN env var."""
-        monkeypatch.setenv("HF_TOKEN", "hf_test_sentinel")
-        captured: dict[str, object] = {}
-
-        def spy(*args: object, **kwargs: object) -> object:
-            captured["token"] = kwargs.get("token")
-            raise RuntimeError("stop after capture")
-
-        monkeypatch.setattr("transformers.AutoConfig.from_pretrained", spy)
-        with pytest.raises(RuntimeError, match="stop after capture"):
-            _run_verification("fake/model", bits=4, threshold=0.99)
-
-        assert captured["token"] == "hf_test_sentinel"
-
-    def test_token_none_when_env_unset(self, monkeypatch: pytest.MonkeyPatch) -> None:
-        """token= is None when HF_TOKEN is not in the environment."""
-        monkeypatch.delenv("HF_TOKEN", raising=False)
-        captured: dict[str, object] = {}
-
-        def spy(*args: object, **kwargs: object) -> object:
-            captured["token"] = kwargs.get("token")
-            raise RuntimeError("stop after capture")
-
-        monkeypatch.setattr("transformers.AutoConfig.from_pretrained", spy)
-        with pytest.raises(RuntimeError, match="stop after capture"):
-            _run_verification("fake/model", bits=4, threshold=0.99)
-
-        assert captured["token"] is None
+# TestHFTokenPassthrough extracted to test_verify_token.py (GH-67 test maturity)

tests/test_verify_token.py

@@ -0,0 +1,43 @@
+"""Tests for HF_TOKEN passthrough in verify CLI (extracted from test_verify.py)."""
+
+from __future__ import annotations
+
+import pytest
+
+from turboquant_vllm.verify import _run_verification
+
+pytestmark = [pytest.mark.unit]
+
+
+class TestHFTokenPassthrough:
+    """Verify that HF_TOKEN is forwarded to all from_pretrained calls."""
+
+    def test_token_passed_to_auto_config(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """AutoConfig.from_pretrained receives token= from HF_TOKEN env var."""
+        monkeypatch.setenv("HF_TOKEN", "hf_test_sentinel")
+        captured: dict[str, object] = {}
+
+        def spy(*args: object, **kwargs: object) -> object:
+            captured["token"] = kwargs.get("token")
+            raise RuntimeError("stop after capture")
+
+        monkeypatch.setattr("transformers.AutoConfig.from_pretrained", spy)
+        with pytest.raises(RuntimeError, match="stop after capture"):
+            _run_verification("fake/model", bits=4, threshold=0.99)
+
+        assert captured["token"] == "hf_test_sentinel"
+
+    def test_token_none_when_env_unset(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        """token= is None when HF_TOKEN is not in the environment."""
+        monkeypatch.delenv("HF_TOKEN", raising=False)
+        captured: dict[str, object] = {}
+
+        def spy(*args: object, **kwargs: object) -> object:
+            captured["token"] = kwargs.get("token")
+            raise RuntimeError("stop after capture")
+
+        monkeypatch.setattr("transformers.AutoConfig.from_pretrained", spy)
+        with pytest.raises(RuntimeError, match="stop after capture"):
+            _run_verification("fake/model", bits=4, threshold=0.99)
+
+        assert captured["token"] is None