fix: fix ep v2 for sm100

Author: niushengxiao

lightllm/common/quantization/deepgemm.py

@@ -131,6 +131,7 @@ class DeepGEMMFP8FP4B32QuantizationMethod(DeepGEMMBaseQuantizationMethod):
     def __init__(self):
         super().__init__()
         self.block_size = 32
+        self.ue8m0_pack_factor = 4
         self.weight_suffix = "weight"
         self.weight_zero_point_suffix = None
         self.weight_scale_suffix = None
@@ -179,14 +180,26 @@ def apply(
     def _create_weight(
         self, out_dims: Union[int, List[int]], in_dim: int, dtype: torch.dtype, device_id: int, num_experts: int = 1
     ) -> Tuple[WeightPack, List[WeightPack]]:
+        import deep_gemm
+
         out_dim = sum(out_dims) if isinstance(out_dims, list) else out_dims
         assert in_dim % 2 == 0, "FP4 packed weight requires even input dimension"
         assert in_dim % self.block_size == 0, "FP4 scale dimension must be divisible by block_size"
+        assert (
+            in_dim % (self.block_size * self.ue8m0_pack_factor) == 0
+        ), "SM100 FP4 scale layout requires input dimension divisible by 128"
         expert_prefix = (num_experts,) if num_experts > 1 else ()
         weight = torch.empty(expert_prefix + (out_dim, in_dim // 2), dtype=torch.int8).cuda(device_id)
-        weight_scale = torch.empty(expert_prefix + (out_dim, in_dim // self.block_size), dtype=torch.int32).cuda(
+        raw_weight_scale = torch.empty(expert_prefix + (out_dim, in_dim // self.block_size), dtype=torch.float32).cuda(
             device_id
         )
+        weight_scale = deep_gemm.transform_sf_into_required_layout(
+            raw_weight_scale,
+            out_dim,
+            in_dim,
+            (1, self.block_size),
+            num_experts if num_experts > 1 else None,
+        )
         mm_param = WeightPack(weight=weight, weight_scale=weight_scale)
         mm_param_list = self._split_weight_pack(
             mm_param,

lightllm/distributed/communication_op.py

@@ -156,28 +156,10 @@ def new_deepep_group(
         self.ll_decode_num_tokens = decode_num_max_dispatch_tokens_per_rank
         self.ll_hidden = hidden_size
         self.ll_num_experts = n_routed_experts + get_redundancy_expert_num() * global_world_size
-        self.ep_buffer = deep_ep.ElasticBuffer(
-            deepep_group,
-            num_max_tokens_per_rank=self.ll_num_tokens,
-            hidden=self.ll_hidden,
-            num_topk=num_experts_per_tok,
-            use_fp8_dispatch=True,
-            allow_multiple_reduction=False,
-        )
-        self.ep_mega_moe_buffer = None
         self.ep_low_latency_buffer = None
-        if not is_sm100_gpu():
-            num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
-                self.ll_decode_num_tokens, self.ll_hidden, global_world_size, self.ll_num_experts
-            )
-            self.ep_low_latency_buffer = deep_ep.Buffer(
-                deepep_group,
-                int(1e9),
-                num_rdma_bytes,
-                low_latency_mode=True,
-                num_qps_per_rank=(self.ll_num_experts // global_world_size),
-            )
-        else:
+        self.ep_mega_moe_buffer = None
+        if is_sm100_gpu():
+            self.ep_buffer = None
             if moe_intermediate_size is None:
                 raise ValueError("SM100 Mega MoE requires moe_intermediate_size or intermediate_size in model config")
 
@@ -191,6 +173,28 @@ def new_deepep_group(
                 self.ll_hidden,
                 moe_intermediate_size,
             )
+            self._set_num_sms_for_deep_gemm(0)
+            logger.info("SM100 detected: skip DeepEP ElasticBuffer init and use Mega MoE buffer only.")
+            return
+
+        self.ep_buffer = deep_ep.ElasticBuffer(
+            deepep_group,
+            num_max_tokens_per_rank=self.ll_num_tokens,
+            hidden=self.ll_hidden,
+            num_topk=num_experts_per_tok,
+            use_fp8_dispatch=True,
+            allow_multiple_reduction=False,
+        )
+        num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
+            self.ll_decode_num_tokens, self.ll_hidden, global_world_size, self.ll_num_experts
+        )
+        self.ep_low_latency_buffer = deep_ep.Buffer(
+            deepep_group,
+            int(1e9),
+            num_rdma_bytes,
+            low_latency_mode=True,
+            num_qps_per_rank=(self.ll_num_experts // global_world_size),
+        )
         theoretical_sms = self.ep_buffer.get_theoretical_num_sms(self.ll_num_experts, num_experts_per_tok)
         self._set_num_sms_for_deep_gemm(theoretical_sms)
 

lightllm/server/api_start.py

@@ -25,6 +25,7 @@
     auto_set_max_req_total_len,
 )
 from lightllm.utils.dist_check_utils import auto_configure_allreduce_flags_from_args
+from lightllm.utils.device_utils import is_sm100_gpu
 
 logger = init_logger(__name__)
 
@@ -36,7 +37,9 @@ def _auto_set_fa4_page_size(args, requested_backends):
     from lightllm.utils.fa4_utils import infer_fa4_page_size
 
     page_size = infer_fa4_page_size(args.model_dir)
-    if page_size is None:
+    if is_sm100_gpu():
+        page_size = 128
+    elif page_size is None:
         return
 
     set_page_size(page_size)

lightllm/server/router/dynamic_prompt/paged_radix_cache.py

@@ -524,8 +524,8 @@ def _print_helper(self, node: TreeNode, indent):
 
     def free_radix_cache_to_get_enough_token(self, need_token_num):
         assert self.mem_manager is not None
-        if need_token_num > self.mem_manager.can_use_mem_size:
-            need_evict_token_num = need_token_num - self.mem_manager.can_use_mem_size
+        if need_token_num > self.mem_manager.allocator.can_use_mem_size:
+            need_evict_token_num = need_token_num - self.mem_manager.allocator.can_use_mem_size
             release_mems = []
 
             def release_mem(mem_index):