提交vllm0.11.0开发分支

vllm_kunlun/ops/quantization/awq.py

@@ -1,128 +0,0 @@
-#
-# Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
-# Author: Li Wei, Pan Xiakai, You Zeyu
-# Email: liwei157@baidu.com
-# This file is a part of the vllm-kunlun project.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-
-from typing import Optional
-from vllm.model_executor.layers.quantization.awq import AWQLinearMethod
-
-
-def repack_int4_for_kunlun(self, packed: torch.Tensor, num_bits: int = 4):
-    """Convert AWQ-packed int4 weights to Kunlun XPU format.
-    Input:  packed[N, K], dtype=int32, saved as AWQ order
-    Output: packed_reordered[N, K], dtype=int32, saved as Kunlun order
-    """
-    N, K = packed.shape
-    self.align_type = 1 if K % 8 == 0 else 0
-    assert num_bits == 4, "Only int4 supported now"
-    shifts = torch.arange(0, 32, num_bits, device=packed.device, dtype=torch.int32)
-
-    if self.align_type == 0:  # NORMAL MODE
-        # Unpack AWQ order:[0, 2, 4, 6, 1, 3, 5, 7]
-        unpacked_awq = (packed.unsqueeze(-1) >> shifts) & 0xF  # [N, K, 8]
-
-        # Reverse AWQ order and convert to KUNLUN order
-        AWQ_TO_KUNLUN_ORDER_NORMAL = [4, 0, 5, 1, 6, 2, 7, 3]
-        # [0,2,4,6,1,3,5,7] --> [1, 0, 3, 2, 5, 4, 7, 6]
-        unpacked_kunlun = unpacked_awq[..., AWQ_TO_KUNLUN_ORDER_NORMAL]  # [N, K, 8]
-
-        # Pack to int32, order[6, 7, 4, 5, 2, 3, 0, 1]
-        packed_kunlun = (unpacked_kunlun << shifts).sum(
-            dim=-1, dtype=torch.int32
-        )  # [N, K]
-    elif self.align_type == 1:  # FAST MODEL
-        # Unpack AWQ order
-        unpacked_awq = (
-            packed.view(N, K // 8, 8).unsqueeze(-1) >> shifts
-        ) & 0xF  # [N, K//8, 8, 8]
-
-        # Reverse AWQ order and convert to KUNLUN order
-        AWQ_TO_KUNLUN_ORDER_FAST = [
-            32, 0, 36, 4, 33, 1, 37, 5,
-            34, 2, 38, 6, 35, 3, 39, 7,
-            40, 8, 44, 12, 41, 9, 45, 13,
-            42, 10, 46, 14, 43, 11, 47, 15,
-            48, 16, 52, 20, 49, 17, 53, 21,
-            50, 18, 54, 22, 51, 19, 55, 23,
-            56, 24, 60, 28, 57, 25, 61, 29,
-            58, 26, 62, 30, 59, 27, 63, 31
-        ]
-        unpacked_awq = unpacked_awq.reshape(N, K // 8, 64)
-        unpacked_kunlun = unpacked_awq[..., AWQ_TO_KUNLUN_ORDER_FAST]  # [N, K//8, 64]
-
-        # Pack to int32
-        unpacked_kunlun = unpacked_kunlun.reshape(N, K // 8, 8, 8)
-        packed_kunlun = (
-            (unpacked_kunlun << shifts).sum(dim=-1, dtype=torch.int32).reshape(N, K)
-        )  # [N, K]
-    else:
-        raise NotImplementedError
-
-    return packed_kunlun
-
-
-def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-    layer.qweight = torch.nn.Parameter(
-        (
-            self.repack_int4_for_kunlun(layer.qweight.data)
-            if layer.qweight.data.dtype == torch.int32
-            else layer.qweight.data
-        ),
-        requires_grad=False,
-    )
-    layer.qzeros = torch.nn.Parameter(
-        (
-            self.repack_int4_for_kunlun(layer.qzeros.data)
-            if layer.qzeros.data.dtype == torch.int32
-            else layer.qzeros.data
-        ),
-        requires_grad=False,
-    )
-    layer.scales = torch.nn.Parameter(layer.scales.data, requires_grad=False)
-
-
-def apply(
-    self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None
-) -> torch.Tensor:
-    qweight = layer.qweight
-    scales = layer.scales
-    qzeros = layer.qzeros
-    pack_factor = self.quant_config.pack_factor
-    out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
-    reshaped_x = x.reshape(-1, x.shape[-1])
-
-    # num_tokens >= threshold
-    FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
-
-    if FP16_MATMUL_HEURISTIC_CONDITION:
-        out = torch.ops._C.awq_dequantize(
-            qweight, scales, qzeros, quant_type=0, align_type=self.align_type
-        )
-        out = torch.matmul(reshaped_x, out)
-    else:
-        out = torch.ops._C.awq_gemm(
-            reshaped_x, qweight, scales, qzeros, align_type=self.align_type
-        )
-    if bias is not None:
-        out.add_(bias)
-    return out.reshape(out_shape)
-
-
-AWQLinearMethod.repack_int4_for_kunlun = repack_int4_for_kunlun
-AWQLinearMethod.process_weights_after_loading = process_weights_after_loading
-AWQLinearMethod.apply = apply

vllm_kunlun/ops/quantization/compressed_tensors_moe.py

@@ -1,37 +1,14 @@
-#
-# Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
-#
-# This file is a part of the vllm-kunlun project.
-# Author: Chen Zhennan, Dong Xinyu
-# Email: chenzhennan@baidu.com
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import torch
 from typing import Any, Literal, Optional, cast, Callable, Optional
 
-from compressed_tensors.config import (
-    CompressionFormat,
-    SparsityCompressionConfig,
-    SparsityStructure,
-)
-from compressed_tensors.quantization import ActivationOrdering, QuantizationStrategy
-from vllm.model_executor.layers.fused_moe import (
-    FusedMoE,
-    FusedMoEMethodBase,
-    FusedMoeWeightScaleSupported,
-)
+from compressed_tensors.config import (CompressionFormat,
+                                       SparsityCompressionConfig,
+                                       SparsityStructure)
+from compressed_tensors.quantization import (ActivationOrdering,
+                                             QuantizationStrategy)
+from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
+                                                  FusedMoeWeightScaleSupported)
 from vllm.model_executor.layers.quantization.utils import replace_parameter
-
 # TODO: import position will be changed after 0.9.0
 # vllm.model_executor.layers.fused_moe.fused_moe --> vllm.model_executor.layers.fused_moe
 
@@ -42,7 +19,6 @@ import xtorch_ops
 
 from safetensors.torch import load_file as safe_load_file
 
-
 class CompressedTensorsMoEMethod(FusedMoEMethodBase):
 
     def get_moe_method(quant_config, layer) -> "CompressedTensorsMoEMethod":
@@ -50,239 +26,177 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
         linear_cfg = None
         for k in ("Linear", "FusedMoE", "MoE", "Moe", "Experts"):
             if k in tsm and isinstance(tsm[k], dict):
-                linear_cfg = tsm[k]
-                break
+                linear_cfg = tsm[k]; break
         if not linear_cfg:
             # print("target_scheme_map missing; fallback to INT8(W8A8) method")
             return CompressedTensorsW8A8Int8MoEMethod(quant_config)
-        wq = linear_cfg.get("weights")
-        aq = linear_cfg.get("input_activations")
+        wq = linear_cfg.get("weights"); aq = linear_cfg.get("input_activations")
         if not wq or not aq:
             # print("incomplete scheme; fallback to INT8(W8A8)")
             return CompressedTensorsW8A8Int8MoEMethod(quant_config)
-
-        # Other branches are handled as needed; default fallback:
+        # 其它分流按需；默认回落：
         return CompressedTensorsW8A8Int8MoEMethod(quant_config)
 
-
 # copied from vllm 0.9.0
 class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
 
     def __init__(
-        self, quant_config: "CompressedTensorsConfig"  # type: ignore # noqa E501
+            self,
+            quant_config: "CompressedTensorsConfig"  # type: ignore # noqa E501
     ):
         self.quant_config = quant_config
-
-        # Directly create a default quantization config dictionary to avoid validation issues with QuantizationArgs
+        
+        # 直接创建默认的量化配置字典，避免 QuantizationArgs 的验证问题
         # print("Creating default INT8 quantization config for MoE")
+        
+        # 创建默认的权重量化配置字典
+        self.weight_quant = type('WeightQuant', (), {
+            'type': 'int',
+            'num_bits': 8,
+            'strategy': 'channel',
+            'group_size': 128,
+            'symmetric': True,
+            'dynamic': False,
+            'actorder': 'none',
+            'observer': None,
+            'observer_kwargs': {},
+            'block_structure': None
+        })()
+        
+        # 创建默认的输入激活量化配置字典
+        self.input_quant = type('InputQuant', (), {
+            'type': 'int',
+            'num_bits': 8,
+            'strategy': 'token',
+            'group_size': 128,
+            'symmetric': True,
+            'dynamic': True,
+            'actorder': 'none',
+            'observer': None,
+            'observer_kwargs': {},
+            'block_structure': None
+        })()
 
-        # Create a default weight quantization config dictionary
-        self.weight_quant = type(
-            "WeightQuant",
-            (),
-            {
-                "type": "int",
-                "num_bits": 8,
-                "strategy": "channel",
-                "group_size": 128,
-                "symmetric": True,
-                "dynamic": False,
-                "actorder": "none",
-                "observer": None,
-                "observer_kwargs": {},
-                "block_structure": None,
-            },
-        )()
-
-        # Create a default input activation quantization config dictionary
-        self.input_quant = type(
-            "InputQuant",
-            (),
-            {
-                "type": "int",
-                "num_bits": 8,
-                "strategy": "token",
-                "group_size": 128,
-                "symmetric": True,
-                "dynamic": True,
-                "actorder": "none",
-                "observer": None,
-                "observer_kwargs": {},
-                "block_structure": None,
-            },
-        )()
-
-        # Change comparison method to directly compare strings
+        # 修改比较方式，直接比较字符串
         per_channel = (
             self.weight_quant.strategy == "channel"
-            and self.input_quant.strategy == "token"
-        )
+            and self.input_quant.strategy == "token")
         if not per_channel:
             raise ValueError(
                 "For INT8 Fused MoE layers, we require channelwise, "
                 "dynamic per token quantization. Found "
-                f"{self.weight_quant}, {self.input_quant}"
-            )
+                f"{self.weight_quant}, {self.input_quant}")
 
         self.static_input_scales = not self.input_quant.dynamic
         if self.static_input_scales:
             raise ValueError(
                 "For INT8 Fused MoE layers, we require channelwise, "
-                "dynamic per token quantization. Found static input scales."
-            )
+                "dynamic per token quantization. Found static input scales.")
 
-    def create_weights1(
-        self,
-        layer: torch.nn.Module,
-        num_experts: int,
-        hidden_size: int,
-        intermediate_size_per_partition: int,
-        params_dtype: torch.dtype,
-        **extra_weight_attrs,
-    ):
-        # Use float32 as a placeholder for weights to facilitate loading original weights from ckpt
-        w13_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                2 * intermediate_size_per_partition,
-                hidden_size,
-                dtype=params_dtype,
-            ),  #  generally is torch.bfloat16
-            requires_grad=False,
-        )
+    def create_weights1(self, layer: torch.nn.Module, num_experts: int, hidden_size: int, intermediate_size_per_partition: int, params_dtype: torch.dtype, **extra_weight_attrs):
+        # 权重先用浮点占位，便于从 ckpt 加载原始权重
+        w13_weight = torch.nn.Parameter(torch.empty(
+            num_experts,
+            2 * intermediate_size_per_partition,
+            hidden_size,
+            dtype=params_dtype),  # 通常是 torch.bfloat16
+            requires_grad=False)
         layer.register_parameter("w13_weight", w13_weight)
         set_weight_attrs(w13_weight, extra_weight_attrs)
 
-        w2_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                hidden_size,
-                intermediate_size_per_partition,
-                dtype=params_dtype,
-            ),
-            requires_grad=False,
-        )
+        w2_weight = torch.nn.Parameter(torch.empty(
+            num_experts,
+            hidden_size,
+            intermediate_size_per_partition,
+            dtype=params_dtype),
+            requires_grad=False)
         layer.register_parameter("w2_weight", w2_weight)
         set_weight_attrs(w2_weight, extra_weight_attrs)
 
-        # Channel scale: float32 + 2D [E, out] (aligned with fused_moe/UT)
+        # 通道 scale：float32 + 二维 [E, out]（与 fused_moe/UT 对齐）
         w13_weight_scale = torch.nn.Parameter(
-            torch.empty(
-                num_experts, 2 * intermediate_size_per_partition, dtype=torch.float32
-            ),
-            requires_grad=False,
-        )
+            torch.empty(num_experts, 2 * intermediate_size_per_partition, dtype=torch.float32),
+            requires_grad=False)
         w2_weight_scale = torch.nn.Parameter(
             torch.empty(num_experts, hidden_size, dtype=torch.float32),
-            requires_grad=False,
-        )
+            requires_grad=False)
         layer.register_parameter("w13_weight_scale", w13_weight_scale)
         layer.register_parameter("w2_weight_scale", w2_weight_scale)
 
-        # Input scale can be dynamically calculated
+        # 输入 scale 动态计算即可
         layer.w13_input_scale = None
         layer.w2_input_scale = None
 
-    def create_weights(
-        self,
-        layer: torch.nn.Module,
-        num_experts: int,
-        hidden_size: int,
-        intermediate_size_per_partition: int,
-        params_dtype: torch.dtype,
-        **extra_weight_attrs,
-    ):
-        w13_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                2 * intermediate_size_per_partition,
-                hidden_size,
-                dtype=torch.int8,
-            ),  # directly use int8
-            requires_grad=False,
-        )
+    def create_weights(self, layer: torch.nn.Module, num_experts: int, hidden_size: int, intermediate_size_per_partition: int, params_dtype: torch.dtype, **extra_weight_attrs):
+        w13_weight = torch.nn.Parameter(torch.empty(
+            num_experts,
+            2 * intermediate_size_per_partition,
+            hidden_size,
+            dtype=torch.int8),  # 直接使用 int8
+            requires_grad=False)
         layer.register_parameter("w13_weight", w13_weight)
         set_weight_attrs(w13_weight, extra_weight_attrs)
 
-        w2_weight = torch.nn.Parameter(
-            torch.empty(
-                num_experts,
-                hidden_size,
-                intermediate_size_per_partition,
-                dtype=torch.int8,
-            ),  # directly use int8
-            requires_grad=False,
-        )
+        w2_weight = torch.nn.Parameter(torch.empty(
+            num_experts,
+            hidden_size,
+            intermediate_size_per_partition,
+            dtype=torch.int8),  # 直接使用 int8
+            requires_grad=False)
         layer.register_parameter("w2_weight", w2_weight)
         set_weight_attrs(w2_weight, extra_weight_attrs)
 
-        # Scale factors
+        # 缩放因子
         w13_weight_scale = torch.nn.Parameter(
-            torch.empty(
-                num_experts, 2 * intermediate_size_per_partition, dtype=torch.float32
-            ),
-            requires_grad=False,
-        )
+            torch.empty(num_experts, 2 * intermediate_size_per_partition, dtype=torch.float32),
+            requires_grad=False)
         w2_weight_scale = torch.nn.Parameter(
             torch.empty(num_experts, hidden_size, dtype=torch.float32),
-            requires_grad=False,
-        )
+            requires_grad=False)
         layer.register_parameter("w13_weight_scale", w13_weight_scale)
         layer.register_parameter("w2_weight_scale", w2_weight_scale)
 
-        # Input scale can be dynamically calculated
+        # 输入 scale 动态计算
         layer.w13_input_scale = None
         layer.w2_input_scale = None
-
+        
     @torch.no_grad()
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         return
-        # Convert original weights to float32 for more robust statistics
+        #原始权重转 float32 做统计更稳健
         w13_f = layer.w13_weight.float()
-        w2_f = layer.w2_weight.float()
+        w2_f  = layer.w2_weight.float()
 
-        # Each column (abs_max) -> per-column scale (out dimension is dim=1, column is dim=-1)
+        # 每列(abs_max) -> per-column scale（out 维在 dim=1，列在 dim=-1）
         qmax = 127.0
         w13_abs_max = torch.amax(torch.abs(w13_f), dim=-1)  # [E, 2N]
-        w2_abs_max = torch.amax(torch.abs(w2_f), dim=-1)  # [E, H]
+        w2_abs_max  = torch.amax(torch.abs(w2_f),  dim=-1)  # [E, H]
 
         w13_scale_2d = torch.clamp(w13_abs_max, min=1e-6) / qmax  # [E, 2N], float32
-        w2_scale_2d = torch.clamp(w2_abs_max, min=1e-6) / qmax  # [E, H],  float32
+        w2_scale_2d  = torch.clamp(w2_abs_max,  min=1e-6) / qmax  # [E, H],  float32
 
-        # Quantization: broadcast 3D scale and store back to 2D scale
+        # 量化：用 3D scale 广播，存回 2D scale
         w13_scale_3d = w13_scale_2d.unsqueeze(-1)  # [E, 2N, 1]
-        w2_scale_3d = w2_scale_2d.unsqueeze(-1)  # [E, H, 1]
+        w2_scale_3d  = w2_scale_2d.unsqueeze(-1)   # [E, H, 1]
 
         w13_q = torch.round(w13_f / w13_scale_3d).clamp_(-128, 127).to(torch.int8)
-        w2_q = torch.round(w2_f / w2_scale_3d).clamp_(-128, 127).to(torch.int8)
+        w2_q  = torch.round(w2_f  / w2_scale_3d ).clamp_(-128, 127).to(torch.int8)
 
-        # Optional: If your fused/kernel expects scale pre-multiplied by 127 (to be consistent with some UT backends), uncomment the following two lines:
+        # 可选：若你的 fused/kernel 期望 scale 预乘 127（与某些 UT 后端一致），打开下面两行：
         w13_scale_2d = w13_scale_2d * 127.0
-        w2_scale_2d = w2_scale_2d * 127.0
+        w2_scale_2d  = w2_scale_2d  * 127.0
 
-        # Write back parameters: weight int8; scale uses float32 + 2D
-        replace_parameter(
-            layer, "w13_weight", torch.nn.Parameter(w13_q, requires_grad=False)
-        )
-        replace_parameter(
-            layer, "w2_weight", torch.nn.Parameter(w2_q, requires_grad=False)
-        )
-        replace_parameter(
-            layer,
-            "w13_weight_scale",
-            torch.nn.Parameter(w13_scale_2d.contiguous(), requires_grad=False),
-        )
-        replace_parameter(
-            layer,
-            "w2_weight_scale",
-            torch.nn.Parameter(w2_scale_2d.contiguous(), requires_grad=False),
-        )
-
-        # Brief check
-        print(
-            f"w13: {w13_q.shape}, w13_s: {w13_scale_2d.shape}, w2: {w2_q.shape}, w2_s: {w2_scale_2d.shape}"
-        )
+        # 回写参数：权重 int8；scale 用 float32 + 2D
+        replace_parameter(layer, 'w13_weight', torch.nn.Parameter(w13_q, requires_grad=False))
+        replace_parameter(layer, 'w2_weight',  torch.nn.Parameter(w2_q,  requires_grad=False))
+        replace_parameter(layer, 'w13_weight_scale',
+                        torch.nn.Parameter(w13_scale_2d.contiguous(), requires_grad=False))
+        replace_parameter(layer, 'w2_weight_scale',
+                        torch.nn.Parameter(w2_scale_2d.contiguous(),  requires_grad=False))
 
+        # 简要检查
+        print(f"w13: {w13_q.shape}, w13_s: {w13_scale_2d.shape}, w2: {w2_q.shape}, w2_s: {w2_scale_2d.shape}")
+ 
     def apply(
         self,
         layer: torch.nn.Module,
@@ -300,11 +214,11 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
         e_score_correction_bias: Optional[torch.Tensor] = None,
         apply_router_weight_on_input: bool = False,
         activation: str = "silu",
-        enable_eplb: bool = False,  # Add this parameter
-        expert_load_view: Optional[torch.Tensor] = None,  # Add this parameter
-        logical_to_physical_map: Optional[torch.Tensor] = None,  # Add this parameter
-        logical_replica_count: Optional[torch.Tensor] = None,  # Add this parameter
-        linear_weights: Optional[torch.Tensor] = None,  # Add this parameter
+        enable_eplb: bool = False,  # 添加这个参数
+        expert_load_view: Optional[torch.Tensor] = None,  # 添加这个参数
+        logical_to_physical_map: Optional[torch.Tensor] = None,  # 添加这个参数
+        logical_replica_count: Optional[torch.Tensor] = None,  # 添加这个参数
+        linear_weights: Optional[torch.Tensor] = None,  # 添加这个参数
     ) -> torch.Tensor:
 
         output = torch.empty_like(x)
@@ -326,8 +240,5 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
         )
         return output
 
-
-print(
-    "[Monkey Patch Applied] >>> vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsMoEMethod \
-      --> vllm_xpu.model_executor.layers.quantization.compressed_tensors_moe.py:CompressedTensorsMoEMethod"
-)
+print("[Monkey Patch Applied] >>> vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.CompressedTensorsMoEMethod \
+      --> vllm_xpu.model_executor.layers.quantization.compressed_tensors_moe.py:CompressedTensorsMoEMethod")

vllm_kunlun/ops/quantization/gptq.py

@@ -1,108 +0,0 @@
-#
-# Copyright (c) 2025 Baidu, Inc. All Rights Reserved.
-# Author: Li Wei, You Zeyu
-# Email: liwei157@baidu.com, youzeyu@baidu.com
-# This file is a part of the vllm-kunlun project.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import torch
-
-from torch.nn.parameter import Parameter
-from typing import Optional
-from vllm.model_executor.layers.quantization.gptq import GPTQLinearMethod, ExllamaState
-
-
-def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-    # for torch.compile
-    layer.qzeros = Parameter(
-        self.repack_int4_for_kunlun(layer.qzeros.data, self.quant_config.weight_bits)
-        if self.quant_config.weight_bits == 4 else layer.qzeros.data,
-        requires_grad=False
-    )
-    layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
-    layer.g_idx = Parameter(layer.g_idx.data, requires_grad=False)
-    layer.scales = Parameter(layer.scales.data, requires_grad=False)
-
-    # exllama needs to shuffle the weight after the weight is loaded
-    # here we do the shuffle on first forward pass
-    if layer.exllama_state == ExllamaState.UNINITIALIZED:
-        if self.quant_config.desc_act:
-            layer.g_idx.data = torch.argsort(layer.g_idx).to(torch.int)
-        else:
-            layer.g_idx.data = torch.empty((0, ),
-                                            dtype=torch.int,
-                                            device=layer.g_idx.device)
-        layer.exllama_state = ExllamaState.READY
-
-        # No need shuffle on xpu
-        # ops.gptq_shuffle(layer.qweight, layer.g_idx,
-        #                  self.quant_config.weight_bits)
-
-
-def repack_int4_for_kunlun(self, packed: torch.Tensor, num_bits: int = 4):
-    N, K = packed.shape
-    assert num_bits == 4, "Only int4 supported now"
-    shifts = torch.arange(0, 32, num_bits, device=packed.device, dtype=torch.int32)
-
-    # Unpack int32 to int4 values
-    unpacked_gptq = (
-        packed.view(N, K // 8, 8).unsqueeze(-1) >> shifts
-    ) & 0xF  # [N, K//8, 8, 8]
-
-    # Convert to KUNLUN order
-    GPTQ_TO_KUNLUN_ORDER_FAST = [
-        32, 0, 33, 1, 34, 2, 35, 3,
-        36, 4, 37, 5, 38, 6, 39, 7,
-        40, 8, 41, 9, 42, 10, 43, 11,
-        44, 12, 45, 13, 46, 14, 47, 15,
-        48, 16, 49, 17, 50, 18, 51, 19,
-        52, 20, 53, 21, 54, 22, 55, 23,
-        56, 24, 57, 25, 58, 26, 59, 27,
-        60, 28, 61, 29, 62, 30, 63, 31,
-    ]
-    unpacked_gptq = unpacked_gptq.reshape(N, K // 8, 64)
-    unpacked_kunlun = unpacked_gptq[..., GPTQ_TO_KUNLUN_ORDER_FAST]  # [N, K//8, 64]
-
-    # Pack to int32
-    unpacked_kunlun = unpacked_kunlun.reshape(N, K // 8, 8, 8)
-    packed_kunlun = (
-        (unpacked_kunlun << shifts).sum(dim=-1, dtype=torch.int32).reshape(N, K)
-    )  # [N, K]
-
-    return packed_kunlun
-
-
-def apply(
-    self, layer: torch.nn.Module, x: torch.Tensor, bias: Optional[torch.Tensor] = None
-) -> torch.Tensor:
-    out_shape = x.shape[:-1] + (layer.qweight.shape[-1], )
-    reshaped_x = x.reshape(-1, x.shape[-1])
-
-    output = torch.ops.xspeedgate_ops.gptq_gemm(
-        reshaped_x,
-        layer.qweight,
-        layer.qzeros,
-        layer.scales,
-        layer.g_idx,
-        layer.exllama_state == ExllamaState.READY,
-        self.quant_config.weight_bits,
-    )
-    if bias is not None:
-        output.add_(bias)
-    return output.reshape(out_shape)
-
-
-GPTQLinearMethod.repack_int4_for_kunlun = repack_int4_for_kunlun
-GPTQLinearMethod.process_weights_after_loading = process_weights_after_loading
-GPTQLinearMethod.apply = apply