diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 0a03f5dc4..5db8967c4 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -10,35 +10,12 @@ import torch.nn.functional as F from torch.nn import Module from torch.nn.parameter import Parameter -from sglang.srt.distributed import get_tp_group -from sglang.srt.distributed.device_communicators.pynccl_allocator import ( - use_symmetric_memory, -) -from sglang.srt.layers.dp_attention import is_allocation_symmetric - -try: - from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( - apply_fp8_marlin_linear, - prepare_fp8_layer_for_marlin, - ) - - MARLIN_FP8_AVAILABLE = True -except ImportError: - MARLIN_FP8_AVAILABLE = False - - def dummy_func(*args, **kwargs): - raise ImportError( - "marlin FP8 requires some operators from vllm. Please install vllm." - ) - - apply_fp8_marlin_linear = prepare_fp8_layer_for_marlin = dummy_func - - from sglang.srt.distributed import get_tensor_model_parallel_world_size, get_tp_group from sglang.srt.distributed.device_communicators.pynccl_allocator import ( use_symmetric_memory, ) from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading +from sglang.srt.layers.dp_attention import is_allocation_symmetric from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo @@ -69,6 +46,10 @@ from sglang.srt.layers.quantization.fp8_utils import ( normalize_e4m3fn_to_e4m3fnuz, ) from sglang.srt.layers.quantization.kv_cache import BaseKVCacheMethod +from sglang.srt.layers.quantization.marlin_utils_fp8 import ( + apply_fp8_marlin_linear, + prepare_fp8_layer_for_marlin, +) from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.quantization.utils import ( all_close_1d, @@ -228,7 +209,7 @@ class Fp8LinearMethod(LinearMethodBase): # For GPUs that lack FP8 hardware support, we can leverage the Marlin # kernel for fast weight-only FP8 quantization self.use_marlin = False - if _is_cuda and MARLIN_FP8_AVAILABLE: + if _is_cuda: force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN") auto_enable = can_auto_enable_marlin_fp8() self.use_marlin = force_marlin or auto_enable