Skip to content

GitLab

Unverified Commit a8ba3279 authored by fzyzcjy's avatar fzyzcjy Committed by GitHub
Browse files

Fix triton_kernels import error on some hardwares (#11831)

parent 3b80232d
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 4 deletions
python/sglang/srt/layers/quantization/unquant.py
View file @ a8ba3279
from __future__ import annotations from __future__ import annotations
import importlib.util
from typing import TYPE_CHECKING, List, Optional from typing import TYPE_CHECKING, List, Optional
import torch import torch
...@@ -31,8 +30,6 @@ if TYPE_CHECKING: ...@@ -31,8 +30,6 @@ if TYPE_CHECKING:
StandardDispatchOutput, StandardDispatchOutput,
) )
has_triton_kernels = importlib.util.find_spec("triton_kernels") is not None
_is_cpu_amx_available = cpu_has_amx_support() _is_cpu_amx_available = cpu_has_amx_support()
_is_hip = is_hip() _is_hip = is_hip()
...@@ -143,7 +140,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp): ...@@ -143,7 +140,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
self.triton_kernel_moe_forward = None self.triton_kernel_moe_forward = None
self.triton_kernel_moe_with_bias_forward = None self.triton_kernel_moe_with_bias_forward = None
if torch.cuda.is_available() and has_triton_kernels: if torch.cuda.is_available() and use_triton_kernels:
from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import ( from sglang.srt.layers.moe.fused_moe_triton.triton_kernels_moe import (
triton_kernel_moe_forward as _tk_forward, triton_kernel_moe_forward as _tk_forward,
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment