A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 06:15 UTC
[Bugfix] Handle Asym W4A16 (ConchLinearKernel) for CT (#33200)
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
Matthias Gehre committed
Jan 31, 2026 at 07:21 UTC
73419abfae97c0b797c7e1fe997913e9ebfeff68
1 file changed
+42
vllm/model_executor/layers/quantization/kernels/mixed_precision/conch.py
+42
@@ -56,6 +56,7 @@ class ConchLinearKernel(MPLinearKernel):
56
# note assumes that
57
# `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
58
# `weight_scale` is: {input_dim = 0, output_dim = 1}
59
+ # `weight_zero_point` is: {input_dim = 1, output_dim = 0, packed_dim = 0}
60
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
61
def transform_w_q(x):
62
assert isinstance(x, BasevLLMParameter)
@@ -69,8 +70,49 @@ class ConchLinearKernel(MPLinearKernel):
70
x.data = x.data.contiguous()
71
return x
72
73
+ def transform_w_zp(x):
74
+ # Zero points are stored PACKED as [N//pack_factor, K//G]
75
+ # The Conch kernel expects UNPACKED zeros: [K//G, N]
76
+ # We need to unpack and reorder
77
+ assert isinstance(x, BasevLLMParameter)
78
+ packed = x.data # shape: [N//pack_factor, K//G], dtype: int32
79
+
80
+ # Determine packing based on weight bit width
81
+ size_bits = self.config.weight_type.size_bits
82
+ pack_factor = 32 // size_bits # 8 for 4-bit, 4 for 8-bit
83
+ mask = (1 << size_bits) - 1 # 0xF for 4-bit, 0xFF for 8-bit
84
+
85
+ n_packed, k_groups = packed.shape
86
+ n_full = n_packed * pack_factor
87
+
88
+ # Unpack using vectorized bitwise ops
89
+ # shifts = [0, size_bits, 2*size_bits, ...] for each packed position
90
+ shifts = torch.arange(
91
+ 0, 32, size_bits, dtype=torch.int32, device=packed.device
92
+ )
93
+ # packed: [N//pack_factor, K//G] -> [N//pack_factor, K//G, 1]
94
+ # shifts: [pack_factor] -> [1, 1, pack_factor]
95
+ # Result: [N//pack_factor, K//G, pack_factor]
96
+ unpacked = (packed.unsqueeze(-1) >> shifts) & mask
97
+
98
+ # Permute to [K//G, N//pack_factor, pack_factor] then reshape to [K//G, N]
99
+ unpacked = unpacked.permute(1, 0, 2).reshape(k_groups, n_full)
100
+
101
+ x.data = unpacked.to(torch.uint8).contiguous()
102
+
103
+ # Update metadata - zeros are no longer packed
104
+ if hasattr(x, "_input_dim"):
105
+ x._input_dim = 0
106
+ if hasattr(x, "_output_dim"):
107
+ x._output_dim = 1
108
+ if hasattr(x, "_packed_factor"):
109
+ x._packed_factor = 1
110
+ return x
111
+
112
self._transform_param(layer, self.w_q_name, transform_w_q)
113
self._transform_param(layer, self.w_s_name, transform_w_s)
114
+ if self.config.zero_points:
115
+ self._transform_param(layer, self.w_zp_name, transform_w_zp)
116
117
def apply_weights(
118
self,