fix: fix ep v2 for sm100

Author: niushengxiao

lightllm/common/basemodel/attention/fa4/fp.py

@@ -59,6 +59,8 @@ class Fa4PrefillAttState(PagedFa3PrefillAttState):
     def _normal_prefill_att(
         self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, att_control: AttControl, alloc_func=torch.empty
     ) -> torch.Tensor:
+        import triton
+
         if att_control.use_sliding_window:
             window_size = att_control.sliding_window
         else:
@@ -81,7 +83,7 @@ def _normal_prefill_att(
             cu_seqlens_q=self.cu_seqlens_q,
             seqused_k=self.infer_state.b_seq_len.int(),
             max_seqlen_q=self.infer_state.max_q_seq_len,
-            max_seqlen_k=self.infer_state.max_kv_seq_len,
+            max_seqlen_k=triton.cdiv(self.infer_state.max_kv_seq_len, self.backend.page_size) * self.backend.page_size,
             page_table=self.page_table,
             softmax_scale=softmax_scale,
             causal=True,

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py

@@ -154,6 +154,7 @@ def experts(
         topk_group: int,
         num_expert_group: int,
         is_prefill: Optional[bool] = None,
+        is_cuda_graph: bool = False,
     ) -> torch.Tensor:
         """Backward compatible method that routes to platform-specific implementation."""
         return self.fuse_moe_impl(
@@ -169,6 +170,7 @@ def experts(
             topk_group=topk_group,
             num_expert_group=num_expert_group,
             is_prefill=is_prefill,
+            is_cuda_graph=is_cuda_graph,
         )
 
     def use_sm100_mega_moe(self) -> bool:

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py

@@ -99,6 +99,54 @@ def _mega_moe(
         )
         return output
 
+    def _sm100_fp4_cuda_graph_moe(
+        self,
+        hidden_states: torch.Tensor,
+        w13: WeightPack,
+        w2: WeightPack,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ) -> torch.Tensor:
+        from deep_gemm.utils import per_token_cast_to_fp8
+
+        buffer = getattr(dist_group_manager, "ep_buffer", None)
+        if buffer is None:
+            raise RuntimeError("SM100 CUDA graph MoE fallback requires dist_group_manager.ep_buffer")
+
+        num_max_tokens_per_rank = get_deepep_num_max_dispatch_tokens_per_rank_decode()
+        qinput_tensor = per_token_cast_to_fp8(
+            hidden_states,
+            use_ue8m0=True,
+            gran_k=self.quant_method.block_size,
+            use_packed_ue8m0=True,
+        )
+        alignment = getattr(dist_group_manager, "ep_expert_alignment", 128)
+        cumulative_stats = self._get_mega_moe_stats(w13.weight.shape[0], hidden_states.device)
+        recv_x, recv_topk_idx, recv_topk_weights, handle, _ = buffer.dispatch(
+            qinput_tensor,
+            topk_idx=topk_ids,
+            topk_weights=topk_weights,
+            cumulative_local_expert_recv_stats=cumulative_stats,
+            num_experts=self.total_expert_num_contain_redundancy,
+            num_max_tokens_per_rank=num_max_tokens_per_rank,
+            expert_alignment=alignment,
+            do_cpu_sync=False,
+            do_handle_copy=False,
+            do_expand=True,
+            use_tma_aligned_col_major_sf=True,
+        )
+        gemm_out = self.prefilled_group_gemm(
+            handle.psum_num_recv_tokens_per_expert,
+            recv_x,
+            recv_topk_idx,
+            recv_topk_weights,
+            w13,
+            w2,
+            hidden_states.dtype,
+        )
+        combined_x, _, _ = buffer.combine(gemm_out, handle=handle, topk_weights=None)
+        return combined_x
+
     def _select_experts(
         self,
         input_tensor: torch.Tensor,
@@ -147,11 +195,17 @@ def _fused_experts(
         topk_ids: torch.Tensor,
         router_logits: Optional[torch.Tensor] = None,
         is_prefill: Optional[bool] = None,
+        is_cuda_graph: bool = False,
     ):
 
         w13_weight, w13_scale = w13.weight, w13.weight_scale
         w2_weight, w2_scale = w2.weight, w2.weight_scale
         if self._use_sm100_fp4_moe():
+            # DeepGEMM's official Mega MoE example is an eager fused path. For
+            # decode CUDA graph, use the official ElasticBuffer + grouped GEMM
+            # baseline instead of capturing Mega MoE's NVLink barrier kernel.
+            if is_cuda_graph and not is_prefill:
+                return self._sm100_fp4_cuda_graph_moe(input_tensor, w13, w2, topk_weights, topk_ids.to(torch.long))
             return self._mega_moe(input_tensor, w13, w2, topk_weights, topk_ids.to(torch.long))
 
         use_fp8_w8a8 = self.quant_method.method_name != "none"

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/marlin_impl.py

@@ -29,6 +29,7 @@ def _fused_experts(
         topk_ids: torch.Tensor,
         router_logits: Optional[torch.Tensor] = None,
         is_prefill: Optional[bool] = None,
+        is_cuda_graph: bool = False,
     ):
 
         w1_weight, w1_scale, w1_zero_point = w13.weight, w13.weight_scale, w13.weight_zero_point

lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py

@@ -91,6 +91,7 @@ def _fused_experts(
         topk_ids: torch.Tensor,
         router_logits: Optional[torch.Tensor] = None,
         is_prefill: bool = False,
+        is_cuda_graph: bool = False,
     ):
         w13_weight, w13_scale = w13.weight, w13.weight_scale
         w2_weight, w2_scale = w2.weight, w2.weight_scale
@@ -125,6 +126,7 @@ def __call__(
         topk_group: int,
         num_expert_group: int,
         is_prefill: Optional[bool] = None,
+        is_cuda_graph: bool = False,
     ):
         topk_weights, topk_ids = self._select_experts(
             input_tensor=input_tensor,
@@ -145,5 +147,6 @@ def __call__(
             topk_ids=topk_ids,
             router_logits=router_logits,
             is_prefill=is_prefill,
+            is_cuda_graph=is_cuda_graph,
         )
         return output

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

@@ -111,6 +111,7 @@ def __init__(self):
         self.ep_low_latency_buffer = None
         self.ep_mega_moe_buffer = None
         self.ep_num_sms = None
+        self.ep_expert_alignment = 128
 
     def __len__(self):
         return len(self.groups)
@@ -156,33 +157,26 @@ 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():
             if moe_intermediate_size is None:
                 raise ValueError("SM100 Mega MoE requires moe_intermediate_size or intermediate_size in model config")
 
             import deep_gemm
 
+            self.ep_expert_alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout()
+            deep_gemm.set_mk_alignment_for_contiguous_layout(self.ep_expert_alignment)
+            # Mega MoE is the eager fast path, while ElasticBuffer provides the official
+            # CUDA-graph-compatible baseline for decode.
+            self.ep_buffer = deep_ep.ElasticBuffer(
+                deepep_group,
+                num_max_tokens_per_rank=self.ll_decode_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 = deep_gemm.get_symm_buffer_for_mega_moe(
                 deepep_group,
                 self.ll_num_experts,
@@ -191,6 +185,28 @@ def new_deepep_group(
                 self.ll_hidden,
                 moe_intermediate_size,
             )
+            self._set_num_sms_for_deep_gemm(0)
+            logger.info("SM100 detected: use Mega MoE for eager path and ElasticBuffer for CUDA graph decode.")
+            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/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -258,6 +258,7 @@ def _moe_ffn_edp(
             topk_group=self.topk_group,
             num_expert_group=self.n_group,
             is_prefill=infer_state.is_prefill,
+            is_cuda_graph=infer_state.is_cuda_graph,
         )
 
         if self.n_shared_experts is not None:

lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py

@@ -104,6 +104,7 @@ def _moe_ffn_edp(
             topk_group=None,
             num_expert_group=None,
             is_prefill=infer_state.is_prefill,
+            is_cuda_graph=infer_state.is_cuda_graph,
         )
 
         ep_output = ep_output.view(token_num, hidden_dim)

lightllm/models/qwen3next/layer_infer/transformer_layer_infer.py

@@ -156,6 +156,7 @@ def _moe_ffn_edp(
             topk_group=None,
             num_expert_group=None,
             is_prefill=infer_state.is_prefill,
+            is_cuda_graph=infer_state.is_cuda_graph,
         )
         ep_output = ep_output.view(token_num, hidden_dim)
         ep_output.add_(shared_expert_out)