[Feature] Enable return routed experts

python/sglang/srt/layers/moe/routed_experts_capturer.py

@@ -0,0 +1,180 @@
+import logging
+from abc import ABC
+from typing import Optional
+
+import numpy as np
+import torch
+
+from sglang.srt.configs.model_config import ModelConfig
+from sglang.srt.server_args import get_global_server_args
+
+# from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+# from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
+
+logger = logging.getLogger(__name__)
+
+GB = 1024 * 1024 * 1024
+MB = 1024 * 1024
+
+
+def get_tensor_size_bytes(t: torch.Tensor):
+    return np.prod(t.shape) * t.dtype.itemsize
+
+
+class RoutedExpertsDeviceCache:
+    def __init__(
+        self, model_config: ModelConfig, max_running_requests: int, device: str
+    ) -> None:
+        self.buffer = torch.zeros(
+            (
+                max(
+                    get_global_server_args().chunked_prefill_size, max_running_requests
+                ),
+                model_config.hf_text_config.num_hidden_layers,
+                model_config.hf_text_config.num_experts_per_tok,
+            ),
+            dtype=torch.int32,
+            device=device,
+        )
+        self._finalize_allocation_log()
+
+    def get_buffer_size_bytes(self):
+        assert hasattr(self, "buffer")
+        return get_tensor_size_bytes(self.buffer)
+
+    def capture_fwd_routed_experts(self, layer_id: int, topk_ids: torch.Tensor):
+        assert layer_id is not None, "capturing routing experts but get layer_id None"
+        batch, _ = topk_ids.shape
+        self.buffer[:batch, layer_id, :] = topk_ids
+
+    def _finalize_allocation_log(self):
+        """Common logging and memory usage computation for captured experts buffers."""
+        buffer_size_MB = self.get_buffer_size_bytes() / MB
+        logger.info(
+            f"Routing experts device buffer allocated. #shape: {tuple(self.buffer.shape)}, size: {buffer_size_MB:.2f} MB"
+        )
+
+
+class RoutedExpertsHostCache:
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        num_tokens: int,
+    ) -> None:
+        self.num_tokens = num_tokens
+        self.buffer = torch.zeros(
+            (
+                num_tokens,
+                model_config.hf_text_config.num_hidden_layers,
+                model_config.hf_text_config.num_experts_per_tok,
+            ),
+            dtype=torch.int32,
+            device="cpu",
+        )
+        self._finalize_allocation_log()
+
+    def get_buffer_size_bytes(self):
+        assert hasattr(self, "buffer")
+        return get_tensor_size_bytes(self.buffer)
+
+    def set_experts_buffer(self, layer_id: int, loc: torch.Tensor, top_k: torch.Tensor):
+        self.buffer[layer_id, loc, :] = top_k.cpu()
+
+    def _finalize_allocation_log(self):
+        """Common logging and memory usage computation for captured experts buffers."""
+        buffer_size_GB = self.get_buffer_size_bytes() / GB
+        logger.info(
+            f"Routing experts host buffer allocated. #tokens: {self.num_tokens}, size: {buffer_size_GB:.2f} GB"
+        )
+
+
+class RoutedExpertsCapturer(ABC):
+    @staticmethod
+    def create(
+        enable: bool,
+        model_config: ModelConfig,
+        num_tokens: int,
+        max_running_requests: int,
+        device: str,
+    ):
+        if enable:
+            return _RoutedExpertsCapturerReal(
+                model_config,
+                num_tokens=num_tokens,
+                max_running_requests=max_running_requests,
+                device=device,
+            )
+        else:
+            return _RoutedExpertsCapturerNoop()
+
+    def capture(self, layer_id: int, topk_ids: torch.Tensor):
+        raise NotImplementedError
+
+    def sync_fwd_experts_buffer_DtoH(self, batch: int, loc: torch.Tensor):
+        raise NotImplementedError
+
+    def get_host_cache(self):
+        raise NotImplementedError
+
+    def get_device_cache(self):
+        raise NotImplementedError
+
+
+class _RoutedExpertsCapturerReal(RoutedExpertsCapturer):
+    """Capturer for routed experts with host buffer"""
+
+    def __init__(
+        self,
+        model_config: ModelConfig,
+        num_tokens: int,
+        max_running_requests: int,
+        device: str,
+    ):
+
+        self.host_cache = RoutedExpertsHostCache(model_config, num_tokens)
+
+        self.device_cache = RoutedExpertsDeviceCache(
+            model_config, max_running_requests, device
+        )
+
+    def capture(self, layer_id: int, topk_ids: torch.Tensor):
+        self.device_cache.capture_fwd_routed_experts(layer_id, topk_ids)
+
+    def sync_fwd_experts_buffer_DtoH(self, loc: torch.Tensor):
+        batch = loc.shape[0]
+        self.host_cache.buffer[loc] = self.device_cache.buffer[:batch].cpu()
+
+    def get_host_cache(self):
+        return self.host_cache
+
+    def get_device_cache(self):
+        return self.device_cache
+
+
+class _RoutedExpertsCapturerNoop(RoutedExpertsCapturer):
+    def __init__(self):
+        pass
+
+    def capture(self, layer_id: int, topk_ids: torch.Tensor):
+        pass
+
+    def sync_fwd_experts_buffer_DtoH(self, loc: torch.Tensor):
+        pass
+
+    def get_host_cache(self):
+        pass
+
+    def get_device_cache(self):
+        pass
+
+
+_global_expert_capturer: Optional[RoutedExpertsCapturer] = _RoutedExpertsCapturerNoop()
+
+
+def get_global_experts_capturer():
+    return _global_expert_capturer
+
+
+def set_global_experts_capturer(capturer: RoutedExpertsCapturer):
+    global _global_expert_capturer
+    _global_expert_capturer = capturer

python/sglang/srt/layers/moe/topk.py

@@ -42,6 +42,7 @@
     get_moe_runner_backend,
     should_use_flashinfer_trtllm_moe,
 )
+from sglang.srt.layers.moe.routed_experts_capturer import get_global_experts_capturer
 from sglang.srt.utils import (
     cpu_has_amx_support,
     get_bool_env_var,
@@ -193,6 +194,7 @@ def __init__(
         self,
         top_k: int,
         *,
+        layer_id: Optional[int] = None,
         use_grouped_topk: bool = False,
         topk_group: Optional[int] = None,
         num_expert_group: Optional[int] = None,
@@ -213,6 +215,7 @@ def __init__(
         if use_grouped_topk:
             assert num_expert_group is not None and topk_group is not None
 
+        self.layer_id = layer_id
         self.topk_config = TopKConfig(
             top_k=top_k,
             use_grouped_topk=use_grouped_topk,
@@ -238,6 +241,7 @@ def forward_native(
         self.topk_config.torch_native = True
         return select_experts(
             hidden_states=hidden_states,
+            layer_id=self.layer_id,
             router_logits=router_logits,
             topk_config=self.topk_config,
             num_token_non_padded=num_token_non_padded,
@@ -284,6 +288,7 @@ def forward_cuda(
             self.topk_config.torch_native = False
             return select_experts(
                 hidden_states=hidden_states,
+                layer_id=self.layer_id,
                 router_logits=router_logits,
                 topk_config=self.topk_config,
                 num_token_non_padded=num_token_non_padded,
@@ -300,6 +305,7 @@ def forward_cpu(
     ) -> TopKOutput:
         return select_experts(
             hidden_states=hidden_states,
+            layer_id=self.layer_id,
             router_logits=router_logits,
             topk_config=self.topk_config,
             num_token_non_padded=num_token_non_padded,
@@ -356,6 +362,7 @@ def forward_npu(
             self.topk_config.torch_native = True
             return select_experts(
                 hidden_states=hidden_states,
+                layer_id=self.layer_id,
                 router_logits=router_logits,
                 topk_config=self.topk_config,
                 num_token_non_padded=num_token_non_padded,
@@ -789,6 +796,7 @@ def select_experts(
     router_logits: torch.Tensor,
     topk_config: TopKConfig,
     *,
+    layer_id: Optional[int] = None,
     num_token_non_padded: Optional[torch.Tensor] = None,
     expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
 ) -> StandardTopKOutput:
@@ -886,5 +894,8 @@ def select_experts(
         )
 
     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)

python/sglang/srt/managers/detokenizer_manager.py

@@ -271,6 +271,7 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
             output_token_ids_logprobs_idx=recv_obj.output_token_ids_logprobs_idx,
             output_token_entropy_val=recv_obj.output_token_entropy_val,
             output_hidden_states=recv_obj.output_hidden_states,
+            output_routed_experts=recv_obj.output_routed_experts,
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,
             token_steps=recv_obj.token_steps,

python/sglang/srt/managers/io_struct.py

@@ -126,6 +126,8 @@ class GenerateReqInput(BaseReq):
     log_metrics: bool = True
     # Whether to return hidden states
     return_hidden_states: Union[List[bool], bool] = False
+    # Whether to return captured routed experts
+    return_routed_experts: bool = False
 
     # The modalities of the image data [image, multi-images, video]
     modalities: Optional[List[str]] = None
@@ -540,6 +542,7 @@ def __getitem__(self, i):
                 if isinstance(self.return_hidden_states, list)
                 else self.return_hidden_states
             ),
+            return_routed_experts=self.return_routed_experts,
             modalities=self.modalities[i] if self.modalities else None,
             session_params=self.session_params,
             lora_path=self.lora_path[i] if self.lora_path is not None else None,
@@ -602,6 +605,9 @@ class TokenizedGenerateReqInput(BaseReq):
     # Whether to return hidden states
     return_hidden_states: bool = False
 
+    # Whether to return captured routed experts
+    return_routed_experts: bool = False
+
     # The input embeds
     input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None
 
@@ -848,6 +854,9 @@ class BatchTokenIDOutput(BaseBatchReq):
     # Hidden states
     output_hidden_states: List[List[float]]
 
+    # The routed experts for each output token
+    output_routed_experts: List[List[int]]
+
     # The information of placeholder tokens (e.g., image token)
     # idx is the index of the token in the prompt after expansion.
     # val is the length of padded tokens after expansion.
@@ -924,6 +933,9 @@ class BatchStrOutput(BaseBatchReq):
     # Hidden states
     output_hidden_states: List[List[float]]
 
+    # The routed experts for each output token
+    output_routed_experts: List[List[int]]
+
     # The information of placeholder tokens (e.g., image token)
     # idx is the index of the token in the prompt after expansion.
     # val is the length of padded tokens after expansion.

python/sglang/srt/managers/schedule_batch.py

@@ -428,6 +428,7 @@ def __init__(
         session_id: Optional[str] = None,
         custom_logit_processor: Optional[str] = None,
         return_hidden_states: bool = False,
+        return_routed_experts: bool = False,
         eos_token_ids: Optional[Set[int]] = None,
         bootstrap_host: Optional[str] = None,
         bootstrap_port: Optional[int] = None,
@@ -599,6 +600,10 @@ def __init__(
         self.output_topk_p = None
         self.output_topk_index = None
 
+        # capture routed experts
+        self.return_routed_experts = return_routed_experts
+        self.routed_experts = []
+
         # Embedding (return values)
         self.embedding = None
 
@@ -876,6 +881,7 @@ def check_finished(self, new_accepted_len: int = 1):
 
     def reset_for_retract(self):
         self.prefix_indices = torch.empty((0,), dtype=torch.int64)
+        self.routed_experts = []
         self.last_node = None
         self.swa_uuid_for_lock = None
         self.extend_input_len = 0
@@ -1027,6 +1033,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
     # Whether to return hidden states
     return_hidden_states: bool = False
 
+    # Whether to return captured experts
+    return_routed_experts: bool = False
+
     # Whether this batch is prefill-only (no token generation needed)
     is_prefill_only: bool = False
 
@@ -1070,6 +1079,7 @@ def init_new(
             device=req_to_token_pool.device,
             spec_algorithm=spec_algorithm,
             return_hidden_states=any(req.return_hidden_states for req in reqs),
+            return_routed_experts=any(req.return_routed_experts for req in reqs),
             is_prefill_only=all(req.is_prefill_only for req in reqs),
             chunked_req=chunked_req,
         )
@@ -1789,7 +1799,8 @@ def _is_available_size_sufficient(self, num_tokens: int) -> bool:
     def __str__(self):
         return (
             f"ScheduleBatch(forward_mode={self.forward_mode.name if self.forward_mode else 'None'}, "
-            f"#req={(len(self.reqs))})"
+            f"#req={(len(self.reqs))}), "
+            f"#out_cache_loc={self.out_cache_loc})"
         )