fix: reduce memory occupation and raise performance for multi-nodes

.dockerignore

@@ -0,0 +1,30 @@
+.git
+.github
+.conda
+.venv
+.idea
+.vscode
+
+__pycache__
+*.py[cod]
+.pytest_cache
+.mypy_cache
+.ruff_cache
+
+build
+dist
+*.egg-info
+docs
+test
+unit_tests
+benchmark
+logs
+tmp
+
+*.bin
+*.ckpt
+*.gguf
+*.onnx
+*.pt
+*.pth
+*.safetensors

docker/Dockerfile

@@ -7,9 +7,8 @@ ARG VLLM_VERSION=0.21.0
 ARG NIXL_REF=v1.2.0
 ARG FLASH_MLA_REF=47c35a7
 ARG DEEPGEMM_REF=891d57b4db1071624b5c8fa0d1e51cb317fa709f
-ARG DEEPEP_REF=099d5f2bad488b9c534ea785062b12f2e91d1d41
+ARG DEEPEP_REF=60d44037a702f651a6e18bd4aea65ed8409051c2
 ARG DEEPEP_NCCL_VERSION=2.30.4
-ARG DEEPEP_NVSHMEM_VERSION=3.3.24
 ARG TARGETPLATFORM
 ARG ENABLE_DEEPEP=1
 ARG ENABLE_NIXL=1
@@ -92,8 +91,7 @@ RUN if [ "${ENABLE_DEEPEP}" = "1" ]; then \
       set -e; \
       ln -sf /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so; \
       python -m pip install --upgrade --no-deps \
-        "nvidia-nccl-cu13==${DEEPEP_NCCL_VERSION}" \
-        "nvidia-nvshmem-cu13==${DEEPEP_NVSHMEM_VERSION}"; \
+        "nvidia-nccl-cu13==${DEEPEP_NCCL_VERSION}"; \
       cd /root && git clone https://github.com/deepseek-ai/DeepEP.git && cd DeepEP && git checkout ${DEEPEP_REF}; \
       ln -sf /opt/conda/lib/python${PYTHON_VERSION}/site-packages/nvidia/nvshmem/lib/libnvshmem_host.so.3 /opt/conda/lib/python${PYTHON_VERSION}/site-packages/nvidia/nvshmem/lib/libnvshmem_host.so; \
       ln -sf /opt/conda/lib/python${PYTHON_VERSION}/site-packages/nvidia/nccl/lib/libnccl.so.2 /opt/conda/lib/python${PYTHON_VERSION}/site-packages/nvidia/nccl/lib/libnccl.so; \

lightllm/__init__.py

@@ -2,3 +2,7 @@
 
 if is_musa():
     import torchada  # noqa: F401
+else:
+    import torch
+
+    torch._C._accelerator_setAllocatorSettings("expandable_segments:True")

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

@@ -222,20 +222,28 @@ def masked_group_gemm(
     def prefilled_group_gemm(
         self,
         num_recv_tokens_per_expert_list,
+        num_unaligned_recv_tokens_per_expert: torch.Tensor,
+        recv_src_metadata: torch.Tensor,
         recv_x: Tuple[torch.Tensor],
         recv_topk_idx: torch.Tensor,
         recv_topk_weights: torch.Tensor,
         hidden_dtype=torch.bfloat16,
+        workspace_index: int = 0,
+        workspace_count: int = 1,
     ):
         assert self.enable_ep_moe, "prefilled_group_gemm is only supported when enable_ep_moe is True"
         return self.fuse_moe_impl.prefilled_group_gemm(
             num_recv_tokens_per_expert_list=num_recv_tokens_per_expert_list,
+            num_unaligned_recv_tokens_per_expert=num_unaligned_recv_tokens_per_expert,
+            recv_src_metadata=recv_src_metadata,
             recv_x=recv_x,
             recv_topk_idx=recv_topk_idx,
             recv_topk_weights=recv_topk_weights,
             w13=self.w13,
             w2=self.w2,
             hidden_dtype=hidden_dtype,
+            workspace_index=workspace_index,
+            workspace_count=workspace_count,
         )
 
     def low_latency_combine(

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

@@ -2,7 +2,6 @@
 from typing import Optional, Tuple, Any
 from .triton_impl import FuseMoeTriton
 from lightllm.distributed import dist_group_manager
-from lightllm.common.triton_utils.autotuner import Autotuner
 from lightllm.common.quantization.quantize_method import WeightPack
 from lightllm.utils.envs_utils import (
     get_deepep_num_max_dispatch_tokens_per_rank_prefill,
@@ -12,15 +11,10 @@
     fused_experts,
     get_ep_num_sms,
     masked_group_gemm,
-    deepgemm_grouped_fp8_nt_contiguous,
+    get_prefill_moe_workspace,
+    expanded_moe_chunked_reduce,
     quantize_fused_experts_input,
 )
-from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import (
-    per_token_group_quant_fp8,
-    tma_align_input_scale,
-)
-from lightllm.common.basemodel.triton_kernel.fused_moe.deepep_scatter_gather import ep_scatter, ep_gather
-from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd
 from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair
 
 
@@ -179,6 +173,8 @@ def dispatch(
             allocate_on_comm_stream=True,
             do_cpu_sync=True,
             do_handle_copy=False,
+            do_expand=True,
+            use_tma_aligned_col_major_sf=True,
         )
 
         def hook():
@@ -211,87 +207,35 @@ def masked_group_gemm(
     def prefilled_group_gemm(
         self,
         num_recv_tokens_per_expert_list,
+        num_unaligned_recv_tokens_per_expert: torch.Tensor,
+        recv_src_metadata: torch.Tensor,
         recv_x: Tuple[torch.Tensor],
         recv_topk_idx: torch.Tensor,
         recv_topk_weights: torch.Tensor,
         w13: WeightPack,
         w2: WeightPack,
         hidden_dtype=torch.bfloat16,
+        workspace_index: int = 0,
+        workspace_count: int = 1,
     ):
-        device = recv_x[0].device
         w13_weight, w13_scale = w13.weight, w13.weight_scale
         w2_weight, w2_scale = w2.weight, w2.weight_scale
-        _, K = recv_x[0].shape
-        _, N, _ = w13_weight.shape
-        block_size = self.quant_method.block_size
-        # scatter
-        all_tokens = sum(num_recv_tokens_per_expert_list)  # calcu padding all nums.
-        # gather_out shape [recive_num_tokens, hidden]
-        gather_out = torch.empty_like(recv_x[0], device=device, dtype=hidden_dtype)
-        if all_tokens > 0:
-            input_tensor = [
-                torch.empty((all_tokens, K), device=device, dtype=recv_x[0].dtype),
-                torch.empty((all_tokens, K // 128), device=device, dtype=torch.float32),
-            ]
-            # when m_indices is filled ok.
-            # m_indices show token use which expert, example, [0, 0, 0, 0, .... 1, 1, 1, 1,...., cur_expert_num - 1, ..]
-            # the count of 0 is num_recv_tokens_per_expert_list[0], the count of 1 is num_recv_tokens_per_expert_list[1]
-            # ...
-            m_indices = torch.empty(all_tokens, device=device, dtype=torch.int32)
-            # output_index shape [recive_num_tokens, topk_num]
-            # output_index use to show the token index in input_tensor
-            output_index = torch.empty_like(recv_topk_idx)
-
-            num_recv_tokens_per_expert = torch.tensor(
-                num_recv_tokens_per_expert_list, dtype=torch.int32, pin_memory=True, device="cpu"
-            ).cuda(non_blocking=True)
-
-            expert_start_loc = torch.empty_like(num_recv_tokens_per_expert)
-
-            ep_scatter(
-                recv_x[0],
-                recv_x[1],
-                recv_topk_idx,
-                num_recv_tokens_per_expert,
-                expert_start_loc,
-                input_tensor[0],
-                input_tensor[1],
-                m_indices,
-                output_index,
-            )
-            input_tensor[1] = tma_align_input_scale(input_tensor[1])
-            # groupgemm (contiguous layout)
-            gemm_out_a = torch.empty((all_tokens, N), device=device, dtype=hidden_dtype)
-
-            deepgemm_grouped_fp8_nt_contiguous(input_tensor, (w13_weight, w13_scale), gemm_out_a, m_indices)
-
-            # silu_and_mul_fwd + qaunt
-            # TODO fused kernel
-            silu_out = torch.empty((all_tokens, N // 2), device=device, dtype=hidden_dtype)
-
-            silu_and_mul_fwd(gemm_out_a.view(-1, N), silu_out)
-            qsilu_out, qsilu_out_scale = per_token_group_quant_fp8(
-                silu_out, block_size, dtype=w13_weight.dtype, column_major_scales=True, scale_tma_aligned=True
-            )
-
-            # groupgemm (contiguous layout)
-            gemm_out_b = torch.empty((all_tokens, K), device=device, dtype=hidden_dtype)
-
-            deepgemm_grouped_fp8_nt_contiguous(
-                (qsilu_out, qsilu_out_scale), (w2_weight, w2_scale), gemm_out_b, m_indices
-            )
-            # gather and local reduce
-            ep_gather(gemm_out_b, recv_topk_idx, recv_topk_weights, output_index, gather_out)
-        else:
-            ######################################## warning ##################################################
-            # here is used to match autotune feature, make moe model run same triton kernel in different rank.
-            # in some special case, one rank will recv 0 token, so add a token to make it run triton kernel.
-            if Autotuner.is_autotune_warmup():
-                _gemm_out_a = torch.zeros((1, N), device=device, dtype=hidden_dtype)
-                _silu_out = torch.zeros((1, N // 2), device=device, dtype=hidden_dtype)
-                silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out)
-                _gemm_out_a, _silu_out = None, None
-
+        assert recv_topk_idx is None
+        gather_out = expanded_moe_chunked_reduce(
+            num_recv_tokens_per_expert_list,
+            num_unaligned_recv_tokens_per_expert,
+            recv_x,
+            recv_topk_weights,
+            recv_src_metadata,
+            w13_weight,
+            w13_scale,
+            w2_weight,
+            w2_scale,
+            self.quant_method.block_size,
+            get_prefill_moe_workspace(workspace_index, workspace_count),
+            hidden_dtype,
+        )
+        del recv_x
         return gather_out
 
     def low_latency_combine(
@@ -312,7 +256,8 @@ def combine(
         handle: Any,
         overlap_event: Optional[Any] = None,
     ):
-        # normal combine
+        # The prefill kernel keeps expanded routing metadata while pointing its
+        # single valid slot at each pre-reduced dense row.
         combined_x, _, event = dist_group_manager.ep_buffer.combine(
             gemm_out_b,
             handle,