From 0fe282543fbb36eaa3066bdf0555bc882d3ce0bc Mon Sep 17 00:00:00 2001
From: jiashaokun-1 <1114621279@qq.com>
Date: Sun, 1 Feb 2026 18:31:39 +0800
Subject: [PATCH] [NPU] support the Enable return routed experts (#17025)

---
 python/sglang/srt/hardware_backend/npu/moe/topk.py | 7 +++++++
 python/sglang/srt/layers/moe/topk.py               | 1 +
 2 files changed, 8 insertions(+)

diff --git a/python/sglang/srt/hardware_backend/npu/moe/topk.py b/python/sglang/srt/hardware_backend/npu/moe/topk.py
index 813c12f6a..920787ed2 100644
--- a/python/sglang/srt/hardware_backend/npu/moe/topk.py
+++ b/python/sglang/srt/hardware_backend/npu/moe/topk.py
@@ -5,6 +5,7 @@ from sgl_kernel_npu.norm.l1_norm import l1_norm
 
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.eplb.expert_location_dispatch import topk_ids_logical_to_physical
+from sglang.srt.layers.moe.routed_experts_capturer import get_global_experts_capturer
 from sglang.srt.layers.moe.topk import StandardTopKOutput, select_experts
 
 if TYPE_CHECKING:
@@ -18,6 +19,7 @@ def fused_topk_npu(
     topk_config: "TopKConfig",
     num_token_non_padded: Optional[torch.Tensor] = None,
     expert_location_dispatch_info: Optional["ExpertLocationDispatchInfo"] = None,
+    layer_id: Optional[int] = None,
 ) -> "TopKOutput":
 
     use_grouped_topk = topk_config.use_grouped_topk
@@ -59,6 +61,7 @@ def fused_topk_npu(
         topk_config.torch_native = True
         return select_experts(
             hidden_states=hidden_states,
+            layer_id=layer_id,
             router_logits=router_logits,
             topk_config=topk_config,
             num_token_non_padded=num_token_non_padded,
@@ -68,5 +71,9 @@ def fused_topk_npu(
     if expert_location_dispatch_info is not None:
         topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
     get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
+    get_global_experts_capturer().capture(
+        layer_id=layer_id,
+        topk_ids=topk_ids,
+    )
 
     return StandardTopKOutput(topk_weights, topk_ids, router_logits)
diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py
index 419786c2f..c1ee9e461 100644
--- a/python/sglang/srt/layers/moe/topk.py
+++ b/python/sglang/srt/layers/moe/topk.py
@@ -351,6 +351,7 @@ class TopK(MultiPlatformOp):
             topk_config=self.topk_config,
             num_token_non_padded=num_token_non_padded,
             expert_location_dispatch_info=expert_location_dispatch_info,
+            layer_id=self.layer_id,
         )
 
     def empty_topk_output(self, device: torch.device) -> TopKOutput: