Tiny refactor condition to requant scale ue8m0 (#13286)

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.
 

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 (
+        if not ENABLE_FLASHINFER_FP8_GEMM and should_deepgemm_weight_requant_ue8m0(
-            not ENABLE_FLASHINFER_FP8_GEMM
+            weight_block_size=getattr(self.quant_config, "weight_block_size", None)
-            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
         ):
             self._weight_requant_ue8m0(is_nextn)
 

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 (
+        if should_deepgemm_weight_requant_ue8m0(
-            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+            weight_block_size=getattr(self.quant_config, "weight_block_size", None)
-            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
         ):
             self._weight_requant_ue8m0()
 

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 (
+        if should_deepgemm_weight_requant_ue8m0(
-            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+            weight_block_size=getattr(self.quant_config, "weight_block_size", None)
-            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
         ):
             self._weight_requant_ue8m0()