From 7319fbcc994710d194229b2fc6c564a23f275e97 Mon Sep 17 00:00:00 2001 From: "yujiong.zhang" Date: Tue, 25 Aug 2026 13:45:35 +0800 Subject: [PATCH 1/3] fix: free QQQ workspace after module quantization --- gptqmodel/looper/qqq_processor.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/gptqmodel/looper/qqq_processor.py b/gptqmodel/looper/qqq_processor.py index c34571e4f..d42c7e7bc 100644 --- a/gptqmodel/looper/qqq_processor.py +++ b/gptqmodel/looper/qqq_processor.py @@ -162,7 +162,10 @@ def process( # logger.info(f"Quantizing module START: {name}, {gptq[name].shape()}") ## Need to return the quantized_weight for offloading q = qqq[module.name] - wq, q_scales, q_zeros, q_g_idx, duration, avg_loss, damp_percent, q_scales_extra, nsamples = q.quantize() + try: + wq, q_scales, q_zeros, q_g_idx, duration, avg_loss, damp_percent, q_scales_extra, nsamples = q.quantize() + finally: + q.free() q_scales = q_scales.to(CPU) q_zeros = q_zeros.to(CPU) From 655a6f4d6114940eead2b9a7b0408887a8ec775c Mon Sep 17 00:00:00 2001 From: "yujiong.zhang" Date: Tue, 25 Aug 2026 13:46:14 +0800 Subject: [PATCH 2/3] perf: use chunked Torch packing for QQQ weights --- gptqmodel/nn_modules/qlinear/qqq.py | 170 ++++++++++++++++------------ 1 file changed, 95 insertions(+), 75 deletions(-) diff --git a/gptqmodel/nn_modules/qlinear/qqq.py b/gptqmodel/nn_modules/qlinear/qqq.py index 0be56f99a..c0e01c36e 100644 --- a/gptqmodel/nn_modules/qlinear/qqq.py +++ b/gptqmodel/nn_modules/qlinear/qqq.py @@ -287,13 +287,8 @@ def list_buffers(self) -> List: buf.append(self.reduce_buffer) return buf - #def pack(self, linear: nn.Module, scales: t.Tensor, zeros: t.Tensor, g_idx: t.Tensor = None): def pack(self, linear: torch.nn.Module, scales: torch.Tensor, s_extra=None): - """Pack a fake-quantized linear layer into this actual Marlin representation. - @linear: fake-quantized `torch.nn.Linear` layer to convert (must be of type `torch.float16`) - @scales: corresponding quantization scales of shape `(infeatures, groups)` - @s_extra: corresponding quantization scales of shape `(1, outfeatures)` - """ + """Pack a fake-quantized linear layer into the Marlin representation.""" if self.group_size != self.in_features: assert s_extra is not None, "s_extra is needed" if linear.weight.dtype != torch.float16: @@ -302,87 +297,112 @@ def pack(self, linear: torch.nn.Module, scales: torch.Tensor, s_extra=None): If you can ensure your GEMM results don't overflow torch.float16, it will still function correctly. Otherwise, it will yield incorrect results.""" ) - s = scales.t() - w = linear.weight.data.t() - if self.group_size != self.in_features: - w = w.reshape((-1, self.group_size, self.out_features)) - w = w.permute(1, 0, 2) - w = w.reshape((self.group_size, -1)) - s = s.reshape((1, -1)) - w = torch.round(w / s).int() - if self.group_size != self.in_features: - w += (self.maxq + 1) // 2 - w = torch.clamp(w, 0, self.maxq) - else: - w = torch.clamp(w, -self.maxq, self.maxq) + + raw_scales = scales.t() if self.group_size != self.in_features: s_extra = s_extra.reshape(1, -1).to(dtype=torch.float32) - s = (s.reshape(-1, self.out_features) / s_extra).to(dtype=torch.float16) - - w = w.reshape((self.group_size, -1, self.out_features)) - w = w.permute(1, 0, 2) - w = w.reshape((self.in_features, self.out_features)).contiguous() - s = s.reshape((-1, len(self._scale_perm)))[:, self._scale_perm] - s_extra = s_extra.reshape((-1, len(self._scale_perm_single)))[ - :, self._scale_perm_single - ] - s_extra = s_extra.reshape((-1, self.out_features)).contiguous() + packed_s_group = (raw_scales / s_extra).to(dtype=torch.float16) + packed_s_group = packed_s_group.reshape( + (-1, len(self._scale_perm)) + )[:, self._scale_perm].reshape((-1, self.out_features)).contiguous() + packed_s_channel = s_extra.reshape( + (-1, len(self._scale_perm_single)) + )[:, self._scale_perm_single].reshape((-1, self.out_features)).contiguous() else: - # NOTE(zhangying): div 2 ** (8 - self.bits)) to deal with right_shift in unpacking - s = ( - (s / (2 ** (8 - self.bits))) + packed_s_group = None + packed_s_channel = ( + (raw_scales / (2 ** (8 - self.bits))) .reshape((-1, len(self._scale_perm_single)))[:, self._scale_perm_single] .to(dtype=torch.float32) + .reshape((-1, self.out_features)) + .contiguous() ) - s = s.reshape((-1, self.out_features)).contiguous() - w = w.reshape( - ( - self.in_features // self.tile, - self.tile, - self.out_features // self.tile, - self.tile, + + input_chunk_size = min(1024, self.in_features) + output_chunk_size = min(256, self.out_features) + input_chunk_size -= input_chunk_size % self.tile + output_chunk_size -= output_chunk_size % 64 + if input_chunk_size == 0 or output_chunk_size == 0: + raise ValueError( + f"QQQ pack requires dimensions divisible by {self.tile} and 64, " + f"got in_features={self.in_features}, out_features={self.out_features}" ) + + packed_weight = torch.empty( + (self.in_features // self.tile, self.out_features * 2), + dtype=torch.int32, + device=CPU, ) - w = w.permute((0, 2, 1, 3)) - w = w.reshape((self.in_features // self.tile, self.out_features * self.tile)) - res = w - res = res.reshape((-1, self._perm.numel()))[:, self._perm].reshape(res.shape) - q = np.zeros((res.shape[0], res.shape[1] // 8), dtype=np.uint32) - res = res.cpu().numpy().astype(np.uint32) - if self.group_size != self.in_features: - for i in range(8): - q |= res[:, i::8] << 4 * i - else: - for i in range(8): - q |= (res[:, i::8] & 0xF) << 4 * i - q = torch.from_numpy(q.astype(np.int32)).to(CPU) + weight = linear.weight.data + perm = self._perm.to(device=weight.device) + for input_start in range(0, self.in_features, input_chunk_size): + input_end = min(input_start + input_chunk_size, self.in_features) + input_size = input_end - input_start + if input_size % self.tile != 0: + raise ValueError("QQQ pack input chunk is not tile aligned") + + input_indices = torch.arange( + input_start, + input_end, + device=weight.device, + dtype=torch.long, + ) + if self.group_size != self.in_features: + group_indices = torch.div(input_indices, self.group_size, rounding_mode="floor") + scale_chunk = raw_scales.index_select(0, group_indices) + else: + scale_chunk = raw_scales.expand(input_size, -1) + + for output_start in range(0, self.out_features, output_chunk_size): + output_end = min(output_start + output_chunk_size, self.out_features) + output_size = output_end - output_start + if output_size % 64 != 0: + raise ValueError("QQQ pack output chunk is not permutation aligned") + + weight_chunk = weight[output_start:output_end, input_start:input_end].transpose(0, 1) + scale_chunk_view = scale_chunk[:, output_start:output_end] + codes = torch.round(weight_chunk / scale_chunk_view).to(dtype=torch.int32) + if self.group_size != self.in_features: + codes.add_((self.maxq + 1) // 2).clamp_(0, self.maxq) + else: + codes.clamp_(-self.maxq, self.maxq) + + transformed = codes.reshape( + input_size // self.tile, + self.tile, + output_size // self.tile, + self.tile, + ).permute(0, 2, 1, 3).reshape( + input_size // self.tile, + output_size * self.tile, + ) + transformed = transformed.reshape(-1, perm.numel()).index_select(1, perm).reshape( + input_size // self.tile, + output_size * self.tile, + ) + packed_chunk = torch.zeros( + (input_size // self.tile, output_size * 2), + dtype=torch.int32, + device=weight.device, + ) + for lane in range(8): + packed_chunk.bitwise_or_( + (transformed[:, lane::8] & 0xF) << (4 * lane) + ) - #self.B[:, :] = q.to(self.B.device) - self.register_buffer("B", q) + packed_weight[ + input_start // self.tile:input_end // self.tile, + output_start * 2:output_end * 2, + ].copy_(packed_chunk.to(device=CPU)) + self.register_buffer("B", packed_weight) if self.group_size != self.in_features: - #self.s_group[:, :] = s.to(self.s_group.device) - self.register_buffer("s_group", s.to(CPU)) - - #self.s_channel[:, :] = s_extra.to(self.s_channel.device) - self.register_buffer("s_channel", s_extra.to(CPU)) + self.register_buffer("s_group", packed_s_group) + self.register_buffer("s_channel", packed_s_channel) else: - # self.s_group = torch.tensor( - # [], dtype=torch.float16, device=self.s_channel.device - # ) - # self.register_buffer("s_group", torch.tensor( - # [], dtype=torch.float16, device=CPU - # )) - - #self.s_channel[:, :] = s.to(self.s_channel.device) - self.register_buffer("s_channel", s.to(CPU)) - if linear.bias is not None: - if self.bias is not None: - # self.bias[:] = linear.bias.data.to(self.bias.device).to(torch.float16) - self.register_buffer("bias", linear.bias.data.to(self.bias.device).to(torch.float16)) - # else: - # # self.bias = linear.bias.clone().to(torch.float16) - # self.register_buffer("bias", linear.bias.clone().to(torch.float16)) + self.register_buffer("s_channel", packed_s_channel) + if linear.bias is not None and self.bias is not None: + self.register_buffer("bias", linear.bias.data.to(self.bias.device).to(torch.float16)) # activation int8 quantization def dynamic_quant(self, x: torch.Tensor): From 4e3510af547ac8db725fc0f7b6b6df43950bbd7b Mon Sep 17 00:00:00 2001 From: "yujiong.zhang" Date: Tue, 25 Aug 2026 13:58:22 +0800 Subject: [PATCH 3/3] add qqq jit ut --- tests/test_qqq_jit.py | 107 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) diff --git a/tests/test_qqq_jit.py b/tests/test_qqq_jit.py index 5e216d858..6352bb82b 100644 --- a/tests/test_qqq_jit.py +++ b/tests/test_qqq_jit.py @@ -75,3 +75,110 @@ def test_qqq_forward_raises_runtime_error_when_jit_ops_missing(monkeypatch): with pytest.raises(ModuleNotFoundError, match="missing qqq jit ops"): module(torch.randn((1, module.in_features), dtype=torch.float16)) + + +def _build_grouped_parity_modules( + cases: list[tuple[int, float, int]], +) -> tuple[qqq_module.QQQLinear, qqq_module.QQQTorchLinear]: + in_features = 256 + out_features = 128 + group_size = 128 + linear = torch.nn.Linear(in_features, out_features, bias=False, dtype=torch.float16) + linear.weight.data.zero_() + scales = torch.ones((out_features, in_features // group_size), dtype=torch.float16) + + for output_index, (code, scale, _) in enumerate(cases): + scales[output_index, 0] = scale + linear.weight.data[output_index, 0] = (code - 8) * scale + + s_channel = torch.ones(out_features, dtype=torch.float32) + modules = [] + for module_cls in (qqq_module.QQQLinear, qqq_module.QQQTorchLinear): + module = module_cls( + bits=4, + group_size=group_size, + sym=True, + desc_act=False, + in_features=in_features, + out_features=out_features, + bias=False, + register_buffers=True, + ) + module.pack(linear, scales, s_channel) + module.post_init() + modules.append(module.eval()) + + return modules[0].cuda(), modules[1] + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qqq_grouped_cuda_matches_torch_rounding_and_saturation(): + if torch.cuda.get_device_capability()[0] < 8: + pytest.skip("QQQ CUDA requires compute capability >= 8.0") + if not qqq_module.qqq_runtime_available(): + pytest.skip(qqq_module.qqq_runtime_error()) + + cases = [ + (15, 22.0, 127), + (0, 20.0, -128), + (15, 18.0, 126), + (9, 0.5, 0), + (11, 0.5, 2), + (13, 0.5, 2), + (7, 0.5, 0), + (5, 0.5, -2), + (3, 0.5, -2), + ] + cuda_module, torch_module = _build_grouped_parity_modules(cases) + x = torch.zeros((1, cuda_module.in_features), dtype=torch.float16) + x[0, 0] = 1.0 + + expected = torch.tensor([case[2] for case in cases], dtype=torch.float16) + cuda_output = cuda_module(x.cuda()).cpu()[0, : len(cases)] + torch_output = torch_module(x)[0, : len(cases)] + + torch.testing.assert_close(torch_output, expected, rtol=0, atol=0) + torch.testing.assert_close(cuda_output, expected, rtol=0, atol=0) + + +@pytest.mark.cuda +@pytest.mark.parametrize("tokens", [1, 17]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_qqq_grouped_cuda_matches_torch_for_regular_values(tokens, dtype): + if torch.cuda.get_device_capability()[0] < 8: + pytest.skip("QQQ CUDA requires compute capability >= 8.0") + if not qqq_module.qqq_runtime_available(): + pytest.skip(qqq_module.qqq_runtime_error()) + + torch.manual_seed(42) + in_features = 256 + out_features = 128 + linear = torch.nn.Linear(in_features, out_features, bias=False, dtype=torch.float16) + linear.weight.data.normal_(0, 0.08) + scales = torch.rand((out_features, 2), dtype=torch.float16) * 0.02 + 0.01 + s_channel = torch.rand(out_features, dtype=torch.float32) * 0.5 + 0.75 + + modules = [] + for module_cls in (qqq_module.QQQLinear, qqq_module.QQQTorchLinear): + module = module_cls( + bits=4, + group_size=128, + sym=True, + desc_act=False, + in_features=in_features, + out_features=out_features, + bias=False, + register_buffers=True, + ) + module.pack(linear, scales, s_channel) + module.post_init() + modules.append(module.eval()) + + x = torch.randn((tokens, in_features), dtype=dtype) + cuda_output = modules[0].cuda()(x.cuda()).cpu() + torch_output = modules[1](x) + + assert cuda_output.dtype == dtype + torch.testing.assert_close(cuda_output, torch_output, rtol=0.02, atol=0.02)