fix(openai): avoid duplicate routed expert in response when return_meta_info = True (#35323)
This commit is contained in:
@@ -342,7 +342,7 @@ for chunk in stream:
|
||||
|
||||
#### Returning Routed Experts (MoE Models)
|
||||
|
||||
For MoE models, set `return_routed_experts: true` in `extra_body` to return expert routing data. Requires `--enable-return-routed-experts` server flag. The `routed_experts` field will be returned in the `sgl_ext` object on each choice, containing base64-encoded int32 expert IDs as a flattened array with logical shape `[num_tokens, num_layers, top_k]`. By default this returns `[0, seqlen - 1)`, the full available sequence, because RL workflows need routed experts for the full sequence. Set `routed_experts_start_len` in `extra_body` to an absolute prefix length to return only `[routed_experts_start_len, seqlen - 1)`. For example, in multi-turn RL rollouts, routed experts for tokens from previous turns have already been collected, so setting this value avoids unnecessary transfer that cause bottlenecks.
|
||||
For MoE models, set `return_routed_experts: true` in `extra_body` to return expert routing data. Requires `--enable-return-routed-experts` server flag. By default, `routed_experts` is returned in the response-level `sglext` object as base64-encoded int32 expert IDs in a flattened array with logical shape `[num_tokens, num_layers, top_k]`. For non-streaming chat responses with `return_meta_info: true`, `routed_experts` is returned only in each choice's `meta_info` object and omitted from response-level `sglext` to avoid duplicating the routing data. Other requested extension fields, such as `cached_tokens_details` and `spec_tokens_details`, remain in response-level `sglext`. By default, routed experts cover `[0, seqlen - 1)`, the full available sequence, because RL workflows need them for the full sequence. Set `routed_experts_start_len` in `extra_body` to an absolute prefix length to return only `[routed_experts_start_len, seqlen - 1)`. For example, in multi-turn RL rollouts, routed experts for tokens from previous turns have already been collected, so setting this value avoids unnecessary transfer that causes bottlenecks.
|
||||
|
||||
```python Example
|
||||
# Example with logit_bias parameter for completions API
|
||||
|
||||
@@ -1819,7 +1819,11 @@ class OpenAIServingChat(OpenAIServingBase):
|
||||
|
||||
# Build sglext at response level (from first ret_item, as these are per-request)
|
||||
first_ret = ret[0]
|
||||
routed_experts = process_routed_experts_from_ret(first_ret, request)
|
||||
routed_experts = (
|
||||
None
|
||||
if request.return_meta_info
|
||||
else process_routed_experts_from_ret(first_ret, request)
|
||||
)
|
||||
cached_tokens_details = process_cached_tokens_details_from_ret(
|
||||
first_ret, request
|
||||
)
|
||||
|
||||
@@ -2477,13 +2477,13 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
"empty-delta logprobs chunk emitted without a parser; would break client chunk-shape assumptions",
|
||||
)
|
||||
|
||||
def test_non_streaming_cached_tokens_details_emits_sglext(self):
|
||||
"""Test that non-streaming chat responses emit cached token details in sglext."""
|
||||
|
||||
def test_non_streaming_extension_fields_emit_sglext_without_meta_info(self):
|
||||
req = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "Hi?"}],
|
||||
max_tokens=100,
|
||||
return_meta_info=False,
|
||||
return_routed_experts=True,
|
||||
return_cached_tokens_details=True,
|
||||
)
|
||||
ret = [
|
||||
@@ -2500,6 +2500,7 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
"storage": 1,
|
||||
"storage_backend": "file",
|
||||
},
|
||||
"routed_experts": "cm91dGUtYQ==",
|
||||
"finish_reason": {"type": "stop", "matched": None},
|
||||
"weight_version": "default",
|
||||
},
|
||||
@@ -2509,6 +2510,7 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
response = self.chat._build_chat_response(req, ret, 1234567890)
|
||||
|
||||
self.assertIsNotNone(response.sglext)
|
||||
self.assertEqual(response.sglext.routed_experts, "cm91dGUtYQ==")
|
||||
self.assertEqual(
|
||||
response.sglext.cached_tokens_details.model_dump(exclude_none=True),
|
||||
{
|
||||
@@ -2518,6 +2520,105 @@ class ServingChatTestCase(unittest.TestCase):
|
||||
"storage_backend": "file",
|
||||
},
|
||||
)
|
||||
self.assertIsNone(response.choices[0].meta_info)
|
||||
dumped_response = response.model_dump()
|
||||
self.assertIn("sglext", dumped_response)
|
||||
self.assertNotIn("meta_info", dumped_response["choices"][0])
|
||||
|
||||
def test_non_streaming_meta_info_omits_response_level_routed_experts(self):
|
||||
req = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "Hi?"}],
|
||||
max_tokens=100,
|
||||
n=2,
|
||||
return_meta_info=True,
|
||||
return_routed_experts=True,
|
||||
)
|
||||
ret = [
|
||||
{
|
||||
"text": f"Response {index}",
|
||||
"meta_info": {
|
||||
"id": "chatcmpl-meta-test",
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 2,
|
||||
"cached_tokens": index,
|
||||
"routed_experts": routed_experts,
|
||||
"finish_reason": {"type": "stop", "matched": None},
|
||||
"weight_version": "default",
|
||||
},
|
||||
}
|
||||
for index, routed_experts in enumerate(["cm91dGUtYQ==", "cm91dGUtYg=="])
|
||||
]
|
||||
|
||||
response = self.chat._build_chat_response(req, ret, 1234567890)
|
||||
|
||||
self.assertIsNone(
|
||||
response.sglext,
|
||||
"sglext is absent only when routed_experts is the sole extension",
|
||||
)
|
||||
self.assertEqual(
|
||||
[choice.meta_info for choice in response.choices],
|
||||
[ret_item["meta_info"] for ret_item in ret],
|
||||
)
|
||||
dumped_response = response.model_dump()
|
||||
self.assertNotIn("sglext", dumped_response)
|
||||
self.assertEqual(
|
||||
[choice["meta_info"] for choice in dumped_response["choices"]],
|
||||
[ret_item["meta_info"] for ret_item in ret],
|
||||
)
|
||||
serialized_response = json.dumps(dumped_response)
|
||||
self.assertEqual(serialized_response.count('"routed_experts"'), 2)
|
||||
|
||||
def test_non_streaming_meta_info_preserves_cache_and_spec_in_sglext(self):
|
||||
req = ChatCompletionRequest(
|
||||
model="x",
|
||||
messages=[{"role": "user", "content": "Hi?"}],
|
||||
max_tokens=100,
|
||||
n=2,
|
||||
return_meta_info=True,
|
||||
return_routed_experts=True,
|
||||
return_cached_tokens_details=True,
|
||||
return_spec_tokens_details=True,
|
||||
)
|
||||
routed_experts = ["cm91dGUtYQ==", "cm91dGUtYg=="]
|
||||
ret = [_spec_result(index) for index in range(2)]
|
||||
for index, ret_item in enumerate(ret):
|
||||
ret_item["meta_info"].update(
|
||||
{
|
||||
"cached_tokens_details": {
|
||||
"device": 4 - index,
|
||||
"host": index,
|
||||
},
|
||||
"routed_experts": routed_experts[index],
|
||||
}
|
||||
)
|
||||
|
||||
response = self.chat._build_chat_response(req, ret, 1234567890)
|
||||
|
||||
self.assertIsNotNone(response.sglext)
|
||||
self.assertIsNone(response.sglext.routed_experts)
|
||||
self.assertEqual(
|
||||
response.sglext.cached_tokens_details.model_dump(exclude_none=True),
|
||||
{"device": 4, "host": 0},
|
||||
)
|
||||
self.assertEqual(
|
||||
[item.spec_cap_length for item in response.sglext.spec_tokens_details],
|
||||
[1.0, 2.0],
|
||||
)
|
||||
self.assertEqual(
|
||||
[choice.meta_info for choice in response.choices],
|
||||
[ret_item["meta_info"] for ret_item in ret],
|
||||
)
|
||||
dumped_response = response.model_dump()
|
||||
self.assertIn("sglext", dumped_response)
|
||||
self.assertIn("spec_tokens_details", dumped_response["sglext"])
|
||||
self.assertNotIn("routed_experts", dumped_response["sglext"])
|
||||
self.assertEqual(
|
||||
dumped_response["sglext"]["cached_tokens_details"],
|
||||
{"device": 4, "host": 0},
|
||||
)
|
||||
serialized_response = json.dumps(dumped_response)
|
||||
self.assertEqual(serialized_response.count('"routed_experts"'), 2)
|
||||
|
||||
def test_parallel_sampling_returns_spec_details_per_choice(self):
|
||||
req = ChatCompletionRequest(
|
||||
|
||||
Reference in New Issue
Block a user