@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 90969fb39a
Fork 0

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 15, 16:00 UTC

[Kernel] Add more dtype support for GGUF dequantization (#15879)

Signed-off-by: lukas.bluebaum <lukas.bluebaum@aleph-alpha.com>

LukasBluebaum committed Apr 2, 2025 at 10:58 UTC 90969fb39a58593515f6a087d9200bc72333ab9a
9 files changed +80 -50
csrc/ops.h
+2 -1
@@ -145,7 +145,8 @@ torch::Tensor permute_cols(torch::Tensor const& A, torch::Tensor const& perm);
145 #endif
146
147 torch::Tensor ggml_dequantize(torch::Tensor W, int64_t type, int64_t m,
148 - int64_t n);
148 + int64_t n,
149 + std::optional<at::ScalarType> const& dtype);
150
151 torch::Tensor ggml_mul_mat_vec_a8(torch::Tensor W, torch::Tensor X,
152 int64_t type, int64_t row);
csrc/quantization/gguf/dequantize.cuh
+34 -31
@@ -94,8 +94,8 @@ static __global__ void dequantize_block(const void * __restrict__ vx, dst_t * __
94 dfloat2 v;
95 dequantize_kernel(vx, ib, iqs, v);
96
97 - y[iybs + iqs + 0] = v.x;
98 - y[iybs + iqs + y_offset] = v.y;
97 + y[iybs + iqs + 0] = convert_from_half<dst_t>(v.x);
98 + y[iybs + iqs + y_offset] = convert_from_half<dst_t>(v.y);
99 }
100
101 template<typename dst_t>
@@ -114,10 +114,10 @@ static __global__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t
114
115 half dall = __low2half(x[i].dm);
116 half dmin = __high2half(x[i].dm);
117 - y[l+ 0] = __hsub(__hmul(dall, __int2half_rn((x[i].scales[is+0] & 0xF) * ((q >> 0) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+0] >> 4)));
118 - y[l+32] = __hsub(__hmul(dall, __int2half_rn((x[i].scales[is+2] & 0xF) * ((q >> 2) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+2] >> 4)));
119 - y[l+64] = __hsub(__hmul(dall, __int2half_rn((x[i].scales[is+4] & 0xF) * ((q >> 4) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+4] >> 4)));
120 - y[l+96] = __hsub(__hmul(dall, __int2half_rn((x[i].scales[is+6] & 0xF) * ((q >> 6) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+6] >> 4)));
117 + y[l+ 0] = convert_from_half<dst_t>(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+0] & 0xF) * ((q >> 0) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+0] >> 4))));
118 + y[l+32] = convert_from_half<dst_t>(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+2] & 0xF) * ((q >> 2) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+2] >> 4))));
119 + y[l+64] = convert_from_half<dst_t>(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+4] & 0xF) * ((q >> 4) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+4] >> 4))));
120 + y[l+96] = convert_from_half<dst_t>(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+6] & 0xF) * ((q >> 6) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+6] >> 4))));
121 }
122
123 template<typename dst_t>
@@ -148,7 +148,9 @@ static __global__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t
148 const uint8_t * q = x[i].qs + 32*n;
149 const uint8_t * hm = x[i].hmask;
150
151 - for (int l = l0; l < l0+4; ++l) y[l] = __hmul(dl, __int2half_rn((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)));
151 + for (int l = l0; l < l0+4; ++l) {
152 + y[l] = convert_from_half<dst_t>(__hmul(dl, __int2half_rn((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4))));
153 + }
154 }
155
156 static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) {
@@ -188,8 +190,8 @@ static __global__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t
190 const half d2 = __hmul(dall, __int2half_rn(sc));
191 const half m2 = __hmul(dmin, __int2half_rn(m));
192 for (int l = 0; l < n; ++l) {
191 - y[l + 0] = __hsub(__hmul(d1, __int2half_rn(q[l] & 0xF)), m1);
192 - y[l +32] = __hsub(__hmul(d2, __int2half_rn(q[l] >> 4)), m2);
193 + y[l + 0] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn(q[l] & 0xF)), m1));
194 + y[l +32] = convert_from_half<dst_t>(__hsub(__hmul(d2, __int2half_rn(q[l] >> 4)), m2));
195 }
196 }
197
@@ -220,11 +222,11 @@ static __global__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t
222 const half d2 = __hmul(dall, __int2half_rn(sc)); const half m2 = __hmul(dmin, __int2half_rn(m));
223
224 uint8_t hm = 1 << (2*il);
223 - y[ 0] = __hsub(__hmul(d1, __int2half_rn((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0))), m1);
224 - y[ 1] = __hsub(__hmul(d1, __int2half_rn((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0))), m1);
225 + y[ 0] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0))), m1));
226 + y[ 1] = convert_from_half<dst_t>(__hsub(__hmul(d1, __int2half_rn((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0))), m1));
227 hm <<= 1;
226 - y[32] = __hsub(__hmul(d2, __int2half_rn((ql[0] >> 4) + (qh[0] & hm ? 16 : 0))), m2);
227 - y[33] = __hsub(__hmul(d2, __int2half_rn((ql[1] >> 4) + (qh[1] & hm ? 16 : 0))), m2);
228 + y[32] = convert_from_half<dst_t>(__hsub(__hmul(d2, __int2half_rn((ql[0] >> 4) + (qh[0] & hm ? 16 : 0))), m2));
229 + y[33] = convert_from_half<dst_t>(__hsub(__hmul(d2, __int2half_rn((ql[1] >> 4) + (qh[1] & hm ? 16 : 0))), m2));
230 }
231
232 template<typename dst_t>
@@ -247,10 +249,10 @@ static __global__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t
249 const uint8_t qh = x[i].qh[32*ip + il];
250 const int8_t * sc = x[i].scales + is;
251
250 - y[ 0] = __hmul(d, __int2half_rn(sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32)));
251 - y[32] = __hmul(d, __int2half_rn(sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32)));
252 - y[64] = __hmul(d, __int2half_rn(sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32)));
253 - y[96] = __hmul(d, __int2half_rn(sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32)));
252 + y[ 0] = convert_from_half<dst_t>(__hmul(d, __int2half_rn(sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32))));
253 + y[32] = convert_from_half<dst_t>(__hmul(d, __int2half_rn(sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32))));
254 + y[64] = convert_from_half<dst_t>(__hmul(d, __int2half_rn(sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32))));
255 + y[96] = convert_from_half<dst_t>(__hmul(d, __int2half_rn(sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32))));
256 }
257
258 template<typename dst_t>
@@ -269,7 +271,7 @@ static __global__ void dequantize_block_iq2_xxs(const void * __restrict__ vx, ds
271 const uint32_t aux32 = q2[2] | (q2[3] << 16);
272 const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.25f;
273 const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
272 - for (int j = 0; j < 8; ++j) y[j] = __float2half(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
274 + for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
275 }
276
277 template<typename dst_t>
@@ -286,7 +288,7 @@ static __global__ void dequantize_block_iq2_xs(const void * __restrict__ vx, dst
288 const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511));
289 const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
290 const uint8_t signs = ksigns_iq2xs[q2[il] >> 9];
289 - for (int j = 0; j < 8; ++j) y[j] = __float2half(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
291 + for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
292
293 }
294
@@ -303,7 +305,7 @@ static __global__ void dequantize_block_iq2_s(const void * __restrict__ vx, dst_
305 const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[i].qs[4*ib+il] | ((x[i].qh[ib] << (8-2*il)) & 0x300)));
306 const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f;
307 const uint8_t signs = x[i].qs[QK_K/8+4*ib+il];
306 - for (int j = 0; j < 8; ++j) y[j] = __float2half(d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f));
308 + for (int j = 0; j < 8; ++j) y[j] = d * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
309 }
310
311 template<typename dst_t>
@@ -324,8 +326,8 @@ static __global__ void dequantize_block_iq3_xxs(const void * __restrict__ vx, ds
326 const float d = __half2float(x[i].d) * (0.5f + (aux32 >> 28)) * 0.5f;
327 const uint8_t signs = ksigns_iq2xs[(aux32 >> 7*il) & 127];
328 for (int j = 0; j < 4; ++j) {
327 - y[j+0] = __float2half(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
328 - y[j+4] = __float2half(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
329 + y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
330 + y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
331 }
332 }
333
@@ -345,8 +347,8 @@ static __global__ void dequantize_block_iq3_s(const void * __restrict__ vx, dst_
347 const float d = __half2float(x[i].d) * (0.5f + ((x[i].scales[ib/2] >> 4*(ib%2)) & 0xf)) * 0.5f;
348 const uint8_t signs = x[i].signs[4*ib + il];
349 for (int j = 0; j < 4; ++j) {
348 - y[j+0] = __float2half(d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f));
349 - y[j+4] = __float2half(d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f));
350 + y[j+0] = d * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
351 + y[j+4] = d * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
352 }
353 }
354
@@ -367,7 +369,7 @@ static __global__ void dequantize_block_iq1_s(const void * __restrict__ vx, dst_
369 grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
370 grid32[0] &= 0x0f0f0f0f;
371 for (int j = 0; j < 8; ++j) {
370 - y[j] = __float2half(d * (q[j] + delta));
372 + y[j] = d * (q[j] + delta);
373 }
374 }
375
@@ -392,7 +394,7 @@ static __global__ void dequantize_block_iq1_m(const void * __restrict__ vx, dst_
394 grid32[1] = (grid32[0] >> 4) & 0x0f0f0f0f;
395 grid32[0] &= 0x0f0f0f0f;
396 for (int j = 0; j < 8; ++j) {
395 - y[j] = __float2half(d * (q[j] + delta));
397 + y[j] = d * (q[j] + delta);
398 }
399 }
400
@@ -409,8 +411,8 @@ static __global__ void dequantize_block_iq4_nl(const void * __restrict__ vx, dst
411 const uint8_t * q4 = x[ib].qs + 4*il;
412 const float d = __half2float(x[ib].d);
413 for (int j = 0; j < 4; ++j) {
412 - y[j+ 0] = __float2half(d * kvalues_iq4nl[q4[j] & 0xf]);
413 - y[j+16] = __float2half(d * kvalues_iq4nl[q4[j] >> 4]);
414 + y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf];
415 + y[j+16] = d * kvalues_iq4nl[q4[j] >> 4];
416 }
417
418 }
@@ -427,8 +429,8 @@ static __global__ void dequantize_block_iq4_xs(const void * __restrict__ vx, dst
429 const uint8_t * q4 = x[i].qs + 16*ib + 4*il;
430 const float d = __half2float(x[i].d) * ((((x[i].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[i].scales_h >> 2*ib) & 3) << 4)) - 32);
431 for (int j = 0; j < 4; ++j) {
430 - y[j+ 0] = __float2half(d * kvalues_iq4nl[q4[j] & 0xf]);
431 - y[j+16] = __float2half(d * kvalues_iq4nl[q4[j] >> 4]);
432 + y[j+ 0] = d * kvalues_iq4nl[q4[j] & 0xf];
433 + y[j+16] = d * kvalues_iq4nl[q4[j] >> 4];
434 }
435 }
436
@@ -522,7 +524,8 @@ static void dequantize_row_iq4_xs_cuda(const void * vx, dst_t * y, const int k,
524 dequantize_block_iq4_xs<<<nb, 32, 0, stream>>>(vx, y);
525 }
526
525 -static to_fp16_cuda_t ggml_get_to_fp16_cuda(int64_t type) {
527 +template<typename dst_t>
528 +static to_cuda_ggml_t<dst_t> ggml_get_to_cuda(int64_t type) {
529 switch (type) {
530 case 2:
531 return dequantize_block_cuda<QK4_0, QR4_0, dequantize_q4_0>;
csrc/quantization/gguf/ggml-common.h
+16 -1
@@ -1063,7 +1063,8 @@ static const __device__ int8_t kvalues_iq4nl[16] = {-127, -104, -83, -65, -49, -
1063 typedef half dfloat; // dequantize float
1064 typedef half2 dfloat2;
1065 typedef void (*dequantize_kernel_t)(const void * vx, const int ib, const int iqs, dfloat2 & v);
1066 -typedef void (*to_fp16_cuda_t)(const void * __restrict__ x, dfloat * __restrict__ y, int k, cudaStream_t stream);
1066 +template<typename dst_t>
1067 +using to_cuda_ggml_t = void (*)(const void * __restrict__ x, dst_t * __restrict__ y, int k, cudaStream_t stream);
1068 typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs);
1069 typedef void (*allocate_tiles_cuda_t)(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc);
1070 typedef void (*load_tiles_cuda_t)(
@@ -1075,6 +1076,20 @@ typedef float (*vec_dot_q_mul_mat_cuda_t)(
1076
1077 // Utility function
1078
1079 +template<typename dst_t>
1080 +static __device__ __forceinline__ dst_t convert_from_half(half val) {
1081 + return val;
1082 +}
1083 +
1084 +template<>
1085 +__device__ __forceinline__ c10::BFloat16 convert_from_half<c10::BFloat16>(half val) {
1086 +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
1087 + return __float2bfloat16(__half2float(val));
1088 +#else
1089 + return __half2float(val);
1090 +#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
1091 +}
1092 +
1093 #if defined(USE_ROCM)
1094
1095 #ifndef __has_builtin
csrc/quantization/gguf/gguf_kernel.cu
+10 -5
@@ -71,14 +71,19 @@ static void quantize_row_q8_1_cuda(const scalar_t* x, void* vy, const int kx,
71 }
72
73 torch::Tensor ggml_dequantize(torch::Tensor W, // quant weight
74 - int64_t type, int64_t m, int64_t n) {
74 + int64_t type, int64_t m, int64_t n,
75 + std::optional<at::ScalarType> const& dtype) {
76 const at::cuda::OptionalCUDAGuard device_guard(device_of(W));
76 - auto options =
77 - torch::TensorOptions().dtype(torch::kFloat16).device(W.device());
77 + auto dtype_ = dtype.value_or(torch::kFloat16);
78 + auto options = torch::TensorOptions().dtype(dtype_).device(W.device());
79 at::Tensor DW = torch::empty({m, n}, options);
80 cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
80 - const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(type);
81 - to_fp16_cuda((void*)W.data_ptr(), (half*)DW.data_ptr(), m * n, stream);
81 +
82 + VLLM_DISPATCH_FLOATING_TYPES(DW.scalar_type(), "ggml_dequantize", [&] {
83 + auto to_cuda = ggml_get_to_cuda<scalar_t>(type);
84 + to_cuda((void*)W.data_ptr(), (scalar_t*)DW.data_ptr(), m * n, stream);
85 + });
86 +
87 return DW;
88 }
89
csrc/torch_bindings.cpp
+3 -1
@@ -295,7 +295,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
295 #endif
296
297 // Dequantization for GGML.
298 - ops.def("ggml_dequantize(Tensor W, int type, SymInt m, SymInt n) -> Tensor");
298 + ops.def(
299 + "ggml_dequantize(Tensor W, int type, SymInt m, SymInt n, ScalarType? "
300 + "dtype) -> Tensor");
301 ops.impl("ggml_dequantize", torch::kCUDA, &ggml_dequantize);
302
303 // mmvq kernel for GGML.
tests/kernels/test_ggml.py
+2 -1
@@ -15,7 +15,8 @@ def test_ggml_opcheck(quant_type):
15 qweight = torch.randint(0, 100, shape, device='cuda', dtype=torch.uint8)
16 m = qweight.shape[0]
17 n = qweight.shape[1] // type_size * block_size
18 - opcheck(torch.ops._C.ggml_dequantize, (qweight, quant_type, m, n))
18 + opcheck(torch.ops._C.ggml_dequantize,
19 + (qweight, quant_type, m, n, torch.float16))
20
21 x = torch.rand((m, 512), device='cuda', dtype=torch.float16)
22 opcheck(torch.ops._C.ggml_mul_mat_a8,
tests/kernels/test_gguf.py
+2 -2
@@ -65,7 +65,7 @@ QUANT_TYPES = [
65
66
67 @pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
68 -@pytest.mark.parametrize("dtype", [torch.half])
68 +@pytest.mark.parametrize("dtype", DTYPES)
69 @pytest.mark.parametrize("quant_type", QUANT_TYPES)
70 @torch.inference_mode()
71 def test_dequantize(hidden_size: int, dtype: torch.dtype,
@@ -78,7 +78,7 @@ def test_dequantize(hidden_size: int, dtype: torch.dtype,
78 ref_output = torch.tensor(dequantize(tensor.data, quant_type),
79 device="cuda").to(dtype)
80 output = ops.ggml_dequantize(torch.tensor(tensor.data, device="cuda"),
81 - quant_type, *list(shape)).to(dtype)
81 + quant_type, *list(shape), dtype)
82
83 torch.testing.assert_close(output, ref_output, atol=1e-2, rtol=4e-2)
84
vllm/_custom_ops.py
+9 -6
@@ -436,9 +436,12 @@ if hasattr(torch.ops._C, "allspark_w8a16_gemm"):
436 if hasattr(torch.ops._C, "ggml_dequantize"):
437
438 @register_fake("_C::ggml_dequantize")
439 - def _ggml_dequantize_fake(W: torch.Tensor, quant_type: int,
440 - m: torch.SymInt,
441 - n: torch.SymInt) -> torch.Tensor:
439 + def _ggml_dequantize_fake(
440 + W: torch.Tensor,
441 + quant_type: int,
442 + m: torch.SymInt,
443 + n: torch.SymInt,
444 + dtype: Optional[torch.dtype] = None) -> torch.Tensor:
445 return torch.empty((m, n), dtype=torch.float16, device=W.device)
446
447 @register_fake("_C::ggml_mul_mat_vec_a8")
@@ -1097,9 +1100,9 @@ def marlin_qqq_gemm(a: torch.Tensor, b_q_weight: torch.Tensor,
1100
1101
1102 # gguf
1100 -def ggml_dequantize(W: torch.Tensor, quant_type: int, m: int,
1101 - n: int) -> torch.Tensor:
1102 - return torch.ops._C.ggml_dequantize(W, quant_type, m, n)
1103 +def ggml_dequantize(W: torch.Tensor, quant_type: int, m: int, n: int,
1104 + dtype: Optional[torch.dtype]) -> torch.Tensor:
1105 + return torch.ops._C.ggml_dequantize(W, quant_type, m, n, dtype)
1106
1107
1108 def ggml_mul_mat_vec_a8(
vllm/model_executor/layers/quantization/gguf.py
+2 -2
@@ -117,7 +117,7 @@ def _fuse_mul_mat(x: torch.Tensor, qweight: torch.Tensor,
117 elif qweight_type in DEQUANT_TYPES:
118 block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type]
119 shape = (qweight.shape[0], qweight.shape[1] // type_size * block_size)
120 - weight = ops.ggml_dequantize(qweight, qweight_type, *shape)
120 + weight = ops.ggml_dequantize(qweight, qweight_type, *shape, x.dtype)
121 y = x @ weight.T
122 else:
123 # Raise an error if the quantization type is not supported.
@@ -377,7 +377,7 @@ class GGUFEmbeddingMethod(GGUFLinearMethod):
377 x_flat = x.flatten()
378 quant = torch.index_select(qweight, dim=0, index=x_flat)
379 dequant = ops.ggml_dequantize(quant, qweight_type, hidden_size,
380 - x_flat.shape[0]).to(self.params_dtype)
380 + x_flat.shape[0], self.params_dtype)
381 return dequant.view(*x.shape, hidden_size)
382
383