Typical decode scenario: seq_q=1, seq_k=large.
()
| 99 | |
| 100 | |
| 101 | def test_single_decode_token(): |
| 102 | """Typical decode scenario: seq_q=1, seq_k=large.""" |
| 103 | engine = TurboQuantEngine(head_dim=128, total_bits=3, device="cpu") |
| 104 | Q = torch.randn(1, 128).half() |
| 105 | K = torch.randn(2048, 128).half() |
| 106 | |
| 107 | compressed_k = engine.compress_keys_pytorch(K) |
| 108 | scores = engine.attention_scores_pytorch(Q, compressed_k) |
| 109 | |
| 110 | assert scores.shape == (1, 2048) |
| 111 | assert not torch.isnan(scores).any() |
| 112 | assert not torch.isinf(scores).any() |
| 113 | |
| 114 | |
| 115 | def test_needle_in_haystack(): |
nothing calls this directly
no test coverage detected