[Fix] Fix trtllm fp4 moe kernel not found error (#21343)
This commit is contained in:
@@ -76,21 +76,11 @@ from sglang.srt.utils.custom_op import register_custom_op
|
|||||||
if is_flashinfer_available():
|
if is_flashinfer_available():
|
||||||
from flashinfer import fp4_quantize
|
from flashinfer import fp4_quantize
|
||||||
|
|
||||||
# Try to import FP4 TRTLLM function if flashinfer is available
|
|
||||||
trtllm_fp4_block_scale_moe = None
|
|
||||||
if get_moe_runner_backend().is_flashinfer_trtllm():
|
|
||||||
try:
|
|
||||||
from flashinfer.fused_moe import trtllm_fp4_block_scale_moe
|
|
||||||
except ImportError:
|
|
||||||
trtllm_fp4_block_scale_moe = None
|
|
||||||
|
|
||||||
_is_hip = is_hip()
|
_is_hip = is_hip()
|
||||||
_is_cpu_amx_available = cpu_has_amx_support()
|
_is_cpu_amx_available = cpu_has_amx_support()
|
||||||
_is_cpu = is_cpu()
|
_is_cpu = is_cpu()
|
||||||
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
|
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
|
def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
|
||||||
a2a_backend = get_moe_a2a_backend()
|
a2a_backend = get_moe_a2a_backend()
|
||||||
@@ -1329,6 +1319,8 @@ class FlashInferFP4MoE(FusedMoE):
|
|||||||
hidden_states: Input tensor
|
hidden_states: Input tensor
|
||||||
topk_output: TopKOutput object with Bypassed format
|
topk_output: TopKOutput object with Bypassed format
|
||||||
"""
|
"""
|
||||||
|
from flashinfer.fused_moe import trtllm_fp4_block_scale_moe
|
||||||
|
|
||||||
assert isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod)
|
assert isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod)
|
||||||
|
|
||||||
assert (
|
assert (
|
||||||
|
|||||||
Reference in New Issue
Block a user