From 0c0d2013cc473570653d2280c7713306e52be0b1 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Thu, 16 Jul 2026 10:01:22 +0000 Subject: [PATCH 1/7] fix: reduce memory occupation --- lightllm/__init__.py | 4 + .../fused_moe/fused_moe_weight.py | 8 + .../fused_moe/impl/deepgemm_impl.py | 106 ++----- .../fused_moe/deepep_scatter_gather.py | 175 +++++++++++ .../fused_moe/grouped_fused_moe_ep.py | 287 +++++++++++++----- lightllm/distributed/communication_op.py | 29 +- .../layer_infer/transformer_layer_infer.py | 18 +- .../layer_infer/transformer_layer_infer.py | 18 +- 8 files changed, 466 insertions(+), 179 deletions(-) diff --git a/lightllm/__init__.py b/lightllm/__init__.py index e9ba6f3041..bc09ec5a17 100644 --- a/lightllm/__init__.py +++ b/lightllm/__init__.py @@ -2,3 +2,7 @@ if is_musa(): import torchada # noqa: F401 +else: + import torch + + torch._C._accelerator_setAllocatorSettings("expandable_segments:True") diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 26f2b338b7..63fdf32b68 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/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( diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 5da17c57e1..92ffa5ce6e 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/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,6 +256,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, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py b/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py index 101d316937..56ddacd98a 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/deepep_scatter_gather.py +++ 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, diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index cb2e370cb9..78031b405a 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -1,7 +1,8 @@ """Fused MoE kernel.""" import torch import triton -from typing import Any, Callable, Dict, Optional, Tuple +import triton.language as tl +from typing import Any, Callable, Dict, List, Optional, Tuple from lightllm.distributed import dist_group_manager from lightllm.utils.log_utils import init_logger from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd @@ -10,9 +11,12 @@ ) 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.deepep_scatter_gather import ( + ep_accumulate_expanded_chunk, + ep_fill_m_indices, + ep_zero_expanded_padding, +) from lightllm.utils.envs_utils import ( get_deepep_num_max_dispatch_tokens_per_rank_prefill, get_deepep_num_max_dispatch_tokens_per_rank_decode, @@ -75,12 +79,11 @@ def masked_group_gemm( expected_m = min(expected_m, padded_m) qsilu_out_scale = torch.empty((E, padded_m, N // 2 // block_size), device=recv_x[0].device, dtype=torch.float32) qsilu_out = torch.empty((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device) - # groupgemm (masked layout) - gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) - _deepgemm_grouped_fp8_nt_masked(recv_x, (w1, w1_scale), gemm_out_a, masked_m, expected_m) silu_and_mul_masked_post_quant_fwd(gemm_out_a, qsilu_out, qsilu_out_scale, block_size, masked_m) + del gemm_out_a + gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) _deepgemm_grouped_fp8_nt_masked((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, masked_m, expected_m) return gemm_out_b @@ -241,9 +244,6 @@ def fused_experts_impl( assert w2.is_contiguous(), "Expert weights2 must be contiguous" assert hidden_states.dtype in [torch.float32, torch.float16, torch.bfloat16] - M, K = hidden_states.shape - E, N, _ = w1.shape - # qaunt hidden_states assert use_fp8_w8a8 and use_fp8_all2all, "use_fp8_w8a8 and use_fp8_all2all must be True" @@ -258,11 +258,9 @@ def fused_experts_impl( if is_prefill: qinput_tensor, input_scale = per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w1.dtype) allocate_on_comm_stream = previous_event is not None - # normal dispatch - # recv_x [recive_num_tokens, hidden] recv_x_scale [recive_num_tokens, hidden // block_size] - # recv_topk_idx [recive_num_tokens, topk_num] - # recv_topk_weights [recive_num_tokens, topk_num] - # num_recv_tokens_per_expert_list list [cur_node_expert_num] padding with expert_alignment=128 + # Expanded dispatch directly produces expert-contiguous FP8 input and + # TMA-aligned scales for DeepGEMM. DeepEP also keeps the metadata needed + # to reduce the expanded W2 output in combine. recv_x, recv_topk_idx, recv_topk_weights, handle, _ = buffer.dispatch( (qinput_tensor, input_scale), topk_idx=topk_idx, @@ -274,76 +272,32 @@ def fused_experts_impl( allocate_on_comm_stream=allocate_on_comm_stream, do_cpu_sync=True, do_handle_copy=False, + do_expand=True, + use_tma_aligned_col_major_sf=True, + ) + # Dispatch is synchronous in this path. Its FP8 source is no longer + # needed once the received tensors have been produced. + del qinput_tensor, input_scale + + assert recv_topk_idx is None + gather_out = expanded_moe_chunked_reduce( + handle.num_recv_tokens_per_expert_list, + handle.num_unaligned_recv_tokens_per_expert, + recv_x, + recv_topk_weights, + handle.recv_src_metadata, + w1, + w1_scale, + w2, + w2_scale, + block_size_k, + get_prefill_moe_workspace(), + hidden_states.dtype, ) + del recv_x - # scatter - all_tokens = sum(handle.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=hidden_states.device, dtype=hidden_states.dtype) - if all_tokens > 0: - input_tensor = [ - torch.empty((all_tokens, K), device=hidden_states.device, dtype=qinput_tensor.dtype), - torch.empty((all_tokens, K // 128), device=hidden_states.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=hidden_states.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( - handle.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, - ) - - # groupgemm (contiguous layout) - gemm_out_a = torch.empty((all_tokens, N), device=hidden_states.device, dtype=hidden_states.dtype) - input_tensor[1] = tma_align_input_scale(input_tensor[1]) - deepgemm_grouped_fp8_nt_contiguous(input_tensor, (w1, w1_scale), gemm_out_a, m_indices) - - # silu_and_mul_fwd + qaunt - # TODO fused kernel - silu_out = torch.empty((all_tokens, N // 2), device=hidden_states.device, dtype=hidden_states.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_k, dtype=w1.dtype, column_major_scales=True, scale_tma_aligned=True - ) - - # groupgemm (contiguous layout) - gemm_out_b = torch.empty((all_tokens, K), device=hidden_states.device, dtype=hidden_states.dtype) - - deepgemm_grouped_fp8_nt_contiguous((qsilu_out, qsilu_out_scale), (w2, 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=hidden_states.device, dtype=hidden_states.dtype) - _silu_out = torch.zeros((1, N // 2), device=hidden_states.device, dtype=hidden_states.dtype) - silu_and_mul_fwd(_gemm_out_a.view(-1, N), _silu_out) - _gemm_out_a, _silu_out = None, None - - # normal combine + # W2 chunks were reduced to the deduplicated receive-token layout. + handle.do_expand = False combined_x, _, event = buffer.combine( gather_out, handle, @@ -387,6 +341,175 @@ def deepgemm_grouped_fp8_nt_contiguous( raise RuntimeError("deep_gemm does not provide grouped_gemm_fp8 NT contiguous GEMM kernel in this version") +def get_prefill_moe_workspace( + workspace_index: int = 0, + workspace_count: int = 1, +): + """Map prefill MoE temporaries onto the idle low-latency RDMA buffer. + + Prefill uses the ElasticBuffer while decode uses the legacy low-latency + buffer, so their communication phases are mutually exclusive. The model + clears the low-latency buffer after every prefill before decode can use it. + """ + + workspace = dist_group_manager.ep_prefill_workspace + assert 0 <= workspace_index < workspace_count + workspace_size = workspace.numel() // workspace_count + workspace = workspace.narrow(0, workspace_index * workspace_size, workspace_size) + return workspace + + +def expanded_moe_chunked_reduce( + num_recv_tokens_per_expert_list: List[int], + num_unaligned_recv_tokens_per_expert: torch.Tensor, + recv_x: Tuple[torch.Tensor, torch.Tensor], + recv_topk_weights: torch.Tensor, + recv_src_metadata: torch.Tensor, + w1: torch.Tensor, + w1_scale: torch.Tensor, + w2: torch.Tensor, + w2_scale: torch.Tensor, + block_size_k: int, + workspace: torch.Tensor, + hidden_dtype: torch.dtype, +): + """Run expanded W1/W2 in bounded chunks and reduce to dense rows.""" + all_tokens = sum(num_recv_tokens_per_expert_list) + assert all_tokens == recv_x[0].shape[0] + intermediate_twice = w1.shape[1] + intermediate_size = intermediate_twice // 2 + hidden_size = w2.shape[1] + if all_tokens == 0: + if Autotuner.is_autotune_warmup(): + gemm_out = torch.zeros((1, intermediate_twice), device=recv_x[0].device, dtype=hidden_dtype) + silu_out = torch.zeros((1, intermediate_size), device=recv_x[0].device, dtype=hidden_dtype) + silu_and_mul_fwd(gemm_out, silu_out) + return torch.empty((0, hidden_size), device=recv_x[0].device, dtype=hidden_dtype) + + m_indices = torch.empty(all_tokens, device=recv_x[0].device, dtype=torch.int32) + 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 = ep_fill_m_indices(num_recv_tokens_per_expert, m_indices) + ep_zero_expanded_padding( + recv_x[0], + recv_x[1], + recv_topk_weights, + num_recv_tokens_per_expert, + num_unaligned_recv_tokens_per_expert, + expert_start_loc, + ) + del num_recv_tokens_per_expert, expert_start_loc + gather_rows = recv_src_metadata.shape[0] + gather_bytes = gather_rows * hidden_size * hidden_dtype.itemsize + silu_row_bytes = intermediate_size * hidden_dtype.itemsize + gemm_a_row_bytes = intermediate_twice * hidden_dtype.itemsize + gemm_b_row_bytes = hidden_size * hidden_dtype.itemsize + quant_row_bytes = intermediate_size * w2.dtype.itemsize + scale_row_bytes = (intermediate_size // block_size_k) * torch.float32.itemsize + quant_with_scale_row_bytes = quant_row_bytes + scale_row_bytes + # The same region is reused in three non-overlapping phases: + # W1: [SwiGLU output][W1 output] + # quant: [SwiGLU output]...[FP8 output + TMA scales] + # W2: [W2 output]......[FP8 output + TMA scales] + # Keeping the quantized activation at the end lets W2 overwrite the old + # SwiGLU/W1 storage without allocating another tensor from the CUDA heap. + temp_row_bytes = max( + silu_row_bytes + gemm_a_row_bytes, + silu_row_bytes + quant_with_scale_row_bytes, + gemm_b_row_bytes + quant_with_scale_row_bytes, + ) + max_chunk_rows = ((workspace.numel() - gather_bytes) // temp_row_bytes // 128) * 128 + if max_chunk_rows <= 0: + raise RuntimeError( + f"DeepEP workspace cannot hold dense output: need {gather_bytes} bytes, have {workspace.numel()} bytes" + ) + + gather_out = workspace.narrow(0, 0, gather_bytes).view(hidden_dtype).view(gather_rows, hidden_size) + gather_out.zero_() + temp_offset = gather_bytes + + for chunk_start in range(0, all_tokens, max_chunk_rows): + chunk_end = min(chunk_start + max_chunk_rows, all_tokens) + chunk_rows = chunk_end - chunk_start + silu_bytes = chunk_rows * silu_row_bytes + gemm_a_bytes = chunk_rows * gemm_a_row_bytes + gemm_b_bytes = chunk_rows * gemm_b_row_bytes + quant_bytes = chunk_rows * quant_row_bytes + aligned_chunk_rows = (chunk_rows + 3) // 4 * 4 + scale_storage_shape = (intermediate_size // block_size_k, aligned_chunk_rows) + scale_bytes = scale_storage_shape[0] * scale_storage_shape[1] * torch.float32.itemsize + temp_bytes = chunk_rows * temp_row_bytes + silu_out = ( + workspace.narrow(0, temp_offset, silu_bytes) + .view(hidden_dtype) + .view(chunk_rows, intermediate_size) + ) + gemm_region_offset = temp_offset + silu_bytes + gemm_out_a = ( + workspace.narrow(0, gemm_region_offset, gemm_a_bytes) + .view(hidden_dtype) + .view(chunk_rows, intermediate_twice) + ) + deepgemm_grouped_fp8_nt_contiguous( + (recv_x[0][chunk_start:chunk_end], recv_x[1][chunk_start:chunk_end]), + (w1, w1_scale), + gemm_out_a, + m_indices[chunk_start:chunk_end], + ) + silu_and_mul_fwd(gemm_out_a, silu_out) + del gemm_out_a + + quant_offset = temp_offset + temp_bytes - quant_bytes - scale_bytes + qsilu_workspace = ( + workspace.narrow(0, quant_offset, quant_bytes) + .view(w2.dtype) + .view(chunk_rows, intermediate_size) + ) + scale_workspace = ( + workspace.narrow(0, quant_offset + quant_bytes, scale_bytes) + .view(torch.float32) + .view(scale_storage_shape) + ) + + def workspace_quant_alloc(shape, dtype, device): + if tuple(shape) == tuple(qsilu_workspace.shape) and dtype == qsilu_workspace.dtype: + return qsilu_workspace + if tuple(shape) == scale_storage_shape and dtype == torch.float32: + return scale_workspace + raise RuntimeError(f"unexpected prefill quant allocation: shape={shape}, dtype={dtype}") + + qsilu_out, qsilu_out_scale = per_token_group_quant_fp8( + silu_out, + block_size_k, + dtype=w2.dtype, + column_major_scales=True, + scale_tma_aligned=True, + alloc_func=workspace_quant_alloc, + ) + gemm_out_b = ( + workspace.narrow(0, temp_offset, gemm_b_bytes) + .view(hidden_dtype) + .view(chunk_rows, hidden_size) + ) + deepgemm_grouped_fp8_nt_contiguous( + (qsilu_out, qsilu_out_scale), + (w2, w2_scale), + gemm_out_b, + m_indices[chunk_start:chunk_end], + ) + del qsilu_out, qsilu_out_scale, silu_out + ep_accumulate_expanded_chunk( + gemm_out_b, + chunk_start, + recv_topk_weights, + recv_src_metadata, + gather_out, + ) + + return gather_out + + def _deepgemm_grouped_fp8_nt_masked( input_tuple: Tuple[torch.Tensor, torch.Tensor], w_tuple: Tuple[torch.Tensor, torch.Tensor], diff --git a/lightllm/distributed/communication_op.py b/lightllm/distributed/communication_op.py index f15badde25..83dd8932c5 100644 --- a/lightllm/distributed/communication_op.py +++ b/lightllm/distributed/communication_op.py @@ -109,6 +109,7 @@ def __init__(self): self.groups = [] self.ep_buffer = None self.ep_low_latency_buffer = None + self.ep_prefill_workspace = None self.ep_mega_moe_buffer = None self.ep_num_sms = None @@ -145,6 +146,7 @@ def new_deepep_group( if not enable_ep_moe: self.ep_buffer = None self.ep_low_latency_buffer = None + self.ep_prefill_workspace = None self.ep_mega_moe_buffer = None self.ep_num_sms = None return @@ -162,22 +164,23 @@ def new_deepep_group( hidden=self.ll_hidden, num_topk=num_experts_per_tok, use_fp8_dispatch=True, - allow_multiple_reduction=False, + allow_multiple_reduction=True, ) 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: + 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, + num_rdma_bytes=num_rdma_bytes, + low_latency_mode=True, + num_qps_per_rank=(self.ll_num_experts // global_world_size), + ) + self.ep_prefill_workspace = self.ep_low_latency_buffer.get_local_buffer_tensor( + torch.uint8, use_rdma_buffer=True + ) + 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") diff --git a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py index be819c94a0..bd93cd39ab 100644 --- a/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py @@ -499,7 +499,14 @@ def overlap_tpsp_context_forward( # 0 moe calu _0_moe_out = layer_weight.experts.prefilled_group_gemm( - _0_num_recv_tokens_per_expert_list, _0_recv_x, _0_recv_topk_idx, _0_recv_topk_weight + _0_num_recv_tokens_per_expert_list, + _0_handle.num_unaligned_recv_tokens_per_expert, + _0_handle.recv_src_metadata, + _0_recv_x, + _0_recv_topk_idx, + _0_recv_topk_weight, + workspace_index=0, + workspace_count=2, ) # 1 dispatch execute @@ -525,7 +532,14 @@ def overlap_tpsp_context_forward( # 1 moe calc _1_moe_out = layer_weight.experts.prefilled_group_gemm( - _1_num_recv_tokens_per_expert_list, _1_recv_x, _1_recv_topk_idx, _1_recv_topk_weight + _1_num_recv_tokens_per_expert_list, + _1_handle.num_unaligned_recv_tokens_per_expert, + _1_handle.recv_src_metadata, + _1_recv_x, + _1_recv_topk_idx, + _1_recv_topk_weight, + workspace_index=1, + workspace_count=2, ) # wait 0 combine diff --git a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py index 8879aa2d27..1698b850d1 100644 --- a/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_moe/layer_infer/transformer_layer_infer.py @@ -313,7 +313,14 @@ def overlap_tpsp_context_forward( # 0 moe calu _0_moe_out = layer_weight.experts.prefilled_group_gemm( - _0_num_recv_tokens_per_expert_list, _0_recv_x, _0_recv_topk_idx, _0_recv_topk_weight + _0_num_recv_tokens_per_expert_list, + _0_handle.num_unaligned_recv_tokens_per_expert, + _0_handle.recv_src_metadata, + _0_recv_x, + _0_recv_topk_idx, + _0_recv_topk_weight, + workspace_index=0, + workspace_count=2, ) # 1 dispatch execute @@ -339,7 +346,14 @@ def overlap_tpsp_context_forward( # 1 moe calc _1_moe_out = layer_weight.experts.prefilled_group_gemm( - _1_num_recv_tokens_per_expert_list, _1_recv_x, _1_recv_topk_idx, _1_recv_topk_weight + _1_num_recv_tokens_per_expert_list, + _1_handle.num_unaligned_recv_tokens_per_expert, + _1_handle.recv_src_metadata, + _1_recv_x, + _1_recv_topk_idx, + _1_recv_topk_weight, + workspace_index=1, + workspace_count=2, ) # wait 0 combine From 09d209eccfaeffcc674b9abd766d75708adff2a9 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Sat, 4 Jul 2026 08:06:02 +0000 Subject: [PATCH 2/7] feat: opt mega moe perf --- .../fused_moe/grouped_fused_moe_ep.py | 170 +++++++++++++++--- lightllm/common/quantization/__init__.py | 26 +-- lightllm/common/quantization/deepgemm.py | 23 ++- 3 files changed, 184 insertions(+), 35 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 78031b405a..cfab1d573b 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -46,6 +46,154 @@ def use_sm100_mega_moe(quant_method: Any) -> bool: return is_sm100_gpu() and quant_method.method_name == "deepgemm-fp4fp8-b32" +def _per_token_cast_to_fp8_packed_ue8m0(hidden_states: torch.Tensor, gran_k: int): + from deep_gemm.utils import per_token_cast_to_fp8 + + hidden_states, scale = per_token_cast_to_fp8( + hidden_states, + use_ue8m0=True, + gran_k=gran_k, + use_packed_ue8m0=False, + ) + assert scale.size(-1) % 4 == 0, "packed UE8M0 scale requires scale groups divisible by 4" + scale = (scale.view(torch.int32) >> 23).to(torch.uint8).view(torch.int32) + return hidden_states, scale + + +@triton.jit +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True), exp + + +@triton.jit +def _mega_moe_quant_topk_to_buffer_kernel( + x_ptr, + x_out_ptr, + x_sf_out_ptr, + topk_idx_ptr, + topk_idx_out_ptr, + topk_weights_ptr, + topk_weights_out_ptr, + stride_x_m: tl.constexpr, + stride_x_k: tl.constexpr, + stride_x_out_m: tl.constexpr, + stride_x_out_k: tl.constexpr, + stride_x_sf_out_m: tl.constexpr, + stride_x_sf_out_k: tl.constexpr, + stride_topk_idx_m: tl.constexpr, + stride_topk_idx_k: tl.constexpr, + stride_topk_idx_out_m: tl.constexpr, + stride_topk_idx_out_k: tl.constexpr, + stride_topk_weights_m: tl.constexpr, + stride_topk_weights_k: tl.constexpr, + stride_topk_weights_out_m: tl.constexpr, + stride_topk_weights_out_k: tl.constexpr, + FP8_MIN: tl.constexpr, + FP8_MAX: tl.constexpr, + TOPK: tl.constexpr, + GROUP_SIZE: tl.constexpr, + BLOCK: tl.constexpr, +): + token_id = tl.program_id(0) + pack_id = tl.program_id(1) + offsets = tl.arange(0, BLOCK) + cols = pack_id * BLOCK + offsets + + x = tl.load(x_ptr + token_id * stride_x_m + cols * stride_x_k).to(tl.float32) + abs_x = tl.abs(x) + group_id = offsets // GROUP_SIZE + + amax0 = tl.max(tl.where(group_id == 0, abs_x, 0.0)) + amax1 = tl.max(tl.where(group_id == 1, abs_x, 0.0)) + amax2 = tl.max(tl.where(group_id == 2, abs_x, 0.0)) + amax3 = tl.max(tl.where(group_id == 3, abs_x, 0.0)) + + scale0, exp0 = _ceil_to_ue8m0(tl.maximum(amax0, 1.0e-4) / FP8_MAX) + scale1, exp1 = _ceil_to_ue8m0(tl.maximum(amax1, 1.0e-4) / FP8_MAX) + scale2, exp2 = _ceil_to_ue8m0(tl.maximum(amax2, 1.0e-4) / FP8_MAX) + scale3, exp3 = _ceil_to_ue8m0(tl.maximum(amax3, 1.0e-4) / FP8_MAX) + + scale = tl.where( + group_id == 0, + scale0, + tl.where(group_id == 1, scale1, tl.where(group_id == 2, scale2, scale3)), + ) + x_q = tl.clamp(x / scale, FP8_MIN, FP8_MAX).to(x_out_ptr.dtype.element_ty) + tl.store(x_out_ptr + token_id * stride_x_out_m + cols * stride_x_out_k, x_q) + + packed_scale = exp0 | (exp1 << 8) | (exp2 << 16) | (exp3 << 24) + tl.store(x_sf_out_ptr + token_id * stride_x_sf_out_m + pack_id * stride_x_sf_out_k, packed_scale) + + if pack_id == 0: + topk_offsets = tl.arange(0, TOPK) + topk_idx = tl.load(topk_idx_ptr + token_id * stride_topk_idx_m + topk_offsets * stride_topk_idx_k) + topk_weights = tl.load( + topk_weights_ptr + token_id * stride_topk_weights_m + topk_offsets * stride_topk_weights_k + ) + tl.store( + topk_idx_out_ptr + token_id * stride_topk_idx_out_m + topk_offsets * stride_topk_idx_out_k, + topk_idx.to(topk_idx_out_ptr.dtype.element_ty), + ) + tl.store( + topk_weights_out_ptr + + token_id * stride_topk_weights_out_m + + topk_offsets * stride_topk_weights_out_k, + topk_weights.to(topk_weights_out_ptr.dtype.element_ty), + ) + + +def _prepare_mega_moe_buffer( + hidden_states: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + buffer: Any, + group_size: int, +): + num_tokens, hidden_size = hidden_states.shape + if num_tokens == 0: + return + assert hidden_size % (group_size * 4) == 0, "packed UE8M0 scale requires four FP8 groups per int32" + assert hidden_states.is_contiguous(), "hidden_states must be contiguous" + assert buffer.x.shape[0] >= num_tokens and buffer.x.shape[1] == hidden_size + assert buffer.x_sf.shape[0] >= num_tokens and buffer.x_sf.shape[1] == hidden_size // group_size // 4 + + block = group_size * 4 + finfo = torch.finfo(buffer.x.dtype) + _mega_moe_quant_topk_to_buffer_kernel[(num_tokens, hidden_size // block)]( + hidden_states, + buffer.x, + buffer.x_sf, + topk_ids, + buffer.topk_idx, + topk_weights, + buffer.topk_weights, + hidden_states.stride(0), + hidden_states.stride(1), + buffer.x.stride(0), + buffer.x.stride(1), + buffer.x_sf.stride(0), + buffer.x_sf.stride(1), + topk_ids.stride(0), + topk_ids.stride(1), + buffer.topk_idx.stride(0), + buffer.topk_idx.stride(1), + topk_weights.stride(0), + topk_weights.stride(1), + buffer.topk_weights.stride(0), + buffer.topk_weights.stride(1), + FP8_MIN=finfo.min, + FP8_MAX=finfo.max, + TOPK=topk_ids.shape[1], + GROUP_SIZE=group_size, + BLOCK=block, + num_warps=4, + num_stages=4, + ) + + def check_ep_expert_dtype(quant_method: Any): expert_dtype = getattr(quant_method, "method_name", None) if expert_dtype not in SUPPORTED_EP_EXPERT_DTYPES: @@ -126,8 +274,6 @@ def mega_moe_impl( if not (HAS_DEEPGEMM and hasattr(deep_gemm, "fp8_fp4_mega_moe")): raise RuntimeError("deep_gemm does not provide fp8-fp4 Mega MoE kernel") - from deep_gemm.utils import per_token_cast_to_fp8 - buffer = getattr(dist_group_manager, "ep_mega_moe_buffer", None) if buffer is None: raise RuntimeError("SM100 Mega MoE requires dist_group_manager.ep_mega_moe_buffer to be initialized") @@ -138,19 +284,10 @@ def mega_moe_impl( f"Mega MoE got {num_tokens} tokens, exceeding num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" ) - qinput_tensor = per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=quant_method.block_size, - use_packed_ue8m0=True, - ) state = _get_mega_moe_cache_state(w13, w2) l1_weights, l2_weights = _get_mega_moe_weights(w13, w2, state) stats = _get_mega_moe_cumulative_stats(w13.weight.shape[0], hidden_states.device, state) - buffer.x[:num_tokens].copy_(qinput_tensor[0]) - buffer.x_sf[:num_tokens].copy_(qinput_tensor[1]) - buffer.topk_idx[:num_tokens].copy_(topk_ids) - buffer.topk_weights[:num_tokens].copy_(topk_weights) + _prepare_mega_moe_buffer(hidden_states, topk_ids, topk_weights, buffer, quant_method.block_size) output = torch.empty_like(hidden_states) deep_gemm.fp8_fp4_mega_moe( @@ -170,14 +307,7 @@ def quantize_fused_experts_input( ): check_ep_expert_dtype(quant_method) if use_sm100_mega_moe(quant_method): - from deep_gemm.utils import per_token_cast_to_fp8 - - return per_token_cast_to_fp8( - hidden_states, - use_ue8m0=True, - gran_k=quant_method.block_size, - use_packed_ue8m0=True, - ) + return _per_token_cast_to_fp8_packed_ue8m0(hidden_states, quant_method.block_size) block_size_k = 0 if w13.weight.ndim == 3: diff --git a/lightllm/common/quantization/__init__.py b/lightllm/common/quantization/__init__.py index cd534d53ec..f676e35d64 100644 --- a/lightllm/common/quantization/__init__.py +++ b/lightllm/common/quantization/__init__.py @@ -43,6 +43,7 @@ def _parse_network_config(self, network_config): self.quantized_weight = False self.static_activation = False self.hf_quantization_config = None + self._mapping_expert_quant_method() return self.quantized_weight = True activation_scheme = network_config.get("activation_scheme", "dynamic") @@ -50,6 +51,19 @@ def _parse_network_config(self, network_config): self.hf_quantization_config = hf_quantization_config self.hf_quantization_method = hf_quantization_config["quant_method"] self._mapping_quant_method() + self._mapping_expert_quant_method() + + def _mapping_expert_quant_method(self): + expert_dtype = self.expert_dtype or self.network_config_.get("expert_dtype", None) + if expert_dtype is None: + return + target = self._get_expert_quant_type(expert_dtype) + for layer_num in range(self.layer_num): + if self.expert_dtype is not None: + self.quant_cfg[layer_num]["fused_moe"] = target + else: + self.quant_cfg[layer_num].setdefault("fused_moe", target) + logger.info(f"select fused_moe quant way from expert_dtype=`{expert_dtype}`: {target}") def _mapping_quant_method(self): if self.hf_quantization_method == "fp8": @@ -63,18 +77,6 @@ def _mapping_quant_method(self): self.quant_type = "vllm-fp8w8a8-b128" logger.info(f"select fp8w8a8-b128 quant way: {self.quant_type}") - # fp8 量化下，部分 MoE 模型（如 DeepSeek-V4），可以单独声明 expert 权重精度， - # 按其值给 fused_moe 选用对应的 deepgemm 量化方法。 - expert_dtype = self.expert_dtype or self.network_config_.get("expert_dtype", None) - if expert_dtype is None: - return - target = self._get_expert_quant_type(expert_dtype) - for layer_num in range(self.layer_num): - if self.expert_dtype is not None: - self.quant_cfg[layer_num]["fused_moe"] = target - else: - self.quant_cfg[layer_num].setdefault("fused_moe", target) - logger.info(f"select fused_moe quant way from expert_dtype=`{expert_dtype}`: {target}") elif self.hf_quantization_method == "awq": self.quant_type = "awq" if is_awq_marlin_compatible(self.hf_quantization_config): diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index ec1ee90fd4..328e9a71c9 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -182,11 +182,28 @@ def _create_weight( 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" + scales_per_int32 = 4 + scale_layout_k = self.block_size * scales_per_int32 + assert in_dim % scale_layout_k == 0, ( + f"FP4 required scale layout needs input dimension divisible by {scale_layout_k}" + ) 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( - device_id - ) + scale_dim = in_dim // scale_layout_k + if num_experts > 1: + weight_scale = torch.empty_strided( + (num_experts, out_dim, scale_dim), + (out_dim * scale_dim, 1, out_dim), + dtype=torch.int32, + device=f"cuda:{device_id}", + ) + else: + weight_scale = torch.empty_strided( + (out_dim, scale_dim), + (1, out_dim), + dtype=torch.int32, + device=f"cuda:{device_id}", + ) mm_param = WeightPack(weight=weight, weight_scale=weight_scale) mm_param_list = self._split_weight_pack( mm_param, From ba9b148919e0e0cb4351223b600d98cc1d35c54c Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Mon, 6 Jul 2026 12:19:58 +0000 Subject: [PATCH 3/7] feat: support UE8M0 --- .../quantization/fp8act_quant_kernel.py | 77 +++++++++++++------ .../fp8w8a8_block_quant_kernel.py | 35 +++++++-- lightllm/common/quantization/deepgemm.py | 3 +- 3 files changed, 84 insertions(+), 31 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py index 0a68372887..af4aff5c88 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/fp8act_quant_kernel.py @@ -14,6 +14,14 @@ pass +@triton.jit +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True) + + # Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py @triton.jit def _per_token_group_quant_fp8( @@ -25,21 +33,19 @@ def _per_token_group_quant_fp8( eps, fp8_min, fp8_max, - xs_m, xs_n, - xs_row_major: tl.constexpr, + xs_stride_m, + xs_stride_n, BLOCK: tl.constexpr, NEED_MASK: tl.constexpr, + USE_UE8M0_SCALE: tl.constexpr, ): g_id = tl.program_id(0) y_ptr += g_id * y_stride y_q_ptr += g_id * y_stride - if xs_row_major: - y_s_ptr += g_id - else: - row_id = g_id // xs_n - col_id = g_id % xs_n - y_s_ptr += col_id * xs_m + row_id # col major + row_id = g_id // xs_n + col_id = g_id % xs_n + y_s_ptr += row_id * xs_stride_m + col_id * xs_stride_n cols = tl.arange(0, BLOCK) # N <= BLOCK @@ -52,8 +58,11 @@ def _per_token_group_quant_fp8( y = tl.load(y_ptr + cols, mask=mask, other=other).to(tl.float32) # Quant - _absmax = tl.maximum(tl.max(tl.abs(y)), eps) - y_s = _absmax / fp8_max + _absmax = tl.max(tl.abs(y)) + if USE_UE8M0_SCALE: + y_s = _ceil_to_ue8m0(tl.maximum(_absmax, 1.0e-4) / fp8_max) + else: + y_s = tl.maximum(_absmax, eps) / fp8_max y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty) tl.store(y_q_ptr + cols, y_q, mask=mask) @@ -67,6 +76,7 @@ def lightllm_per_token_group_quant_fp8( x_s: torch.Tensor, eps: float = 1e-10, dtype: torch.dtype = torch.float8_e4m3fn, + use_ue8m0_scales: bool = False, ): """group-wise, per-token quantization on input tensor `x`. Args: @@ -80,8 +90,8 @@ def lightllm_per_token_group_quant_fp8( assert x.shape[-1] % group_size == 0, "the last dimension of `x` cannot be divisible by `group_size`" assert x.is_contiguous(), "`x` is not contiguous" - xs_row_major = x_s.is_contiguous() - xs_m, xs_n = x_s.shape + xs_n = x_s.shape[-1] + xs_stride_m, xs_stride_n = x_s.stride() finfo = torch.finfo(dtype) fp8_max = finfo.max @@ -102,11 +112,12 @@ def lightllm_per_token_group_quant_fp8( eps, fp8_min=fp8_min, fp8_max=fp8_max, - xs_m=xs_m, xs_n=xs_n, - xs_row_major=xs_row_major, + xs_stride_m=xs_stride_m, + xs_stride_n=xs_stride_n, BLOCK=BLOCK, NEED_MASK=BLOCK != group_size, + USE_UE8M0_SCALE=use_ue8m0_scales, num_warps=num_warps, num_stages=num_stages, ) @@ -121,12 +132,13 @@ def per_token_group_quant_fp8( column_major_scales: bool = False, scale_tma_aligned: bool = False, alloc_func: Callable = torch.empty, + use_ue8m0_scales: bool = False, ): x_q = alloc_func(x.shape, dtype=dtype, device=x.device) x_s = None # Adapted from # https://github.com/sgl-project/sglang/blob/7e257cd666c0d639626487987ea8e590da1e9395/python/sglang/srt/layers/quantization/fp8_kernel.py#L290 - if HAS_SGL_KERNEL: + if HAS_SGL_KERNEL and not use_ue8m0_scales: finfo = torch.finfo(dtype) fp8_max, fp8_min = finfo.max, finfo.min @@ -157,14 +169,35 @@ def per_token_group_quant_fp8( sgl_ops.sgl_per_token_group_quant_fp8(x, x_q, x_s, group_size, 1e-10, fp8_min, fp8_max, False, enable_v2=True) else: # 使用LightLLM kernel进行量化 - x_s = alloc_func( - x.shape[:-1] + (x.shape[-1] // group_size,), - device=x.device, - dtype=torch.float32, + if column_major_scales: + if scale_tma_aligned: + aligned_size = (x.shape[-2] + 3) // 4 * 4 + x_s = alloc_func( + x.shape[:-2] + (x.shape[-1] // group_size, aligned_size), + device=x.device, + dtype=torch.float32, + ).permute(-1, -2)[: x.shape[-2], :] + else: + x_s = alloc_func( + (x.shape[-1] // group_size,) + x.shape[:-1], + device=x.device, + dtype=torch.float32, + ).permute(-1, -2) + else: + x_s = alloc_func( + x.shape[:-1] + (x.shape[-1] // group_size,), + device=x.device, + dtype=torch.float32, + ) + lightllm_per_token_group_quant_fp8( + x, + group_size, + x_q, + x_s, + eps=eps, + dtype=dtype, + use_ue8m0_scales=use_ue8m0_scales, ) - lightllm_per_token_group_quant_fp8(x, group_size, x_q, x_s, eps=1e-10, dtype=torch.float8_e4m3fn) - if column_major_scales and scale_tma_aligned: - x_s = tma_align_input_scale(x_s) return x_q, x_s diff --git a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py index 3881cfe4b8..1c2caa967d 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py @@ -5,7 +5,15 @@ @triton.jit -def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True) + + +@triton.jit +def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr, USE_UE8M0_SCALE: tl.constexpr): pid_m = tl.program_id(axis=0) pid_n = tl.program_id(axis=1) n_blocks = tl.cdiv(N, BLOCK_SIZE) @@ -20,14 +28,21 @@ def weight_quant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr): amax = tl.max(tl.abs(x)) max_fp8e4m3_val = 448.0 - scale = amax / max_fp8e4m3_val - y = (x / (scale + 1e-6)).to(y_ptr.dtype.element_ty) + if USE_UE8M0_SCALE: + scale = _ceil_to_ue8m0(tl.maximum(amax, 1.0e-4) / max_fp8e4m3_val) + denom = scale + else: + scale = amax / max_fp8e4m3_val + denom = scale + 1e-6 + y = (x / denom).to(y_ptr.dtype.element_ty) tl.store(y_ptr + offs, y, mask=mask) tl.store(s_ptr + pid_m * n_blocks + pid_n, scale) -def mm_weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]: +def mm_weight_quant( + x: torch.Tensor, block_size: int = 128, use_ue8m0_scales: bool = False +) -> tuple[torch.Tensor, torch.Tensor]: assert x.is_contiguous(), "Input tensor must be contiguous" M, N = x.size() @@ -38,11 +53,15 @@ def mm_weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tenso s_scales = torch.empty((num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device) grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"])) - weight_quant_kernel[grid](x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size) + weight_quant_kernel[grid]( + x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size, USE_UE8M0_SCALE=use_ue8m0_scales + ) return y_quant, s_scales -def weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]: +def weight_quant( + x: torch.Tensor, block_size: int = 128, use_ue8m0_scales: bool = False +) -> tuple[torch.Tensor, torch.Tensor]: assert x.is_contiguous(), "Input tensor must be contiguous" x = x.cuda(get_current_device_id()) if x.dim() == 3: @@ -51,8 +70,8 @@ def weight_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, num_blocks_n = triton.cdiv(x.shape[2], block_size) s_scales = torch.empty((x.shape[0], num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device) for i in range(x.shape[0]): - y_quant[i], s_scales[i] = mm_weight_quant(x[i], block_size) + y_quant[i], s_scales[i] = mm_weight_quant(x[i], block_size, use_ue8m0_scales=use_ue8m0_scales) return y_quant, s_scales else: - y_quant, s_scales = mm_weight_quant(x, block_size) + y_quant, s_scales = mm_weight_quant(x, block_size, use_ue8m0_scales=use_ue8m0_scales) return y_quant, s_scales diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index 328e9a71c9..daaceadc36 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -62,7 +62,7 @@ def quantize(self, weight: torch.Tensor, output: WeightPack): from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_quant_kernel import weight_quant device = output.weight.device - weight, scale = weight_quant(weight.cuda(device), self.block_size) + weight, scale = weight_quant(weight.cuda(device), self.block_size, use_ue8m0_scales=True) output.weight.copy_(weight) output.weight_scale.copy_(scale) return @@ -90,6 +90,7 @@ def apply( column_major_scales=True, scale_tma_aligned=True, alloc_func=alloc_func, + use_ue8m0_scales=True, ) if out is None: From 6836cdf098784cbc95b10fa9cf077ce22af2869a Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Thu, 9 Jul 2026 13:05:14 +0800 Subject: [PATCH 4/7] feat: update deepep --- docker/Dockerfile | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 5acee29028..ff661fa1f3 100644 --- a/docker/Dockerfile +++ b/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; \ From f06c6735c5aaba666c2a307af4397f9fd43f3e43 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Mon, 13 Jul 2026 12:20:16 +0800 Subject: [PATCH 5/7] feat: mega moe weight in place --- .../layer_weights/base_layer_weight.py | 1 + .../layer_weights/meta_weights/base_weight.py | 3 ++ .../fused_moe/fused_moe_weight.py | 10 +++++ .../fused_moe/grouped_fused_moe_ep.py | 44 +++++++++---------- 4 files changed, 34 insertions(+), 24 deletions(-) diff --git a/lightllm/common/basemodel/layer_weights/base_layer_weight.py b/lightllm/common/basemodel/layer_weights/base_layer_weight.py index b1d992a7c4..7a49787e20 100644 --- a/lightllm/common/basemodel/layer_weights/base_layer_weight.py +++ b/lightllm/common/basemodel/layer_weights/base_layer_weight.py @@ -38,6 +38,7 @@ def verify_load(self): else: layer_num = None assert attr.verify_load(), f"Loading {attr_name} of layers {layer_num} fails." + attr.finalize_load() def _cuda(self, cpu_tensor): return cpu_tensor.contiguous().to(self.data_type_).cuda(get_current_device_id()) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py index 714e7acf48..d9f2b76ace 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/base_weight.py @@ -21,6 +21,9 @@ def _create_weight(self): def verify_load(self) -> bool: pass + def finalize_load(self) -> None: + pass + class BaseWeightTpl(BaseWeight): def __init__(self, tp_rank: int = None, tp_world_size: int = None, data_type: torch.dtype = None): diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 63fdf32b68..7dfe23800d 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -292,6 +292,16 @@ def verify_load(self): ) return weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok + def finalize_load(self): + if self.enable_ep_moe: + from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( + transform_mega_moe_weights_in_place, + use_sm100_mega_moe, + ) + + if use_sm100_mega_moe(self.quant_method): + transform_mega_moe_weights_in_place(self.w13, self.w2) + def _create_weight(self): intermediate_size = self.split_inter_size self.e_score_correction_bias = None diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index cfab1d573b..4d8952a501 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -25,7 +25,7 @@ from lightllm.utils.device_utils import is_sm100_gpu logger = init_logger(__name__) -_MEGA_MOE_STATES: Dict[Tuple[int, int, int, int], Dict[str, Any]] = {} +_MEGA_MOE_STATS: Dict[Tuple[int, int, int, int], torch.Tensor] = {} SUPPORTED_EP_EXPERT_DTYPES = ("deepgemm-fp8w8a8-b128", "deepgemm-fp4fp8-b32") try: @@ -138,9 +138,7 @@ def _mega_moe_quant_topk_to_buffer_kernel( topk_idx.to(topk_idx_out_ptr.dtype.element_ty), ) tl.store( - topk_weights_out_ptr - + token_id * stride_topk_weights_out_m - + topk_offsets * stride_topk_weights_out_k, + topk_weights_out_ptr + token_id * stride_topk_weights_out_m + topk_offsets * stride_topk_weights_out_k, topk_weights.to(topk_weights_out_ptr.dtype.element_ty), ) @@ -236,31 +234,29 @@ def masked_group_gemm( return gemm_out_b -def _get_mega_moe_cache_state(w13: Any, w2: Any): +def _get_mega_moe_cumulative_stats(w13: Any, w2: Any): state_key = ( w13.weight.data_ptr(), w13.weight_scale.data_ptr(), w2.weight.data_ptr(), w2.weight_scale.data_ptr(), ) - return _MEGA_MOE_STATES.setdefault(state_key, {}) - - -def _get_mega_moe_weights(w13: Any, w2: Any, state: Dict[str, Any]): - if "weight_cache" not in state: - state["weight_cache"] = deep_gemm.transform_weights_for_mega_moe( - (w13.weight, w13.weight_scale), - (w2.weight, w2.weight_scale), - ) - return state["weight_cache"] + stats = _MEGA_MOE_STATS.get(state_key) + if stats is None: + stats = torch.zeros((w13.weight.shape[0],), device=w13.weight.device, dtype=torch.int32) + _MEGA_MOE_STATS[state_key] = stats + return stats -def _get_mega_moe_cumulative_stats(num_local_experts: int, device: torch.device, state: Dict[str, Any]): - stats = state.get("stats") - if stats is None or stats.numel() != num_local_experts or stats.device != device: - stats = torch.zeros((num_local_experts,), device=device, dtype=torch.int32) - state["stats"] = stats - return stats +def transform_mega_moe_weights_in_place(w13: Any, w2: Any): + """Convert to Mega MoE layout without retaining a second weight copy.""" + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( + (w13.weight, w13.weight_scale), + (w2.weight, w2.weight_scale), + ) + w13.weight.copy_(transformed_l1[0]) + w13.weight_scale.copy_(transformed_l1[1]) + w2.weight_scale.copy_(transformed_l2[1]) def mega_moe_impl( @@ -284,9 +280,9 @@ def mega_moe_impl( f"Mega MoE got {num_tokens} tokens, exceeding num_max_tokens_per_rank={buffer.num_max_tokens_per_rank}" ) - state = _get_mega_moe_cache_state(w13, w2) - l1_weights, l2_weights = _get_mega_moe_weights(w13, w2, state) - stats = _get_mega_moe_cumulative_stats(w13.weight.shape[0], hidden_states.device, state) + l1_weights = (w13.weight, w13.weight_scale) + l2_weights = (w2.weight, w2.weight_scale) + stats = _get_mega_moe_cumulative_stats(w13, w2) _prepare_mega_moe_buffer(hidden_states, topk_ids, topk_weights, buffer, quant_method.block_size) output = torch.empty_like(hidden_states) From e8a82d5daf532be81be050bff5d952fd99ac9203 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 14 Jul 2026 14:53:53 +0800 Subject: [PATCH 6/7] refine --- .dockerignore | 30 +++++++++++++++ .../fused_moe/fused_moe_weight.py | 3 +- .../fused_moe/grouped_fused_moe_ep.py | 38 +++++++++---------- .../fp8w8a8_block_quant_kernel.py | 4 +- lightllm/common/quantization/deepgemm.py | 36 +++++++----------- 5 files changed, 63 insertions(+), 48 deletions(-) create mode 100644 .dockerignore diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000000..1ac2bb0d48 --- /dev/null +++ b/.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 diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py index 7dfe23800d..067b2e9ede 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/fused_moe_weight.py @@ -293,7 +293,7 @@ def verify_load(self): return weight_load_ok and per_expert_scale_load_ok and e_score_correction_bias_load_ok def finalize_load(self): - if self.enable_ep_moe: + if self.enable_ep_moe and not getattr(self, "_mega_moe_weights_transformed", False): from lightllm.common.basemodel.triton_kernel.fused_moe.grouped_fused_moe_ep import ( transform_mega_moe_weights_in_place, use_sm100_mega_moe, @@ -301,6 +301,7 @@ def finalize_load(self): if use_sm100_mega_moe(self.quant_method): transform_mega_moe_weights_in_place(self.w13, self.w2) + self._mega_moe_weights_transformed = True def _create_weight(self): intermediate_size = self.split_inter_size diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 4d8952a501..2f0c5332ce 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -2,7 +2,7 @@ import torch import triton import triton.language as tl -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, List, Optional, Tuple from lightllm.distributed import dist_group_manager from lightllm.utils.log_utils import init_logger from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd @@ -25,7 +25,6 @@ from lightllm.utils.device_utils import is_sm100_gpu logger = init_logger(__name__) -_MEGA_MOE_STATS: Dict[Tuple[int, int, int, int], torch.Tensor] = {} SUPPORTED_EP_EXPERT_DTYPES = ("deepgemm-fp8w8a8-b128", "deepgemm-fp4fp8-b32") try: @@ -94,6 +93,7 @@ def _mega_moe_quant_topk_to_buffer_kernel( FP8_MIN: tl.constexpr, FP8_MAX: tl.constexpr, TOPK: tl.constexpr, + TOPK_BLOCK: tl.constexpr, GROUP_SIZE: tl.constexpr, BLOCK: tl.constexpr, ): @@ -128,18 +128,25 @@ def _mega_moe_quant_topk_to_buffer_kernel( tl.store(x_sf_out_ptr + token_id * stride_x_sf_out_m + pack_id * stride_x_sf_out_k, packed_scale) if pack_id == 0: - topk_offsets = tl.arange(0, TOPK) - topk_idx = tl.load(topk_idx_ptr + token_id * stride_topk_idx_m + topk_offsets * stride_topk_idx_k) + topk_offsets = tl.arange(0, TOPK_BLOCK) + topk_mask = topk_offsets < TOPK + topk_idx = tl.load( + topk_idx_ptr + token_id * stride_topk_idx_m + topk_offsets * stride_topk_idx_k, + mask=topk_mask, + ) topk_weights = tl.load( - topk_weights_ptr + token_id * stride_topk_weights_m + topk_offsets * stride_topk_weights_k + topk_weights_ptr + token_id * stride_topk_weights_m + topk_offsets * stride_topk_weights_k, + mask=topk_mask, ) tl.store( topk_idx_out_ptr + token_id * stride_topk_idx_out_m + topk_offsets * stride_topk_idx_out_k, topk_idx.to(topk_idx_out_ptr.dtype.element_ty), + mask=topk_mask, ) tl.store( topk_weights_out_ptr + token_id * stride_topk_weights_out_m + topk_offsets * stride_topk_weights_out_k, topk_weights.to(topk_weights_out_ptr.dtype.element_ty), + mask=topk_mask, ) @@ -155,8 +162,12 @@ def _prepare_mega_moe_buffer( return assert hidden_size % (group_size * 4) == 0, "packed UE8M0 scale requires four FP8 groups per int32" assert hidden_states.is_contiguous(), "hidden_states must be contiguous" + assert topk_ids.shape == topk_weights.shape and topk_ids.shape[0] == num_tokens + assert topk_ids.shape[1] > 0 assert buffer.x.shape[0] >= num_tokens and buffer.x.shape[1] == hidden_size assert buffer.x_sf.shape[0] >= num_tokens and buffer.x_sf.shape[1] == hidden_size // group_size // 4 + assert buffer.topk_idx.shape[0] >= num_tokens and buffer.topk_idx.shape[1] == topk_ids.shape[1] + assert buffer.topk_weights.shape[0] >= num_tokens and buffer.topk_weights.shape[1] == topk_ids.shape[1] block = group_size * 4 finfo = torch.finfo(buffer.x.dtype) @@ -185,6 +196,7 @@ def _prepare_mega_moe_buffer( FP8_MIN=finfo.min, FP8_MAX=finfo.max, TOPK=topk_ids.shape[1], + TOPK_BLOCK=triton.next_power_of_2(topk_ids.shape[1]), GROUP_SIZE=group_size, BLOCK=block, num_warps=4, @@ -234,20 +246,6 @@ def masked_group_gemm( return gemm_out_b -def _get_mega_moe_cumulative_stats(w13: Any, w2: Any): - state_key = ( - w13.weight.data_ptr(), - w13.weight_scale.data_ptr(), - w2.weight.data_ptr(), - w2.weight_scale.data_ptr(), - ) - stats = _MEGA_MOE_STATS.get(state_key) - if stats is None: - stats = torch.zeros((w13.weight.shape[0],), device=w13.weight.device, dtype=torch.int32) - _MEGA_MOE_STATS[state_key] = stats - return stats - - def transform_mega_moe_weights_in_place(w13: Any, w2: Any): """Convert to Mega MoE layout without retaining a second weight copy.""" transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe( @@ -282,7 +280,6 @@ def mega_moe_impl( l1_weights = (w13.weight, w13.weight_scale) l2_weights = (w2.weight, w2.weight_scale) - stats = _get_mega_moe_cumulative_stats(w13, w2) _prepare_mega_moe_buffer(hidden_states, topk_ids, topk_weights, buffer, quant_method.block_size) output = torch.empty_like(hidden_states) @@ -291,7 +288,6 @@ def mega_moe_impl( l1_weights, l2_weights, buffer, - cumulative_local_expert_recv_stats=stats, ) return output diff --git a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py index 1c2caa967d..c711aae850 100644 --- a/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py +++ b/lightllm/common/basemodel/triton_kernel/quantization/fp8w8a8_block_quant_kernel.py @@ -53,9 +53,7 @@ def mm_weight_quant( s_scales = torch.empty((num_blocks_m, num_blocks_n), dtype=torch.float32, device=x.device) grid = lambda meta: (triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"])) - weight_quant_kernel[grid]( - x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size, USE_UE8M0_SCALE=use_ue8m0_scales - ) + weight_quant_kernel[grid](x, s_scales, y_quant, M, N, BLOCK_SIZE=block_size, USE_UE8M0_SCALE=use_ue8m0_scales) return y_quant, s_scales diff --git a/lightllm/common/quantization/deepgemm.py b/lightllm/common/quantization/deepgemm.py index daaceadc36..9d3043ce17 100644 --- a/lightllm/common/quantization/deepgemm.py +++ b/lightllm/common/quantization/deepgemm.py @@ -5,6 +5,7 @@ from lightllm.common.quantization.registry import QUANTMETHODS from lightllm.common.basemodel.triton_kernel.quantization.fp8act_quant_kernel import per_token_group_quant_fp8 from lightllm.utils.log_utils import init_logger +from lightllm.utils.device_utils import is_sm100_gpu logger = init_logger(__name__) @@ -62,7 +63,7 @@ def quantize(self, weight: torch.Tensor, output: WeightPack): from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_quant_kernel import weight_quant device = output.weight.device - weight, scale = weight_quant(weight.cuda(device), self.block_size, use_ue8m0_scales=True) + weight, scale = weight_quant(weight.cuda(device), self.block_size, use_ue8m0_scales=is_sm100_gpu()) output.weight.copy_(weight) output.weight_scale.copy_(scale) return @@ -90,7 +91,7 @@ def apply( column_major_scales=True, scale_tma_aligned=True, alloc_func=alloc_func, - use_ue8m0_scales=True, + use_ue8m0_scales=is_sm100_gpu(), ) if out is None: @@ -181,30 +182,19 @@ 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]]: 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" - scales_per_int32 = 4 - scale_layout_k = self.block_size * scales_per_int32 - assert in_dim % scale_layout_k == 0, ( - f"FP4 required scale layout needs input dimension divisible by {scale_layout_k}" - ) + scale_layout_k = self.block_size * 4 # Each int32 packs four UE8M0 scales. + assert ( + in_dim % scale_layout_k == 0 + ), f"FP4 required scale layout needs input dimension divisible by {scale_layout_k}" 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) scale_dim = in_dim // scale_layout_k - if num_experts > 1: - weight_scale = torch.empty_strided( - (num_experts, out_dim, scale_dim), - (out_dim * scale_dim, 1, out_dim), - dtype=torch.int32, - device=f"cuda:{device_id}", - ) - else: - weight_scale = torch.empty_strided( - (out_dim, scale_dim), - (1, out_dim), - dtype=torch.int32, - device=f"cuda:{device_id}", - ) + weight_scale = torch.empty_strided( + expert_prefix + (out_dim, scale_dim), + (out_dim * scale_dim, 1, out_dim) if num_experts > 1 else (1, out_dim), + dtype=torch.int32, + device=f"cuda:{device_id}", + ) mm_param = WeightPack(weight=weight, weight_scale=weight_scale) mm_param_list = self._split_weight_pack( mm_param, From f4efbc103008019563a04844f38d9f0dc3ece703 Mon Sep 17 00:00:00 2001 From: niushengxiao Date: Tue, 14 Jul 2026 18:29:26 +0800 Subject: [PATCH 7/7] fix: fix bugs for fp8 quant in sm100 --- .../fused_moe/impl/deepgemm_impl.py | 3 +++ .../fused_moe/grouped_fused_moe_ep.py | 20 ++++++++++++++++--- .../moe_silu_and_mul_mix_quant_ep.py | 16 ++++++++++++++- .../test_moe_silu_and_mul_mix_quant_ep.py | 13 ++++++++++-- 4 files changed, 46 insertions(+), 6 deletions(-) diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py index 92ffa5ce6e..07d5a506c0 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/deepgemm_impl.py @@ -16,6 +16,7 @@ quantize_fused_experts_input, ) from lightllm.common.basemodel.triton_kernel.redundancy_topk_ids_repair import redundancy_topk_ids_repair +from lightllm.utils.device_utils import is_sm100_gpu class FuseMoeDeepGEMM(FuseMoeTriton): @@ -119,6 +120,8 @@ def low_latency_dispatch( num_max_dispatch_tokens_per_rank=num_max_dispatch_tokens_per_rank, num_experts=self.total_expert_num_contain_redundancy, use_fp8=use_fp8_w8a8, + round_scale=is_sm100_gpu(), + use_ue8m0=is_sm100_gpu(), async_finish=False, return_recv_hook=True, ) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py index 2f0c5332ce..56f62661df 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/grouped_fused_moe_ep.py @@ -239,7 +239,14 @@ def masked_group_gemm( qsilu_out = torch.empty((E, padded_m, N // 2), dtype=w1.dtype, device=recv_x[0].device) _deepgemm_grouped_fp8_nt_masked(recv_x, (w1, w1_scale), gemm_out_a, masked_m, expected_m) - silu_and_mul_masked_post_quant_fwd(gemm_out_a, qsilu_out, qsilu_out_scale, block_size, masked_m) + silu_and_mul_masked_post_quant_fwd( + gemm_out_a, + qsilu_out, + qsilu_out_scale, + block_size, + masked_m, + use_ue8m0_scales=is_sm100_gpu(), + ) del gemm_out_a gemm_out_b = torch.empty_like(recv_x[0], device=recv_x[0].device, dtype=dtype) _deepgemm_grouped_fp8_nt_masked((qsilu_out, qsilu_out_scale), (w2, w2_scale), gemm_out_b, masked_m, expected_m) @@ -305,7 +312,9 @@ def quantize_fused_experts_input( if w13.weight.ndim == 3: block_size_k = w13.weight.shape[2] // w13.weight_scale.shape[2] assert block_size_k == 128, "block_size_k must be 128" - return per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w13.weight.dtype) + return per_token_group_quant_fp8( + hidden_states, block_size_k, dtype=w13.weight.dtype, use_ue8m0_scales=is_sm100_gpu() + ) def fused_experts( @@ -378,7 +387,9 @@ def fused_experts_impl( combined_x = None if is_prefill: - qinput_tensor, input_scale = per_token_group_quant_fp8(hidden_states, block_size_k, dtype=w1.dtype) + qinput_tensor, input_scale = per_token_group_quant_fp8( + hidden_states, block_size_k, dtype=w1.dtype, use_ue8m0_scales=is_sm100_gpu() + ) allocate_on_comm_stream = previous_event is not None # Expanded dispatch directly produces expert-contiguous FP8 input and # TMA-aligned scales for DeepGEMM. DeepEP also keeps the metadata needed @@ -437,6 +448,8 @@ def fused_experts_impl( num_max_dispatch_tokens_per_rank, num_experts, use_fp8=use_fp8_w8a8, + round_scale=is_sm100_gpu(), + use_ue8m0=is_sm100_gpu(), async_finish=False, return_recv_hook=False, ) @@ -608,6 +621,7 @@ def workspace_quant_alloc(shape, dtype, device): column_major_scales=True, scale_tma_aligned=True, alloc_func=workspace_quant_alloc, + use_ue8m0_scales=is_sm100_gpu(), ) gemm_out_b = ( workspace.narrow(0, temp_offset, gemm_b_bytes) diff --git a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py index aa91f15ed9..7380cd1b78 100644 --- a/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py +++ b/lightllm/common/basemodel/triton_kernel/fused_moe/moe_silu_and_mul_mix_quant_ep.py @@ -6,6 +6,14 @@ from lightllm.utils.config_utils import ffn_use_tanh_approximate_gelu +@triton.jit +def _ceil_to_ue8m0(x): + bits = x.to(tl.float32).to(tl.int32, bitcast=True) + exp = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0) + exp = tl.maximum(tl.minimum(exp, 254), 1) + return (exp << 23).to(tl.float32, bitcast=True) + + @triton.jit def _silu_and_mul_post_quant_kernel( input_ptr, @@ -26,6 +34,7 @@ def _silu_and_mul_post_quant_kernel( fp8_min, BLOCK_N: tl.constexpr, NUM_STAGE: tl.constexpr, + USE_UE8M0_SCALE: tl.constexpr, USE_TANH_APPROXIMATE_GELU: tl.constexpr = False, ): expert_id = tl.program_id(2) @@ -61,7 +70,10 @@ def _silu_and_mul_post_quant_kernel( gate = gate.to(input_ptr.dtype.element_ty) gate_up = up * gate _absmax = tl.maximum(tl.max(tl.abs(gate_up)), 1e-10) - output_s = _absmax / fp8_max + if USE_UE8M0_SCALE: + output_s = _ceil_to_ue8m0(tl.maximum(_absmax, 1.0e-4) / fp8_max) + else: + output_s = _absmax / fp8_max output_q = tl.clamp(gate_up / output_s, fp8_min, fp8_max).to(output_ptr.dtype.element_ty) tl.store( output_ptr_offs + token_index * stride_output_1, @@ -80,6 +92,7 @@ def silu_and_mul_masked_post_quant_fwd( output_scale: torch.Tensor, quant_group_size: int, masked_m: torch.Tensor, + use_ue8m0_scales: bool = False, ): """ input shape [expert_num, token_num_padded, hidden_dim] @@ -135,6 +148,7 @@ def silu_and_mul_masked_post_quant_fwd( fp8_min, BLOCK_N=BLOCK_N, NUM_STAGE=NUM_STAGES, + USE_UE8M0_SCALE=use_ue8m0_scales, USE_TANH_APPROXIMATE_GELU=ffn_use_tanh_approximate_gelu(), num_warps=num_warps, ) diff --git a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py index 8783f35a42..3bebd29198 100644 --- a/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py +++ b/unit_tests/common/fused_moe/test_moe_silu_and_mul_mix_quant_ep.py @@ -36,7 +36,8 @@ def is_fp8_native_supported(): for token_num in range(1, 7, 2) ], ) -def test_silu_and_mul_masked(expert_num, token_num, hidden_dim): +@pytest.mark.parametrize("use_ue8m0_scales", [False, True]) +def test_silu_and_mul_masked(expert_num, token_num, hidden_dim, use_ue8m0_scales): quant_group_size = 128 in_tensor = torch.randn((expert_num, token_num, hidden_dim), dtype=torch.bfloat16, device="cuda") out_tensor = torch.empty((expert_num, token_num, hidden_dim // 2), dtype=torch.float8_e4m3fn, device="cuda") @@ -53,9 +54,17 @@ def test_silu_and_mul_masked(expert_num, token_num, hidden_dim): true_out_tensor_mid.view(-1, hidden_dim // 2), quant_group_size, alloc_func=torch.empty, + use_ue8m0_scales=use_ue8m0_scales, ) - silu_and_mul_masked_post_quant_fwd(in_tensor, out_tensor, out_scale_tensor, quant_group_size, masked_m) + silu_and_mul_masked_post_quant_fwd( + in_tensor, + out_tensor, + out_scale_tensor, + quant_group_size, + masked_m, + use_ue8m0_scales=use_ue8m0_scales, + ) true_out_tensor = true_out_tensor.view(out_tensor.shape) true_out_scale_tensor = true_out_scale_tensor.view(out_scale_tensor.shape)