[Core]refactor aqlm quant ops (#4351)

2024-04-25 19:03:56 +00:00
parent bd7a8eef25
commit f4bc4de1b1
3 changed files with 16 additions and 2 deletions
--- a/benchmarks/kernels/benchmark_aqlm.py
+++ b/benchmarks/kernels/benchmark_aqlm.py
@@ -6,7 +6,7 @@ from typing import Optional
 import torch
 import torch.nn.functional as F

-from vllm._C import ops
+from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.aqlm import (
    dequantize_weight, generic_dequantize_gemm, get_int_dtype,
    optimized_dequantize_gemm)