From 7f45bcdd2ab80e798b9eb0d335736d5e4f4f3b20 Mon Sep 17 00:00:00 2001
From: Siyuan Chen <41201609+SYChen123@users.noreply.github.com>
Date: Mon, 25 May 2026 01:09:41 +0800
Subject: [PATCH] [dsv4] support eplb (#25948)

Co-authored-by: xutizhou <xutingz@nvidia.com>
---
 python/sglang/srt/layers/moe/hash_topk.py |  4 ++++
 python/sglang/srt/models/deepseek_v4.py   | 20 ++++++++++++++------
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/python/sglang/srt/layers/moe/hash_topk.py b/python/sglang/srt/layers/moe/hash_topk.py
index 959880f84..0902403e6 100644
--- a/python/sglang/srt/layers/moe/hash_topk.py
+++ b/python/sglang/srt/layers/moe/hash_topk.py
@@ -7,6 +7,9 @@ import torch
 from torch import nn
 
 from sglang.srt.environ import envs
+from sglang.srt.eplb.expert_distribution import (
+    get_global_expert_distribution_recorder,
+)
 from sglang.srt.eplb.expert_location_dispatch import (
     ExpertLocationDispatchInfo,
     topk_ids_logical_to_physical,
@@ -145,6 +148,7 @@ class HashTopK(nn.Module):
 
         topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
         _mask_topk_ids_padded_region(topk_ids, num_token_non_padded)
+        get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
         topk_output = StandardTopKOutput(
             topk_weights=topk_weights, topk_ids=topk_ids, router_logits=router_logits
         )
diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py
index 86864dcc9..6e8af891c 100644
--- a/python/sglang/srt/models/deepseek_v4.py
+++ b/python/sglang/srt/models/deepseek_v4.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 import concurrent.futures
 import logging
 import time
+from contextlib import nullcontext
 from typing import (
     TYPE_CHECKING,
     Iterable,
@@ -33,6 +34,7 @@ from sglang.srt.distributed import (
     get_tp_group,
 )
 from sglang.srt.environ import envs
+from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
 from sglang.srt.layers.attention.dsa.utils import (
     can_dsa_cp_split,
@@ -1269,13 +1271,19 @@ class DeepseekV4Model(nn.Module):
 
         for i in range(self.start_layer, self.end_layer):
             layer = self.layers[i]
-            hidden_states = layer(
-                positions=positions,
-                hidden_states=hidden_states,
-                forward_batch=forward_batch,
-                input_ids=input_ids,
-                input_ids_global=input_ids_global,
+            ctx = (
+                nullcontext()
+                if not get_global_server_args().disable_piecewise_cuda_graph
+                else get_global_expert_distribution_recorder().with_current_layer(i)
             )
+            with ctx:
+                hidden_states = layer(
+                    positions=positions,
+                    hidden_states=hidden_states,
+                    forward_batch=forward_batch,
+                    input_ids=input_ids,
+                    input_ids_global=input_ids_global,
+                )
 
         # CP all-gather only on the last PP rank; PP IPC carries CP-split tensors.
         if self.pp_group.is_last_rank and dsa_use_prefill_cp(forward_batch):