From 33f08a98b0ea32b1cd921b84512bd18ec2f499fd Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sat, 15 Nov 2025 16:36:00 +0800 Subject: [PATCH] Tiny refactor condition to requant scale ue8m0 (#13286) --- python/sglang/srt/model_loader/utils.py | 10 ++++++++++ python/sglang/srt/models/deepseek_v2.py | 14 +++++++------- python/sglang/srt/models/longcat_flash.py | 13 +++++++------ python/sglang/srt/models/longcat_flash_nextn.py | 8 +++----- 4 files changed, 27 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index 2511c61cc..b0a45148d 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/python/sglang/srt/model_loader/utils.py @@ -12,6 +12,7 @@ from torch import nn from transformers.dynamic_module_utils import get_class_from_dynamic_module from sglang.srt.configs.model_config import ModelConfig, ModelImpl +from sglang.srt.layers import deep_gemm_wrapper logger = logging.getLogger(__name__) @@ -119,6 +120,15 @@ def post_load_weights(model: nn.Module, model_config: ModelConfig): model.post_load_weights() +def should_deepgemm_weight_requant_ue8m0(weight_block_size): + """Should we requant fp8 weights into UE8M0 format when loading the model""" + return ( + deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 + and weight_block_size is not None + ) + + def should_async_load(weight: torch.Tensor) -> bool: """Return True if we should load the given weight asynchronously. diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index b174407c7..ac7e7e2bd 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -110,7 +110,11 @@ from sglang.srt.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors -from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load +from sglang.srt.model_loader.utils import ( + maybe_executor_submit, + should_async_load, + should_deepgemm_weight_requant_ue8m0, +) from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.server_args import get_global_server_args from sglang.srt.single_batch_overlap import SboFlags @@ -3418,12 +3422,8 @@ class DeepseekV2ForCausalLM(nn.Module): self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous()) self_attn.use_deep_gemm_bmm = True - if ( - not ENABLE_FLASHINFER_FP8_GEMM - and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 - and hasattr(self.quant_config, "weight_block_size") - and self.quant_config.weight_block_size is not None + if not ENABLE_FLASHINFER_FP8_GEMM and should_deepgemm_weight_requant_ue8m0( + weight_block_size=getattr(self.quant_config, "weight_block_size", None) ): self._weight_requant_ue8m0(is_nextn) diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index 84aeb8b30..fa7ba01bc 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/python/sglang/srt/models/longcat_flash.py @@ -80,7 +80,11 @@ from sglang.srt.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch -from sglang.srt.model_loader.utils import maybe_executor_submit, should_async_load +from sglang.srt.model_loader.utils import ( + maybe_executor_submit, + should_async_load, + should_deepgemm_weight_requant_ue8m0, +) from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA from sglang.srt.server_args import get_global_server_args @@ -774,11 +778,8 @@ class LongcatFlashForCausalLM(nn.Module): # TODO(linguoyuan) EPMoE not support DEEPGEMM_BLACKWELL, DeepEP needs to be supported in the future deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 = False - if ( - deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 - and hasattr(self.quant_config, "weight_block_size") - and self.quant_config.weight_block_size is not None + if should_deepgemm_weight_requant_ue8m0( + weight_block_size=getattr(self.quant_config, "weight_block_size", None) ): self._weight_requant_ue8m0() diff --git a/python/sglang/srt/models/longcat_flash_nextn.py b/python/sglang/srt/models/longcat_flash_nextn.py index cae974815..b51417c65 100644 --- a/python/sglang/srt/models/longcat_flash_nextn.py +++ b/python/sglang/srt/models/longcat_flash_nextn.py @@ -66,6 +66,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ( VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0 from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA from sglang.srt.models.longcat_flash import LongcatFlashForCausalLM, LongcatFlashMLP @@ -455,11 +456,8 @@ class LongcatFlashForCausalLMNextN(LongcatFlashForCausalLM): self.config.hidden_size / self.config.kv_lora_rank ) ** 0.5 - if ( - deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM - and deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 - and hasattr(self.quant_config, "weight_block_size") - and self.quant_config.weight_block_size is not None + if should_deepgemm_weight_requant_ue8m0( + weight_block_size=getattr(self.quant_config, "weight_block_size", None) ): self._weight_requant_ue8m0()