fix: reduce memory occupation

Author: niushengxiao

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
 from lightllm.utils.device_utils import is_sm100_gpu
 
@@ -182,6 +176,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():
@@ -214,92 +210,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,
-                use_ue8m0_scales=is_sm100_gpu(),
-            )
-
-            # 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(
@@ -320,6 +259,8 @@ def combine(
         handle: Any,
         overlap_event: Optional[Any] = None,
     ):
+        # Chunked W2 has already reduced expanded expert slots to dense rows.
+        handle.do_expand = False
         # normal combine
         combined_x, _, event = dist_group_manager.ep_buffer.combine(
             gemm_out_b,

lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py

@@ -152,6 +152,181 @@ def ep_scatter(
     return
 
 
+@torch.no_grad()
+def ep_fill_m_indices(
+    num_recv_tokens_per_expert: torch.Tensor,
+    m_indices: torch.Tensor,
+):
+    """Build DeepGEMM's contiguous expert index vector without scattering data."""
+    block_e = 128
+    num_experts = num_recv_tokens_per_expert.shape[0]
+    assert m_indices.shape[0] % block_e == 0
+
+    expert_start_loc = torch.empty_like(num_recv_tokens_per_expert)
+    _fwd_kernel_ep_scatter_1[(num_experts,)](
+        num_recv_tokens_per_expert,
+        expert_start_loc,
+        m_indices,
+        num_experts=num_experts,
+        num_warps=8,
+        BLOCK_E=block_e,
+        BLOCK_EXPERT_NUM=triton.next_power_of_2(num_experts),
+    )
+    return expert_start_loc
+
+
+@triton.jit
+def _zero_expanded_padding_kernel(
+    recv_x,
+    recv_x_stride_m,
+    recv_x_stride_k,
+    recv_x_scale,
+    recv_x_scale_stride_m,
+    recv_x_scale_stride_k,
+    recv_topk_weights,
+    num_recv_tokens_per_expert,
+    num_unaligned_recv_tokens_per_expert,
+    expert_start_loc,
+    hidden_size: tl.constexpr,
+    scale_hidden_size: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_SCALE_K: tl.constexpr,
+):
+    expert_id = tl.program_id(0)
+    pad_block_id = tl.program_id(1)
+    hidden_block_id = tl.program_id(2)
+    expert_start = tl.load(expert_start_loc + expert_id)
+    aligned_count = tl.load(num_recv_tokens_per_expert + expert_id)
+    actual_count = tl.load(num_unaligned_recv_tokens_per_expert + expert_id)
+    pad_offsets = pad_block_id * BLOCK_M + tl.arange(0, BLOCK_M)
+    row_offsets = (expert_start + actual_count + pad_offsets).to(tl.int64)
+    row_mask = pad_offsets < aligned_count - actual_count
+
+    hidden_offsets = hidden_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
+    x_ptrs = recv_x + row_offsets[:, None] * recv_x_stride_m + hidden_offsets[None, :] * recv_x_stride_k
+    tl.store(x_ptrs, 0.0, mask=row_mask[:, None] & (hidden_offsets[None, :] < hidden_size))
+    if hidden_block_id == 0:
+        scale_offsets = tl.arange(0, BLOCK_SCALE_K)
+        scale_ptrs = (
+            recv_x_scale
+            + row_offsets[:, None] * recv_x_scale_stride_m
+            + scale_offsets[None, :] * recv_x_scale_stride_k
+        )
+        tl.store(scale_ptrs, 0.0, mask=row_mask[:, None] & (scale_offsets[None, :] < scale_hidden_size))
+        tl.store(recv_topk_weights + row_offsets, 0.0, mask=row_mask)
+
+
+@torch.no_grad()
+def ep_zero_expanded_padding(
+    recv_x: torch.Tensor,
+    recv_x_scale: torch.Tensor,
+    recv_topk_weights: torch.Tensor,
+    num_recv_tokens_per_expert: torch.Tensor,
+    num_unaligned_recv_tokens_per_expert: torch.Tensor,
+    expert_start_loc: torch.Tensor,
+):
+    block_m = 8
+    block_k = 256
+    scale_hidden_size = recv_x_scale.shape[1]
+    grid = (
+        num_recv_tokens_per_expert.shape[0],
+        triton.cdiv(127, block_m),
+        triton.cdiv(recv_x.shape[1], block_k),
+    )
+    _zero_expanded_padding_kernel[grid](
+        recv_x,
+        recv_x.stride(0),
+        recv_x.stride(1),
+        recv_x_scale,
+        recv_x_scale.stride(0),
+        recv_x_scale.stride(1),
+        recv_topk_weights,
+        num_recv_tokens_per_expert,
+        num_unaligned_recv_tokens_per_expert,
+        expert_start_loc,
+        hidden_size=recv_x.shape[1],
+        scale_hidden_size=scale_hidden_size,
+        BLOCK_M=block_m,
+        BLOCK_K=block_k,
+        BLOCK_SCALE_K=triton.next_power_of_2(scale_hidden_size),
+        num_warps=4,
+    )
+
+
+@triton.jit
+def _accumulate_expanded_chunk_kernel(
+    total_recv_tokens,
+    chunk,
+    chunk_stride_m,
+    chunk_stride_k,
+    chunk_start,
+    chunk_end,
+    weights,
+    recv_src_metadata,
+    metadata_stride_m,
+    metadata_stride_k,
+    output,
+    output_stride_m,
+    output_stride_k,
+    TOPK: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    hidden_block_id = tl.program_id(0)
+    start_recv_token_id = tl.program_id(1)
+    recv_token_grid_size = tl.num_programs(1)
+    hidden_offsets = hidden_block_id * BLOCK_D + tl.arange(0, BLOCK_D)
+
+    for recv_token_id in range(start_recv_token_id, total_recv_tokens, recv_token_grid_size):
+        output_ptrs = output + recv_token_id * output_stride_m + hidden_offsets * output_stride_k
+        accumulator = tl.load(output_ptrs).to(tl.float32)
+        for topk_id in range(TOPK):
+            slot = tl.load(
+                recv_src_metadata
+                + recv_token_id * metadata_stride_m
+                + (topk_id + 2) * metadata_stride_k
+            )
+            if slot >= chunk_start and slot < chunk_end:
+                local_row = (slot - chunk_start).to(tl.int64)
+                value = tl.load(chunk + local_row * chunk_stride_m + hidden_offsets * chunk_stride_k)
+                weight = tl.load(weights + slot)
+                accumulator += value.to(tl.float32) * weight
+        tl.store(output_ptrs, accumulator)
+
+
+@torch.no_grad()
+def ep_accumulate_expanded_chunk(
+    chunk: torch.Tensor,
+    chunk_start: int,
+    weights: torch.Tensor,
+    recv_src_metadata: torch.Tensor,
+    output: torch.Tensor,
+):
+    """Accumulate one contiguous expanded W2 chunk into dense receive-token rows."""
+    topk = recv_src_metadata.shape[1] - 2
+    block_d = 1024
+    assert chunk.shape[1] == output.shape[1] and output.shape[1] % block_d == 0
+    grid = (triton.cdiv(output.shape[1], block_d), min(output.shape[0], 1024))
+    _accumulate_expanded_chunk_kernel[grid](
+        output.shape[0],
+        chunk,
+        chunk.stride(0),
+        chunk.stride(1),
+        chunk_start,
+        chunk_start + chunk.shape[0],
+        weights,
+        recv_src_metadata,
+        recv_src_metadata.stride(0),
+        recv_src_metadata.stride(1),
+        output,
+        output.stride(0),
+        output.stride(1),
+        TOPK=topk,
+        BLOCK_D=block_d,
+        num_warps=2,
+    )
+
+
 @triton.jit
 def _fwd_kernel_ep_gather(
     total_token_num,