fix return_routed_experts

Author: sufubao

lightllm/common/basemodel/routing_manager.py

@@ -14,7 +14,7 @@
 
 def routing_dtype_id_to_np(dtype_id: int):
     if dtype_id == 1:
-        return np.int8
+        return np.uint8
     elif dtype_id == 2:
         return np.int16
     return np.int32
@@ -39,8 +39,8 @@ def __init__(
         self.num_experts = num_experts
         self.kv_cache_size = kv_cache_size
 
-        self.dtype = torch.int8 if num_experts <= 127 else torch.int16
-        dtype_bytes = 1 if self.dtype == torch.int8 else 2
+        self.dtype = torch.uint8 if num_experts <= 255 else torch.int16
+        dtype_bytes = 1 if self.dtype == torch.uint8 else 2
 
         # Shape: (num_moe_layers, kv_cache_size, topk) — on CPU to save GPU memory.
         # Written after forward() via flush_to_routing_buffer(), read on request finish.
@@ -57,7 +57,7 @@ def __init__(
             torch.zeros((max_capture_tokens, num_moe_layers, topk), dtype=self.dtype, device="cuda") for _ in range(2)
         ]
 
-        dtype_name = "int8" if self.dtype == torch.int8 else "int16"
+        dtype_name = "uint8" if self.dtype == torch.uint8 else "int16"
         logger.info(
             f"RoutingCaptureManager initialized: {num_moe_layers} MoE layers, topk={topk}, "
             f"routing_buffer(cpu)={routing_buffer_size / 1024 / 1024:.2f}MB, "
@@ -66,11 +66,11 @@ def __init__(
 
     @property
     def np_dtype(self):
-        return np.int8 if self.dtype == torch.int8 else np.int16
+        return np.uint8 if self.dtype == torch.uint8 else np.int16
 
     @property
     def dtype_id(self) -> int:
-        return 1 if self.dtype == torch.int8 else 2
+        return 1 if self.dtype == torch.uint8 else 2
 
     def capture(self, moe_layer_index: int, topk_ids: torch.Tensor, microbatch_index: int = 0) -> None:
         num_tokens = topk_ids.shape[0]

lightllm/server/api_lightllm.py

@@ -35,6 +35,9 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana
     prompt = request_dict.pop("inputs")
     sample_params_dict = request_dict["parameters"]
     return_details = sample_params_dict.pop("return_details", False)
+    return_routed_experts = sample_params_dict.pop(
+        "return_routed_experts", httpserver_manager.args.enable_return_routed_experts
+    )
     sampling_params = SamplingParams()
     sampling_params.init(tokenizer=httpserver_manager.tokenizer, **sample_params_dict)
     sampling_params.verify()
@@ -105,7 +108,7 @@ async def lightllm_generate(request: Request, httpserver_manager: HttpServerMana
         ret["prompt_logprobs"] = prompt_logprobs
     if input_usage is not None:
         ret["input_usage"] = input_usage
-    if routed_experts_data is not None:
+    if return_routed_experts and routed_experts_data is not None:
         ret["routed_experts"] = routed_experts_data
 
     return Response(content=json.dumps(ret, ensure_ascii=False).encode("utf-8"))
@@ -117,6 +120,7 @@ async def lightllm_generate_stream(request: Request, httpserver_manager: HttpSer
     prompt = request_dict.pop("inputs")
     sample_params_dict = request_dict["parameters"]
     _ = sample_params_dict.pop("return_details", False)
+    _ = sample_params_dict.pop("return_routed_experts", None)
     sampling_params = SamplingParams()
     sampling_params.init(tokenizer=httpserver_manager.tokenizer, **sample_params_dict)
     sampling_params.verify()

lightllm/server/core/objs/sampling_params.py

@@ -357,23 +357,28 @@ class SamplingParams(ctypes.Structure):
 
     def init(self, tokenizer, **kwargs):
         super().__init__()
-        self.best_of = kwargs.get("best_of", 1)
-        self.n = kwargs.get("n", self.best_of)
-        self.do_sample = kwargs.get("do_sample", SamplingParams._do_sample)
-        self.presence_penalty = kwargs.get("presence_penalty", SamplingParams._presence_penalty)
-        self.frequency_penalty = kwargs.get("frequency_penalty", SamplingParams._frequency_penalty)
-        self.repetition_penalty = kwargs.get("repetition_penalty", SamplingParams._repetition_penalty)
-        self.temperature = kwargs.get("temperature", SamplingParams._temperature)
-        self.top_p = kwargs.get("top_p", SamplingParams._top_p)
-        self.top_k = kwargs.get("top_k", SamplingParams._top_k)
-        self.ignore_eos = kwargs.get("ignore_eos", False)
-        self.min_pixels = kwargs.get("min_pixels", -1)
-        self.max_pixels = kwargs.get("max_pixels", -1)
-        self.max_new_tokens = kwargs.get("max_new_tokens", 16)
-        self.min_new_tokens = kwargs.get("min_new_tokens", 1)
-        self.input_penalty = kwargs.get("input_penalty", DEFAULT_INPUT_PENALTY)
-        self.group_request_id = kwargs.get("group_request_id", -1)
-        self.suggested_dp_index = kwargs.get("suggested_dp_index", -1)
+
+        def _get(key, default):
+            v = kwargs.get(key)
+            return v if v is not None else default
+
+        self.best_of = _get("best_of", 1)
+        self.n = _get("n", self.best_of)
+        self.do_sample = _get("do_sample", SamplingParams._do_sample)
+        self.presence_penalty = _get("presence_penalty", SamplingParams._presence_penalty)
+        self.frequency_penalty = _get("frequency_penalty", SamplingParams._frequency_penalty)
+        self.repetition_penalty = _get("repetition_penalty", SamplingParams._repetition_penalty)
+        self.temperature = _get("temperature", SamplingParams._temperature)
+        self.top_p = _get("top_p", SamplingParams._top_p)
+        self.top_k = _get("top_k", SamplingParams._top_k)
+        self.ignore_eos = _get("ignore_eos", False)
+        self.min_pixels = _get("min_pixels", -1)
+        self.max_pixels = _get("max_pixels", -1)
+        self.max_new_tokens = _get("max_new_tokens", 16)
+        self.min_new_tokens = _get("min_new_tokens", 1)
+        self.input_penalty = _get("input_penalty", DEFAULT_INPUT_PENALTY)
+        self.group_request_id = _get("group_request_id", -1)
+        self.suggested_dp_index = _get("suggested_dp_index", -1)
 
         self.skip_special_tokens = kwargs.get("skip_special_tokens", SKIP_SPECIAL_TOKENS)
         self.disable_prompt_cache = kwargs.get("disable_prompt_cache", False)