From 90849840a6d37dde74567835e1e94637454d5a5b Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Mon, 2 Nov 2026 17:58:40 +0800 Subject: [PATCH 1/3] perf(dsv4): overlap shared experts with DeepEP combine --- .../meta_weights/fused_moe/fused_moe_weight.py | 4 +++- .../meta_weights/fused_moe/impl/base_impl.py | 2 ++ .../fused_moe/impl/deepgemm_impl.py | 9 ++++++++- .../meta_weights/fused_moe/impl/triton_impl.py | 6 ++++++ .../fused_moe/grouped_fused_moe_ep.py | 17 ++++++++++++++++- .../layer_infer/transformer_layer_infer.py | 14 ++++++++++++++ 6 files changed, 49 insertions(+), 3 deletions(-) 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 72311d22e4..bfaf6b76ba 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 @@ -1,6 +1,6 @@ import torch import threading -from typing import Dict, Any, Optional, Tuple, List +from typing import Dict, Any, Optional, Tuple, List, Callable from lightllm.common.basemodel.layer_weights.meta_weights.base_weight import BaseWeightTpl from lightllm.common.quantization.quantize_method import WeightPack from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_slicer import ( @@ -174,6 +174,7 @@ def experts_with_topk( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, ) -> torch.Tensor: moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_) if moe_capture_callback is not None: @@ -188,6 +189,7 @@ def experts_with_topk( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + shared_expert_func=shared_expert_func, ) def low_latency_dispatch( diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py index e242870d69..aa397c32f9 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/base_impl.py @@ -12,6 +12,8 @@ class FuseMoeBaseImpl: + supports_decode_shared_expert_overlap = False + def __init__( self, n_routed_experts: int, 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 8e4220e4b2..21e37ba969 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 @@ -1,5 +1,5 @@ import torch -from typing import Optional, Tuple, Any +from typing import Optional, Tuple, Any, Callable from .triton_impl import FuseMoeTriton from lightllm.distributed import dist_group_manager from lightllm.common.quantization.quantize_method import WeightPack @@ -14,6 +14,7 @@ masked_group_gemm, chunked_expanded_moe_forward, quantize_fused_experts_input, + use_sm100_mega_moe, ) from lightllm.common.basemodel.triton_kernel.fused_moe.moe_silu_and_mul import silu_and_mul_fwd from lightllm.common.triton_utils.autotuner import Autotuner, AutotuneKernelType @@ -21,6 +22,10 @@ class FuseMoeDeepGEMM(FuseMoeTriton): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.supports_decode_shared_expert_overlap = not use_sm100_mega_moe(self.quant_method) + def _select_experts( self, input_tensor: torch.Tensor, @@ -80,6 +85,7 @@ def _fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, ): output = fused_experts( hidden_states=input_tensor, @@ -94,6 +100,7 @@ def _fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + shared_expert_func=shared_expert_func, ) return output diff --git a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py index 506b711fcf..4c5874c4f8 100644 --- a/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py +++ b/lightllm/common/basemodel/layer_weights/meta_weights/fused_moe/impl/triton_impl.py @@ -101,7 +101,12 @@ def fused_experts_with_topk( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, ): + extra_args = {} + if shared_expert_func is not None: + assert self.supports_decode_shared_expert_overlap + extra_args["shared_expert_func"] = shared_expert_func return self._fused_experts( input_tensor=input_tensor, w13=w13, @@ -112,6 +117,7 @@ def fused_experts_with_topk( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + **extra_args, ) def __call__( 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 61cd81556a..a3ed356889 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 @@ -221,10 +221,13 @@ def fused_experts( alpha: Optional[float] = None, limit: Optional[float] = None, clamp_up_add_one: bool = True, + shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, ): assert (limit is None and alpha is None) or (limit is not None and alpha is not None) + assert shared_expert_func is None or is_prefill is False check_ep_expert_dtype(quant_method) if use_sm100_mega_moe(quant_method): + assert shared_expert_func is None, "Mega MoE does not support shared expert overlap" if limit is not None: raise NotImplementedError("FP4 Mega MoE does not support clamped SwiGLU") return mega_moe_impl(hidden_states, w13, w2, topk_weights, topk_idx, quant_method) @@ -249,6 +252,7 @@ def fused_experts( alpha=alpha, limit=limit, clamp_up_add_one=clamp_up_add_one, + shared_expert_func=shared_expert_func, ) @@ -271,6 +275,7 @@ def fused_experts_impl( limit: Optional[float] = None, clamp_up_add_one: bool = True, block_size_k: int = 128, + shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, ): # Check constraints. assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch" @@ -403,8 +408,18 @@ def fused_experts_impl( ) # low latency combine combined_x, event_overlap, hook = buffer.low_latency_combine( - gemm_out_b, topk_idx, topk_weights, handle, async_finish=False, return_recv_hook=False + gemm_out_b, + topk_idx, + topk_weights, + handle, + async_finish=False, + return_recv_hook=shared_expert_func is not None, ) + if shared_expert_func is not None: + # Launch shared computation while the combine send is in flight. + shared = shared_expert_func() + hook() + return combined_x + shared return combined_x diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py index c5fe0495ee..03c2a9bc23 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py @@ -550,6 +550,20 @@ def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV x = self._tpsp_allgather(input=x, infer_state=infer_state) topk_weights, topk_indices = self._select_experts(x, infer_state, layer_weight) + if not infer_state.is_prefill and layer_weight.experts.fuse_moe_impl.supports_decode_shared_expert_overlap: + out = layer_weight.experts.experts_with_topk( + input_tensor=x, + topk_weights=topk_weights, + topk_ids=topk_indices, + is_prefill=False, + infer_state=infer_state, + alpha=1.0, + limit=self.swiglu_limit, + clamp_up_add_one=False, + shared_expert_func=lambda: self._ffn_tp(x, infer_state, layer_weight), + ) + return out if self.enable_ep_moe else self._tpsp_reduce(input=out, infer_state=infer_state) + # Routed experts can overwrite x in place, so compute the shared expert first. shared = self._ffn_tp(input=x, infer_state=infer_state, layer_weight=layer_weight) routed = layer_weight.experts.experts_with_topk( From 571869590e6075daec208b4e71f1d96633eab94b Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Tue, 3 Nov 2026 15:50:36 +0800 Subject: [PATCH 2/3] refactor(dsv4): expose and validate shared expert overlap controls --- .../fused_moe/fused_moe_weight.py | 14 + .../fused_moe/grouped_fused_moe_ep.py | 3 +- .../layer_infer/transformer_layer_infer.py | 38 +-- .../kernel/benchmark_shared_expert_overlap.py | 300 ++++++++++++++++++ 4 files changed, 336 insertions(+), 19 deletions(-) create mode 100644 test/benchmark/kernel/benchmark_shared_expert_overlap.py 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 bfaf6b76ba..c440dc7ef8 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 @@ -164,6 +164,11 @@ def experts( clamp_up_add_one=clamp_up_add_one, ) + @property + def supports_decode_shared_expert_overlap(self) -> bool: + """Backend capability, independent of the model's overlap enable switch.""" + return self.fuse_moe_impl.supports_decode_shared_expert_overlap + def experts_with_topk( self, input_tensor: torch.Tensor, @@ -176,6 +181,15 @@ def experts_with_topk( clamp_up_add_one: bool = True, shared_expert_func: Optional[Callable[[], torch.Tensor]] = None, ) -> torch.Tensor: + """Run routed experts, optionally including the shared expert in the result. + + The callback requires decode overlap capability. It must treat input_tensor + as read-only and return a matching shared output without modifying routed + state. On successful eager/capture calls it runs exactly once on the current + compute stream, after combine SEND and before RECV. Graph replay executes + the captured GPU work rather than calling Python again. A callback result + is already added to the returned tensor; without one the result is routed-only. + """ moe_capture_callback = get_moe_capture_callback(infer_state, self.layer_num_) if moe_capture_callback is not None: moe_capture_callback(topk_ids) 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 a3ed356889..f3db1352fb 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 @@ -416,7 +416,8 @@ def fused_experts_impl( return_recv_hook=shared_expert_func is not None, ) if shared_expert_func is not None: - # Launch shared computation while the combine send is in flight. + # SEND, shared work, RECV and the final add use the current compute + # stream. The callback is read-only on hidden_states and runs once. shared = shared_expert_func() hook() return combined_x + shared diff --git a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py index 03c2a9bc23..6bef4f34a7 100644 --- a/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/deepseek_v4/layer_infer/transformer_layer_infer.py @@ -1,3 +1,4 @@ +import os import torch import triton from lightllm.common.basemodel import TransformerLayerInferTpl @@ -52,6 +53,11 @@ def __init__(self, layer_num, network_config): self.flashmla_q_head_num_ = self.tp_q_head_num_ self.tp_groups = self.o_groups // self.tp_world_size_ self.enable_ep_moe = get_env_start_args().enable_ep_moe + self.enable_shared_expert_overlap = os.getenv("LIGHTLLM_DSV4_SHARED_EXPERT_OVERLAP", "1").upper() in ( + "1", + "TRUE", + "ON", + ) self.compressor = CompressorInfer( layer_idx=self.layer_num_, network_config=self.network_config_, tp_world_size=self.tp_world_size_ ) @@ -550,23 +556,17 @@ def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV x = self._tpsp_allgather(input=x, infer_state=infer_state) topk_weights, topk_indices = self._select_experts(x, infer_state, layer_weight) - if not infer_state.is_prefill and layer_weight.experts.fuse_moe_impl.supports_decode_shared_expert_overlap: - out = layer_weight.experts.experts_with_topk( - input_tensor=x, - topk_weights=topk_weights, - topk_ids=topk_indices, - is_prefill=False, - infer_state=infer_state, - alpha=1.0, - limit=self.swiglu_limit, - clamp_up_add_one=False, - shared_expert_func=lambda: self._ffn_tp(x, infer_state, layer_weight), - ) - return out if self.enable_ep_moe else self._tpsp_reduce(input=out, infer_state=infer_state) - - # Routed experts can overwrite x in place, so compute the shared expert first. - shared = self._ffn_tp(input=x, infer_state=infer_state, layer_weight=layer_weight) - routed = layer_weight.experts.experts_with_topk( + if ( + self.enable_shared_expert_overlap + and not infer_state.is_prefill + and layer_weight.experts.supports_decode_shared_expert_overlap + ): + shared_expert_func = lambda: self._ffn_tp(x, infer_state, layer_weight) + else: + # Routed experts can overwrite x in place, so compute shared first. + shared = self._ffn_tp(input=x, infer_state=infer_state, layer_weight=layer_weight) + shared_expert_func = None + out = layer_weight.experts.experts_with_topk( input_tensor=x, topk_weights=topk_weights, topk_ids=topk_indices, @@ -575,8 +575,10 @@ def _ffn(self, x, infer_state: DeepseekV4InferStateInfo, layer_weight: DeepseekV alpha=1.0, limit=self.swiglu_limit, clamp_up_add_one=False, + shared_expert_func=shared_expert_func, ) - out = routed + shared + if shared_expert_func is None: + out = out + shared if self.enable_ep_moe: return out return self._tpsp_reduce(input=out, infer_state=infer_state) diff --git a/test/benchmark/kernel/benchmark_shared_expert_overlap.py b/test/benchmark/kernel/benchmark_shared_expert_overlap.py new file mode 100644 index 0000000000..ccc8064ef2 --- /dev/null +++ b/test/benchmark/kernel/benchmark_shared_expert_overlap.py @@ -0,0 +1,300 @@ +"""A/B benchmark of the real DeepEP decode/shared-expert call chain on Hopper. + +Run from the checkout with its pinned DeepEP/DeepGEMM dependencies and RDMA access:: + + PYTHONPATH=. NVSHMEM_QP_DEPTH=1024 python test/benchmark/kernel/benchmark_shared_expert_overlap.py \ + --world-size 8 --output /mnt/vqa/wanzihao/shared-expert-overlap/result.json + +Fixed routes and synthetic FP8 weights isolate a single FFN layer. Defaults use +DSV4 Flash dimensions (4096 hidden, 2048 intermediate, 256 experts over 8 ranks, +top-6). Width 4 tests MTP-expanded row counts, not the full speculative pipeline. +Each mode is checked in eager and CUDA Graph replay, including changed inputs and +routes. Timings exclude routing and report the slowest rank, with alternating A/B +order. This is not a complete model accuracy or end-to-end service benchmark. + +Keep NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE at the service default of 256; +request count times width must fit that capacity. The switch is read when a layer +is created, so each mode owns a separate layer and graph. +""" +import argparse +import json +import os +import socket +import statistics +from datetime import timedelta +from dataclasses import asdict +from pathlib import Path +from types import SimpleNamespace as NS +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + + +def worker(rank, args, port): + torch.cuda.set_device(rank) + from lightllm.utils import dist_utils as du + + du.set_global_rank(rank) + du.set_global_world_size(args.world_size) + du.set_dp_size(args.world_size) + du.set_dp_world_size(1) + du.set_current_rank_in_dp(0) + du.set_current_device_id(rank) + du.set_current_rank_in_node(rank) + du.set_node_world_size(args.world_size) + from lightllm.server.core.objs.start_args_type import StartArgs + + os.environ["LIGHTLLM_START_ARGS"] = json.dumps( + asdict( + StartArgs( + model_dir=args.config_dir, + enable_ep_moe=True, + enable_return_routed_experts=False, + enable_tpsp_mix_mode=False, + batch_max_tokens=8192, + mtp_mode="dspark", + mtp_step=3, + ) + ) + ) + os.environ["NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE"] = str(args.max_rows) + os.environ["LIGHTLLM_USE_UE8M0_SCALES"] = "1" + from lightllm.utils.envs_utils import get_env_start_args + + get_env_start_args.cache_clear() + import deep_ep + from lightllm.distributed import dist_group_manager + import lightllm.common.basemodel + from lightllm.common.quantization.deepgemm import DeepGEMMFP8w8a8B128QuantizationMethod + from lightllm.common.quantization.quantize_method import WeightPack + from lightllm.common.basemodel.triton_kernel.quantization.fp8w8a8_block_quant_kernel import weight_quant + from lightllm.common.basemodel.layer_weights.meta_weights.mm_weight.mm_weight import MMWeightTpl + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight + from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.deepgemm_impl import FuseMoeDeepGEMM + from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer + + dist.init_process_group( + "nccl", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=args.world_size, + timeout=timedelta(seconds=120), + ) + group = dist.new_group(backend="gloo") + dist.barrier(device_ids=[rank]) + experts = args.experts_per_rank * args.world_size + size = deep_ep.Buffer.get_low_latency_rdma_size_hint(args.max_rows, args.hidden, args.world_size, experts) + print(rank, "creating DeepEP buffer", size, flush=True) + buffer = deep_ep.Buffer( + dist.group.WORLD, + num_rdma_bytes=size, + low_latency_mode=True, + num_qps_per_rank=args.experts_per_rank, + explicitly_destroy=True, + ) + dist_group_manager.ep_low_latency_buffer = buffer + print(rank, "DeepEP ready", flush=True) + try: + torch.manual_seed(1013 + rank) + quant = DeepGEMMFP8w8a8B128QuantizationMethod() + quant.use_ue8m0_scales = True + quant.use_packed_ue8m0 = False + + def pack(shape): + value = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.02 + q, s = weight_quant(value, 128, use_ue8m0_scales=True) + return WeightPack(weight=q, weight_scale=s) + + owner = FusedMoeWeight.__new__(FusedMoeWeight) + owner.layer_num_ = 0 + owner.w13 = pack((args.experts_per_rank, 2 * args.intermediate, args.hidden)) + owner.w2 = pack((args.experts_per_rank, args.hidden, args.intermediate)) + owner.fuse_moe_impl = FuseMoeDeepGEMM(experts, 0, 1.0, quant, 0, None, None, False) + assert owner.supports_decode_shared_expert_overlap + + def shared_weight(shape): + w = MMWeightTpl.__new__(MMWeightTpl) + w.quant_method = quant + w.mm_param = pack(shape) + w.bias = None + return w + + weights = NS( + experts=owner, + gate_up_proj=shared_weight((2 * args.intermediate, args.hidden)), + down_proj=shared_weight((args.hidden, args.intermediate)), + ) + config = dict( + rms_norm_eps=1e-6, + hidden_size=args.hidden, + num_attention_heads=32, + head_dim=512, + qk_rope_head_dim=64, + o_groups=8, + hc_mult=4, + hc_sinkhorn_iters=20, + hc_eps=1e-6, + compress_ratios=[0], + num_hash_layers=0, + vocab_size=129280, + n_layer=1, + num_experts_per_tok=args.topk, + routed_scaling_factor=1.0, + swiglu_limit=args.limit, + index_head_dim=128, + index_n_heads=32, + index_topk=512, + ) + layers = {} + for mode in (0, 1): + os.environ["LIGHTLLM_DSV4_SHARED_EXPERT_OVERLAP"] = str(mode) + layer = DeepseekV4TransformerLayerInfer(0, config) + assert layer.enable_shared_expert_overlap == bool(mode) + layers[mode] = layer + original_shared = {mode: layer._ffn_tp for mode, layer in layers.items()} + stats = [] + for requests in args.requests: + for width in args.widths: + rows = requests * width + x = torch.randn(rows, args.hidden, device="cuda", dtype=torch.bfloat16) * 0.1 + ids = torch.rand(rows, experts, device="cuda").topk(args.topk, dim=-1).indices + probs = torch.randn(rows, args.topk, device="cuda").softmax(-1) + infer = NS(is_prefill=False) + counts = {0: 0, 1: 0} + for mode, layer in layers.items(): + layer._select_experts = lambda x, infer, w: (probs, ids) + original = original_shared[mode] + + def shared(*a, mode=mode, original=original, **kw): + assert torch.cuda.current_stream().cuda_stream == stream.cuda_stream + counts[mode] += 1 + result = original(*a, **kw) + return result + + layer._ffn_tp = shared + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + funcs = {mode: (lambda mode=mode: layers[mode]._ffn(x, infer, weights)) for mode in (0, 1)} + print(rank, "case", requests, width, "warming", flush=True) + with torch.cuda.stream(stream): + before = x.clone() + baseline = funcs[0]() + actual = funcs[1]() + stream.synchronize() + torch.testing.assert_close(actual, baseline, rtol=0, atol=0) + torch.testing.assert_close(x, before, rtol=0, atol=0) + assert counts == {0: 1, 1: 1} + for _ in range(args.warmup): + funcs[0]() + funcs[1]() + stream.synchronize() + dist.barrier(group=group) + graphs = {} + outputs = {} + for mode in (0, 1): + graph = torch.cuda.CUDAGraph() + prev = counts[mode] + with torch.cuda.graph(graph, stream=stream): + outputs[mode] = funcs[mode]() + assert counts[mode] == prev + 1 + graphs[mode] = graph + graph.replay() + stream.synchronize() + dist.barrier(group=group) + # A replay must consume changed inputs and routing, without running Python callbacks again. + x.mul_(0.875) + ids.add_(1).remainder_(experts) + counts_before = counts.copy() + graphs[0].replay() + graphs[1].replay() + stream.synchronize() + assert counts == counts_before + torch.testing.assert_close(outputs[1], outputs[0], rtol=0, atol=0) + eager_changed = funcs[0]() + stream.synchronize() + torch.testing.assert_close(outputs[1], eager_changed, rtol=0, atol=0) + assert torch.isfinite(outputs[1]).all() + + def measure(fn): + dist.barrier(group=group) + a, b = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) + a.record(stream) + for _ in range(args.iterations): + fn() + b.record(stream) + b.synchronize() + value = torch.tensor(a.elapsed_time(b) * 1000 / args.iterations, device="cuda") + dist.all_reduce(value, op=dist.ReduceOp.MAX) + return value.item() + + values = {"eager": [], "graph": []} + for kind, call in [("eager", funcs), ("graph", {mode: graphs[mode].replay for mode in (0, 1)})]: + for order in [(0, 1), (1, 0), (0, 1), (1, 0)]: + values[kind].append({str(mode): measure(call[mode]) for mode in order}) + result = { + "requests_per_rank": requests, + "rows_per_request": width, + "token_rows_per_rank": rows, + "exact_output_match": True, + "input_unchanged": True, + "eager_us": { + str(mode): statistics.median(r[str(mode)] for r in values["eager"]) for mode in (0, 1) + }, + "graph_us": { + str(mode): statistics.median(r[str(mode)] for r in values["graph"]) for mode in (0, 1) + }, + "samples": values, + } + stats.append(result) + if rank == 0: + print(json.dumps(result), flush=True) + for graph in graphs.values(): + graph.reset() + torch.cuda.synchronize() + dist.barrier(group=group) + if rank == 0: + output = { + "world_size": args.world_size, + "gpu": torch.cuda.get_device_name(rank), + "torch": torch.__version__, + "deep_ep_module": deep_ep.__file__, + "hidden": args.hidden, + "intermediate": args.intermediate, + "experts": experts, + "topk": args.topk, + "scope": ( + "real single-layer routed/shared FFN; fixed routes, synthetic FP8 weights; " + "no complete model or server" + ), + "cases": stats, + } + Path(args.output).write_text(json.dumps(output, indent=2) + "\n") + finally: + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--world-size", type=int, default=2) + parser.add_argument("--requests", type=int, nargs="+", default=[16, 64]) + parser.add_argument("--widths", type=int, nargs="+", default=[1, 4]) + parser.add_argument("--hidden", type=int, default=4096) + parser.add_argument("--intermediate", type=int, default=2048) + parser.add_argument("--experts-per-rank", type=int, default=32) + parser.add_argument("--topk", type=int, default=6) + parser.add_argument("--limit", type=float, default=10.0) + parser.add_argument("--iterations", type=int, default=10) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--output", required=True) + args = parser.parse_args() + args.max_rows = int(os.getenv("NUM_MAX_DISPATCH_TOKENS_PER_RANK_DECODE", "256")) + assert max(args.requests) * max(args.widths) <= args.max_rows + config_dir = Path(args.output).resolve().parent / "shared-expert-benchmark-config" + config_dir.mkdir(parents=True, exist_ok=True) + (config_dir / "config.json").write_text(json.dumps({"model_type": "deepseek_v4"})) + args.config_dir = str(config_dir) + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + port = s.getsockname()[1] + mp.spawn(worker, args=(args, port), nprocs=args.world_size, join=True) From 1d75d2b363d0d963e80bb322b4657b8bcfa410f5 Mon Sep 17 00:00:00 2001 From: wanzihao <1060304770@qq.com> Date: Tue, 3 Nov 2026 16:52:43 +0800 Subject: [PATCH 3/3] test(dsv4): cover shared expert overlap contracts and switch --- .../deepseek_v4/test_shared_expert_overlap.py | 193 ++++++++++++++++++ 1 file changed, 193 insertions(+) create mode 100644 unit_tests/models/deepseek_v4/test_shared_expert_overlap.py diff --git a/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py b/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py new file mode 100644 index 0000000000..304e5c9b49 --- /dev/null +++ b/unit_tests/models/deepseek_v4/test_shared_expert_overlap.py @@ -0,0 +1,193 @@ +"""Shared-expert callback, dispatch and model fallback contracts.""" + +from types import SimpleNamespace +from unittest.mock import Mock +import pytest +import torch +from lightllm.common.basemodel.triton_kernel.fused_moe import grouped_fused_moe_ep as ep +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl import base_impl, deepgemm_impl +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.impl.triton_impl import FuseMoeTriton +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe.fused_moe_weight import FusedMoeWeight +from lightllm.common.basemodel.layer_weights.meta_weights.fused_moe import fused_moe_weight as weight_module +from lightllm.models.deepseek_v4.layer_infer.transformer_layer_infer import DeepseekV4TransformerLayerInfer + + +@pytest.mark.parametrize("overlap", [False, True]) +def test_real_decode_pipeline_waits_after_shared_computation(monkeypatch, overlap): + events = [] + x = torch.ones(3, 4, dtype=torch.bfloat16) + routed = torch.full_like(x, float("nan")) + + class Buffer: + group_size = 8 + + def low_latency_dispatch(self, *args, **kwargs): + assert kwargs["return_recv_hook"] is False + events.append("dispatch") + return x, None, "handle", None, None + + def low_latency_combine(self, *args, **kwargs): + events.append("combine_send") + assert kwargs["return_recv_hook"] is overlap + + def receive(): + events.append("combine_recv") + routed.fill_(3) + + if not overlap: + receive() + return routed, None, receive + + def gemm(*args, **kwargs): + assert args[7] == 3 # ceil(3 requests * 8 ranks * topk2 / 16 experts) + events.append("gemm") + return x + + def shared(): + assert torch.isnan(routed).all() + assert torch.all(x == 1) + events.append("shared") + return x * 2 + + monkeypatch.setattr(ep, "masked_group_gemm", gemm) + monkeypatch.setattr(ep, "get_env_use_ue8m0_scales", lambda: False) + monkeypatch.setattr(ep, "get_deepep_num_max_dispatch_tokens_per_rank_decode", lambda: 64) + out = ep.fused_experts_impl( + x, + torch.ones(2, 8, 4), + torch.ones(2, 4, 4), + torch.ones(3, 2), + torch.zeros(3, 2, dtype=torch.int64), + 16, + Buffer(), + False, + use_fp8_w8a8=True, + use_fp8_all2all=True, + shared_expert_func=shared if overlap else None, + ) + expected = ["dispatch", "gemm", "combine_send"] + (["shared", "combine_recv"] if overlap else ["combine_recv"]) + assert events == expected + assert torch.all(out == (5 if overlap else 3)) + + +@pytest.mark.parametrize("mega", [False, True]) +def test_capability_is_selected_at_backend_construction(monkeypatch, mega): + monkeypatch.setattr(base_impl, "get_global_rank", lambda: 0) + monkeypatch.setattr(base_impl, "get_global_world_size", lambda: 8) + monkeypatch.setattr(deepgemm_impl, "use_sm100_mega_moe", lambda _: mega) + backend = deepgemm_impl.FuseMoeDeepGEMM(16, 0, 1.0, SimpleNamespace(), 0, None, None, False) + assert backend.supports_decode_shared_expert_overlap is (not mega) + assert base_impl.FuseMoeBaseImpl.supports_decode_shared_expert_overlap is False + + +def test_weight_preserves_route_capture_and_callback(monkeypatch): + owner = FusedMoeWeight.__new__(FusedMoeWeight) + owner.layer_num_ = 0 + owner.w13 = owner.w2 = None + backend = deepgemm_impl.FuseMoeDeepGEMM.__new__(deepgemm_impl.FuseMoeDeepGEMM) + backend.supports_decode_shared_expert_overlap = True + backend._fused_experts = Mock(return_value="output") + owner.fuse_moe_impl = backend + capture = Mock() + monkeypatch.setattr(weight_module, "get_moe_capture_callback", lambda *_: capture) + ids = torch.zeros(3, 2, dtype=torch.int64) + shared = Mock() + assert ( + owner.experts_with_topk(torch.ones(3, 4), torch.ones(3, 2), ids, is_prefill=False, shared_expert_func=shared) + == "output" + ) + capture.assert_called_once_with(ids) + assert backend._fused_experts.call_args.kwargs["shared_expert_func"] is shared + + +def test_existing_inplace_backends_receive_no_new_keyword(): + backend = FuseMoeTriton.__new__(FuseMoeTriton) + backend._fused_experts = Mock(return_value="old-output") + assert backend.fused_experts_with_topk(None, None, None, None, None) == "old-output" + assert "shared_expert_func" not in backend._fused_experts.call_args.kwargs + with pytest.raises(AssertionError): + backend.fused_experts_with_topk(None, None, None, None, None, shared_expert_func=lambda: None) + + +@pytest.mark.parametrize("enabled", [False, True]) +@pytest.mark.parametrize("ep_mode", [False, True]) +@pytest.mark.parametrize("prefill,supports", [(False, True), (False, False), (True, True)]) +def test_model_keeps_shared_before_inplace_fallback(prefill, supports, enabled, ep_mode): + layer = DeepseekV4TransformerLayerInfer.__new__(DeepseekV4TransformerLayerInfer) + layer.embed_dim_ = 4 + layer.enable_ep_moe = ep_mode + layer.swiglu_limit = 7 + layer.enable_shared_expert_overlap = enabled + layer._tpsp_allgather = lambda input, infer_state: input + ids = torch.zeros(3, 2, dtype=torch.int64) + layer._select_experts = lambda *args: (torch.ones(3, 2), ids) + events = [] + + def shared(x=None, *args, **kwargs): + if x is None: + x = kwargs["input"] + events.append("shared") + return x * 2 + + layer._ffn_tp = shared + + def routed(**kwargs): + events.append("routed") + if kwargs.get("shared_expert_func") is not None: + return kwargs["shared_expert_func"]() + 3 + kwargs["input_tensor"].fill_(9) + return torch.full((3, 4), 3.0) + + owner = FusedMoeWeight.__new__(FusedMoeWeight) + owner.fuse_moe_impl = SimpleNamespace(supports_decode_shared_expert_overlap=supports) + owner.experts_with_topk = routed + weights = SimpleNamespace(experts=owner) + reductions = [] + layer._tpsp_reduce = lambda input, infer_state: (reductions.append(input) or input) + result = layer._ffn(torch.ones(3, 4), SimpleNamespace(is_prefill=prefill), weights) + assert torch.all(result == 5) + assert events == (["routed", "shared"] if enabled and supports and not prefill else ["shared", "routed"]) + assert len(reductions) == (0 if ep_mode else 1) + + +@pytest.mark.parametrize("supported", [False, True]) +def test_weight_reports_backend_capability(supported): + owner = FusedMoeWeight.__new__(FusedMoeWeight) + owner.fuse_moe_impl = SimpleNamespace(supports_decode_shared_expert_overlap=supported) + assert owner.supports_decode_shared_expert_overlap is supported + + +@pytest.mark.parametrize("value,enabled", [(None, True), ("0", False), ("1", True), ("TRUE", True)]) +def test_overlap_switch_is_read_when_layer_is_created(monkeypatch, value, enabled): + from lightllm.models.deepseek_v4.layer_infer import transformer_layer_infer + + monkeypatch.setenv("LIGHTLLM_CURRENT_RANK_IN_DP", "0") + monkeypatch.setenv("LIGHTLLM_DP_WORLD_SIZE", "1") + if value is None: + monkeypatch.delenv("LIGHTLLM_DSV4_SHARED_EXPERT_OVERLAP", raising=False) + else: + monkeypatch.setenv("LIGHTLLM_DSV4_SHARED_EXPERT_OVERLAP", value) + monkeypatch.setattr(transformer_layer_infer, "get_env_start_args", lambda: SimpleNamespace(enable_ep_moe=True)) + config = dict( + rms_norm_eps=1e-6, + hidden_size=4096, + num_attention_heads=32, + head_dim=512, + qk_rope_head_dim=64, + o_groups=8, + hc_mult=4, + hc_sinkhorn_iters=20, + hc_eps=1e-6, + compress_ratios=[0], + num_hash_layers=0, + vocab_size=129280, + n_layer=1, + num_experts_per_tok=6, + routed_scaling_factor=1.0, + swiglu_limit=10.0, + index_head_dim=128, + index_n_heads=32, + index_topk=512, + ) + layer = DeepseekV4TransformerLayerInfer(0, config) + assert layer.enable_shared_expert_overlap is enabled