[Fix][XPU/ROCm/NPU] Defer sgl_kernel.quantization import in expert_pack (#36529)
This commit is contained in:
@@ -7,7 +7,6 @@ from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from sgl_kernel.quantization import ggml_moe_a8_vec
|
||||
|
||||
from sglang.kernels.ops.moe.expert_pack_mxfp4 import (
|
||||
mxfp4_matvec,
|
||||
@@ -157,6 +156,10 @@ class ExpertPackMoEMethod(FusedMoEMethodBase):
|
||||
weight_type: int,
|
||||
output_size: int,
|
||||
) -> torch.Tensor:
|
||||
# sgl_kernel.quantization is CUDA/MUSA-only; keep the import local so
|
||||
# this module stays importable on other devices (see gguf.py:44-70).
|
||||
from sgl_kernel.quantization import ggml_moe_a8_vec
|
||||
|
||||
return ggml_moe_a8_vec(
|
||||
inputs,
|
||||
weights,
|
||||
|
||||
@@ -424,16 +424,25 @@ class HFRunner:
|
||||
f"before producing output"
|
||||
)
|
||||
|
||||
def terminate(self):
|
||||
def _stop_model_proc(self):
|
||||
# Fire-and-forget terminate() leaves the child holding the accelerator
|
||||
# during teardown; a follow-on SRTRunner on the same device can then
|
||||
# deadlock in driver init (observed on Intel XPU B580).
|
||||
self.model_proc.terminate()
|
||||
self.model_proc.join(timeout=30)
|
||||
if self.model_proc.is_alive():
|
||||
self.model_proc.kill()
|
||||
self.model_proc.join()
|
||||
self.in_queue = self.out_queue = None
|
||||
|
||||
def terminate(self):
|
||||
self._stop_model_proc()
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
self.model_proc.terminate()
|
||||
self.in_queue = self.out_queue = None
|
||||
self._stop_model_proc()
|
||||
|
||||
@staticmethod
|
||||
def forward_generation_raw(
|
||||
|
||||
Reference in New Issue
Block a user