| import torch | |
| from torch import nn | |
| import cupy as cp | |
| import math | |
| from os import path | |
| class Quantizer(nn.Module): | |
| def __init__(self, config, codebook): | |
| super().__init__() | |
| self.nsq, nc, self.d = codebook.shape | |
| self.b = int(math.log2(nc)) | |
| head_dim = config.hidden_size // config.num_attention_heads | |
| self.head_dim = head_dim | |
| qpk = config.num_attention_heads // config.num_key_value_heads | |
| self.window_length = getattr(config, 'window_length', 32) | |
| self.register_buffer('codebook', codebook) | |
| with open(path.join(path.dirname(__file__), "quantize.cu"), "r") as f: | |
| kernel_code = f.read().replace('__NSQ__', str(self.nsq)).replace('__B__', str(self.b)).replace('__D__', str(self.d)) | |
| self._quantize = cp.RawKernel( | |
| kernel_code, | |
| 'quantize', | |
| backend="nvrtc" | |
| ) | |
| with open(path.join(path.dirname(__file__), "dequantize.cu"), "r") as f: | |
| kernel_code = f.read().replace('__NSQ__', str(self.nsq)).replace('__B__', str(self.b)).replace('__D__', str(self.d)) | |
| self._dequantize = cp.RawKernel( | |
| kernel_code, | |
| 'dequantize', | |
| backend="nvrtc" | |
| ) | |
| with open(path.join(path.dirname(__file__), "dequantize_rope.cu"), "r") as f: | |
| kernel_code = f.read().replace('__NSQ__', str(self.nsq)).replace('__B__', str(self.b)).replace('__D__', str(self.d)).replace('__HEAD_DIM__', str(head_dim)) | |
| self._dequantize_rope = cp.RawKernel( | |
| kernel_code, | |
| 'dequantize_rope', | |
| backend="nvrtc" | |
| ) | |
| with open(path.join(path.dirname(__file__), "fused_rope_mult.cu"), "r") as f: | |
| kernel_code = f.read().replace('__NSQ__', str(self.nsq)).replace('__B__', str(self.b)).replace('__D__', str(self.d)).replace('__HEAD_DIM__', str(head_dim)) | |
| self._fused_rope_mult = cp.RawKernel( | |
| kernel_code, | |
| 'fused_rope_mult', | |
| backend="nvrtc" | |
| ) | |
| with open(path.join(path.dirname(__file__), "fused_rope_pos_mult_mqa.cu"), "r") as f: | |
| kernel_code = f.read().replace('__NSQ__', str(self.nsq)).replace('__B__', str(self.b)).replace('__D__', str(self.d)).replace('__HEAD_DIM__', str(head_dim)).replace('__QPK__', str(qpk)).replace('__ROPE_THETA__', str(config.rope_theta)) | |
| self._fused_rope_pos_mult = cp.RawKernel( | |
| kernel_code, | |
| 'fused_rope_pos_mult', | |
| backend="nvrtc" | |
| ) | |
| with open(path.join(path.dirname(__file__), "fused_mult_len.cu"), "r") as f: | |
| kernel_code = f.read().replace('__NSQ__', str(self.nsq)).replace('__B__', str(self.b)).replace('__D__', str(self.d)).replace('__HEAD_DIM__', str(head_dim)).replace('__QPK__', str(qpk)) | |
| self._fused_mult = cp.RawKernel( | |
| kernel_code, | |
| 'fused_mult', | |
| backend="nvrtc" | |
| ) | |
| def quantize(self, x): | |
| n = x.numel() // x.shape[-1] | |
| codes = torch.empty(self.nsq, n, dtype=torch.uint8, device=x.device) | |
| blocks_per_grid = (self.nsq, ) | |
| threads_per_block = (1024, ) | |
| self._quantize(grid=blocks_per_grid, block=threads_per_block, shared_mem=(2 ** self.b) * self.d * 2, args=[ | |
| self.codebook.data_ptr(), | |
| x.data_ptr(), | |
| codes.data_ptr(), | |
| n | |
| ]) | |
| return codes | |
| def dequantize(self, codes): | |
| n = codes.numel() // codes.shape[0] | |
| x = torch.zeros(n, self.nsq * self.d, dtype=torch.float16, device=codes.device) | |
| blocks_per_grid = (self.nsq, ) | |
| threads_per_block = (1024, ) | |
| self._dequantize(grid=blocks_per_grid, block=threads_per_block, shared_mem=(2 ** self.b) * self.d * 2, args=[ | |
| self.codebook.data_ptr(), | |
| codes.data_ptr(), | |
| x.data_ptr(), | |
| n | |
| ]) | |
| return x | |
| def dequantize_rope(self, codes): | |
| _, batch_size, seq_len = codes.shape | |
| n = batch_size * seq_len | |
| x = torch.zeros(n, self.nsq * self.d, dtype=torch.float16, device=codes.device) | |
| blocks_per_grid = (self.nsq, ) | |
| threads_per_block = (1024, ) | |
| self._dequantize_rope(grid=blocks_per_grid, block=threads_per_block, shared_mem=(2 ** self.b) * self.d * 2, args=[ | |
| self.codebook.data_ptr(), | |
| codes.data_ptr(), | |
| x.data_ptr(), | |
| batch_size, seq_len | |
| ]) | |
| return x | |
| def fused_rope_mult(self, codes, queries): | |
| _, batch_size, k_len = codes.shape | |
| _, n_heads, q_len, _ = queries.shape | |
| out = torch.zeros(batch_size, n_heads, q_len, k_len, dtype=torch.float16, device=codes.device) | |
| blocks_per_grid = (self.nsq, ) | |
| threads_per_block = (1024, ) | |
| self._fused_rope_mult(grid=blocks_per_grid, block=threads_per_block, shared_mem=(2 ** self.b) * self.d * 2, args=[ | |
| self.codebook.data_ptr(), | |
| codes.data_ptr(), | |
| queries.data_ptr(), | |
| out.data_ptr(), | |
| batch_size, q_len, k_len | |
| ]) | |
| return out | |
| def fused_rope_pos_mult(self, codes, queries, position_ids): | |
| _, batch_size, k_len = codes.shape | |
| _, n_heads, q_len, _ = queries.shape | |
| position_offsets = position_ids[:, -1] - k_len + 1 | |
| out = torch.zeros(batch_size, n_heads, q_len, k_len, dtype=torch.float32, device=codes.device) | |
| blocks_per_grid = (self.nsq, ) | |
| threads_per_block = (1024, ) | |
| self._fused_rope_pos_mult(grid=blocks_per_grid, block=threads_per_block, shared_mem=(2 ** self.b) * self.d * 2, args=[ | |
| self.codebook.data_ptr(), | |
| codes.data_ptr(), | |
| position_offsets.data_ptr(), | |
| queries.data_ptr(), | |
| out.data_ptr(), | |
| batch_size, q_len, k_len | |
| ]) | |
| return out | |
| def fused_mult(self, codes, weights, skip_last=0): | |
| batch_size, n_heads, q_len, k_len = weights.shape | |
| out = torch.zeros(batch_size, n_heads, q_len, self.head_dim, dtype=torch.float16, device=codes.device) | |
| blocks_per_grid = (self.nsq, ) | |
| threads_per_block = (min(1024, batch_size), ) | |
| self._fused_mult(grid=blocks_per_grid, block=threads_per_block, shared_mem=(2 ** self.b) * self.d * 2, args=[ | |
| self.codebook.data_ptr(), | |
| codes.data_ptr(), | |
| weights.data_ptr(), | |
| out.data_ptr(), | |
| batch_size, q_len, k_len, k_len - skip_last | |
| ]) | |
| return out | |