(seq_q, seq_k, total_bits)
| 21 | @pytest.mark.parametrize("seq_q,seq_k", [(1, 64), (1, 256), (1, 1024), (16, 512)]) |
| 22 | @pytest.mark.parametrize("total_bits", [3, 4]) |
| 23 | def test_scores_shape(seq_q, seq_k, total_bits): |
| 24 | engine = TurboQuantEngine(head_dim=128, total_bits=total_bits, device="cpu") |
| 25 | |
| 26 | Q = torch.randn(seq_q, 128).half() |
| 27 | K = torch.randn(seq_k, 128).half() |
| 28 | compressed_k = engine.compress_keys_pytorch(K) |
| 29 | scores = engine.attention_scores_pytorch(Q, compressed_k) |
| 30 | |
| 31 | assert scores.shape == (seq_q, seq_k) |
| 32 | |
| 33 | |
| 34 | @pytest.mark.parametrize("total_bits", [3, 4]) |
nothing calls this directly
no test coverage detected