[CPU] refine CPU integration code (#7647)

python/sglang/srt/layers/linear.py

@@ -17,6 +17,7 @@ from sglang.srt.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
 )
+from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
 from sglang.srt.layers.parameter import (
     BasevLLMParameter,
     BlockQuantScaleParameter,
@@ -31,10 +32,10 @@ from sglang.srt.layers.quantization.base_config import (
     QuantizeMethodBase,
 )
 from sglang.srt.utils import (
-    _process_weight_after_loading,
     cpu_has_amx_support,
     is_cpu,
     set_weight_attrs,
+    use_intel_amx_backend,
 )
 
 logger = logging.getLogger(__name__)
@@ -175,7 +176,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
 
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         if _is_cpu and _is_cpu_amx_available:
-            _process_weight_after_loading(layer, ["weight"])
+            _amx_process_weight_after_loading(layer, ["weight"])
 
     def apply(
         self,
@@ -184,7 +185,7 @@ class UnquantizedLinearMethod(LinearMethodBase):
         bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
 
-        if getattr(layer, "use_intel_amx_backend", False):
+        if use_intel_amx_backend(layer):
             return torch.ops.sgl_kernel.weight_packed_linear(
                 x, layer.weight, bias, True  # is_vnni
             )