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 5, 23:30 UTC
[CI] Fix conch kernel crash on 3D input by reshaping to 2D before GEMM (#38178)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 26, 2026 at 11:46 UTC
a8e48a7b85f035df401cab03dfb11b16f98aa413
1 file changed
+5
-2
vllm/model_executor/kernels/linear/mixed_precision/conch.py
+5
-2
@@ -134,8 +134,11 @@ class ConchLinearKernel(MPLinearKernel):
134
if group_size == -1:
135
group_size = x.shape[-1]
136
137
+ x_2d = x.reshape(-1, x.shape[-1])
138
+ out_shape = x.shape[:-1] + (self.config.partition_weight_shape[1],)
139
+
140
output = mixed_precision_gemm(
138
- x=x,
141
+ x=x_2d,
142
w_q_packed=w_q.data,
143
w_s=w_s.data,
144
w_zp=w_zp.data if w_zp is not None else None,
@@ -147,4 +150,4 @@ class ConchLinearKernel(MPLinearKernel):
150
if bias is not None:
151
output.add_(bias) # In-place add
152
150
- return output
153
+ return output.reshape(out_shape)