Precomputed state (Pi, S, codebook) + kernel launch + PyTorch fallbacks.
| 32 | |
| 33 | |
| 34 | class TurboQuantEngine: |
| 35 | """Precomputed state (Pi, S, codebook) + kernel launch + PyTorch fallbacks.""" |
| 36 | |
| 37 | def __init__( |
| 38 | self, |
| 39 | head_dim: int = HEAD_DIM, |
| 40 | total_bits: int = DEFAULT_TOTAL_BITS, |
| 41 | seed: int = DEFAULT_SEED, |
| 42 | device: str = "cpu", |
| 43 | ): |
| 44 | self.head_dim = head_dim |
| 45 | self.total_bits = total_bits |
| 46 | self.mse_bits = max(total_bits - 1, 1) |
| 47 | self.device = device |
| 48 | |
| 49 | self.Pi = _generate_rotation_matrix(head_dim, seed, device) |
| 50 | self.PiT = self.Pi.T.contiguous() |
| 51 | self.S = _generate_qjl_matrix(head_dim, seed, device) |
| 52 | self.ST = self.S.T.contiguous() |
| 53 | |
| 54 | self.key_codebook = LloydMaxCodebook(head_dim, self.mse_bits) |
| 55 | self.val_codebook = LloydMaxCodebook(head_dim, total_bits) |
| 56 | |
| 57 | self.scale = 1.0 / math.sqrt(head_dim) |
| 58 | self.correction_scale = math.sqrt(math.pi / 2) / head_dim |
| 59 | |
| 60 | @torch.no_grad() |
| 61 | def compress_keys_pytorch(self, K: torch.Tensor) -> dict: |
| 62 | """K: (seq_k, head_dim) -> compressed dict.""" |
| 63 | K_f = K.float() |
| 64 | vec_norms = torch.norm(K_f, dim=-1, keepdim=True) |
| 65 | K_normed = K_f / (vec_norms + 1e-8) |
| 66 | |
| 67 | rotated = K_normed @ self.PiT.float() |
| 68 | |
| 69 | centroids = self.key_codebook.centroids.to(K.device) |
| 70 | diffs = rotated.unsqueeze(-1) - centroids |
| 71 | indices = diffs.abs().argmin(dim=-1).to(torch.uint8) |
| 72 | |
| 73 | y_hat = centroids[indices.long()] |
| 74 | k_mse = (y_hat @ self.Pi.float()) * vec_norms |
| 75 | |
| 76 | residual = K_f - k_mse |
| 77 | residual_norms = torch.norm(residual, dim=-1) |
| 78 | |
| 79 | projected = residual @ self.ST.float() |
| 80 | signs = torch.sign(projected).to(torch.int8) |
| 81 | signs[signs == 0] = 1 |
| 82 | |
| 83 | return { |
| 84 | "indices": indices, |
| 85 | "k_mse": k_mse.half(), |
| 86 | "qjl_signs": signs, |
| 87 | "vec_norms": vec_norms.squeeze(-1).half(), |
| 88 | "residual_norms": residual_norms.half(), |
| 89 | } |
| 90 | |
| 91 | @torch.no_grad() |
no outgoing calls