[Fix][XPU/ROCm/NPU] Defer sgl_kernel.quantization import in expert_pack (#36529)

This commit is contained in:
ashwini rathi
2026-08-27 16:17:24 -07:00
committed by GitHub
parent 76217f6603
commit 7cbe564829
6 changed files with 36 additions and 17 deletions
@@ -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,
+12 -3
View File
@@ -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(