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)
|
#### 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
|
```python Example
|
||||||
# Example with logit_bias parameter for completions API
|
# 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)
|
# Build sglext at response level (from first ret_item, as these are per-request)
|
||||||
first_ret = ret[0]
|
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(
|
cached_tokens_details = process_cached_tokens_details_from_ret(
|
||||||
first_ret, request
|
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",
|
"empty-delta logprobs chunk emitted without a parser; would break client chunk-shape assumptions",
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_non_streaming_cached_tokens_details_emits_sglext(self):
|
def test_non_streaming_extension_fields_emit_sglext_without_meta_info(self):
|
||||||
"""Test that non-streaming chat responses emit cached token details in sglext."""
|
|
||||||
|
|
||||||
req = ChatCompletionRequest(
|
req = ChatCompletionRequest(
|
||||||
model="x",
|
model="x",
|
||||||
messages=[{"role": "user", "content": "Hi?"}],
|
messages=[{"role": "user", "content": "Hi?"}],
|
||||||
max_tokens=100,
|
max_tokens=100,
|
||||||
|
return_meta_info=False,
|
||||||
|
return_routed_experts=True,
|
||||||
return_cached_tokens_details=True,
|
return_cached_tokens_details=True,
|
||||||
)
|
)
|
||||||
ret = [
|
ret = [
|
||||||
@@ -2500,6 +2500,7 @@ class ServingChatTestCase(unittest.TestCase):
|
|||||||
"storage": 1,
|
"storage": 1,
|
||||||
"storage_backend": "file",
|
"storage_backend": "file",
|
||||||
},
|
},
|
||||||
|
"routed_experts": "cm91dGUtYQ==",
|
||||||
"finish_reason": {"type": "stop", "matched": None},
|
"finish_reason": {"type": "stop", "matched": None},
|
||||||
"weight_version": "default",
|
"weight_version": "default",
|
||||||
},
|
},
|
||||||
@@ -2509,6 +2510,7 @@ class ServingChatTestCase(unittest.TestCase):
|
|||||||
response = self.chat._build_chat_response(req, ret, 1234567890)
|
response = self.chat._build_chat_response(req, ret, 1234567890)
|
||||||
|
|
||||||
self.assertIsNotNone(response.sglext)
|
self.assertIsNotNone(response.sglext)
|
||||||
|
self.assertEqual(response.sglext.routed_experts, "cm91dGUtYQ==")
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
response.sglext.cached_tokens_details.model_dump(exclude_none=True),
|
response.sglext.cached_tokens_details.model_dump(exclude_none=True),
|
||||||
{
|
{
|
||||||
@@ -2518,6 +2520,105 @@ class ServingChatTestCase(unittest.TestCase):
|
|||||||
"storage_backend": "file",
|
"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):
|
def test_parallel_sampling_returns_spec_details_per_choice(self):
|
||||||
req = ChatCompletionRequest(
|
req = ChatCompletionRequest(
|
||||||
|
|||||||
Reference in New Issue
Block a user