A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 05:30 UTC
[R3] Add routed experts to openai entrypoint (#38939)
Signed-off-by: ahao-anyscale <ahao@anyscale.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Aaron Hao committed
May 20, 2026 at 09:07 UTC
2d6b3489b9a325988ad52507236409747d2098a7
6 files changed
+208
tests/entrypoints/openai/test_return_routed_experts.py
new
+64
@@ -0,0 +1,64 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import io
5
+
6
+import numpy as np
7
+import pybase64 as base64
8
+import pytest
9
+
10
+from ...utils import RemoteOpenAIServer
11
+
12
+MODEL_NAME = "TitanML/tiny-mixtral"
13
+
14
+# tiny-mixtral config: 8 local experts, top-2 routing, 2 hidden layers.
15
+# The published config has sliding_window=4096, which produces
16
+# SlidingWindowSpec kv-cache groups; RoutedExpertsManager requires a
17
+# FullAttentionSpec group, so we override sliding_window=null below.
18
+NUM_LOCAL_EXPERTS = 8
19
+NUM_EXPERTS_PER_TOK = 2
20
+NUM_HIDDEN_LAYERS = 2
21
+
22
+
23
+@pytest.fixture(scope="module")
24
+def server():
25
+ args = [
26
+ "--max-model-len",
27
+ "256",
28
+ "--max-num-seqs",
29
+ "32",
30
+ "--enforce-eager",
31
+ "--enable-return-routed-experts",
32
+ "--hf-overrides",
33
+ '{"sliding_window": null}',
34
+ ]
35
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
36
+ yield remote_server
37
+
38
+
39
+@pytest.mark.asyncio
40
+async def test_routed_experts(server):
41
+ """Test that /v1/completions returns routed_experts when enabled."""
42
+ async with server.get_async_client() as client:
43
+ result = await client.completions.create(
44
+ model=MODEL_NAME,
45
+ prompt="Hello, world",
46
+ max_tokens=10,
47
+ temperature=0,
48
+ extra_body={"return_token_ids": True},
49
+ )
50
+
51
+ choice = result.model_dump()["choices"][0]
52
+
53
+ assert choice["routed_experts"] is not None
54
+ assert choice["token_ids"] is not None
55
+
56
+ # routed_experts is base64-encoded .npy bytes; decode to ndarray.
57
+ routed_experts = np.load(io.BytesIO(base64.b64decode(choice["routed_experts"])))
58
+ assert routed_experts.ndim == 3
59
+ num_tokens, num_layers, topk = routed_experts.shape
60
+ assert num_tokens > 0
61
+ assert num_layers == NUM_HIDDEN_LAYERS
62
+ assert topk == NUM_EXPERTS_PER_TOK
63
+ assert (routed_experts >= 0).all()
64
+ assert (routed_experts < NUM_LOCAL_EXPERTS).all()
tests/entrypoints/serve/disagg/test_return_routed_experts.py
new
+81
@@ -0,0 +1,81 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import io
5
+
6
+import httpx
7
+import numpy as np
8
+import pybase64 as base64
9
+import pytest
10
+import pytest_asyncio
11
+
12
+from tests.utils import RemoteOpenAIServer
13
+
14
+MODEL_NAME = "TitanML/tiny-mixtral"
15
+GEN_ENDPOINT = "/inference/v1/generate"
16
+
17
+# tiny-mixtral config: 8 local experts, top-2 routing, 2 hidden layers.
18
+# The published config has sliding_window=4096, which produces
19
+# SlidingWindowSpec kv-cache groups; RoutedExpertsManager requires a
20
+# FullAttentionSpec group, so we override sliding_window=null below.
21
+NUM_LOCAL_EXPERTS = 8
22
+NUM_EXPERTS_PER_TOK = 2
23
+NUM_HIDDEN_LAYERS = 2
24
+
25
+
26
+@pytest.fixture(scope="module")
27
+def server():
28
+ args = [
29
+ "--max-model-len",
30
+ "256",
31
+ "--max-num-seqs",
32
+ "32",
33
+ "--enforce-eager",
34
+ "--enable-return-routed-experts",
35
+ "--hf-overrides",
36
+ '{"sliding_window": null}',
37
+ ]
38
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
39
+ yield remote_server
40
+
41
+
42
+@pytest_asyncio.fixture
43
+async def client(server: RemoteOpenAIServer):
44
+ transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
45
+ headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
46
+ async with httpx.AsyncClient(
47
+ transport=transport,
48
+ base_url=server.url_root,
49
+ timeout=600,
50
+ headers=headers,
51
+ ) as c:
52
+ yield c
53
+
54
+
55
+@pytest.mark.asyncio
56
+async def test_generate_routed_experts(client):
57
+ """Test that /inference/v1/generate returns routed_experts when enabled."""
58
+ payload = {
59
+ "model": MODEL_NAME,
60
+ "token_ids": [1, 2, 3],
61
+ "sampling_params": {"max_tokens": 10, "temperature": 0.0},
62
+ "stream": False,
63
+ }
64
+ resp = await client.post(GEN_ENDPOINT, json=payload)
65
+ resp.raise_for_status()
66
+ data = resp.json()
67
+
68
+ choice = data["choices"][0]
69
+
70
+ assert choice["routed_experts"] is not None
71
+ assert choice["token_ids"] is not None
72
+
73
+ # routed_experts is base64-encoded .npy bytes; decode to ndarray.
74
+ routed_experts = np.load(io.BytesIO(base64.b64decode(choice["routed_experts"])))
75
+ assert routed_experts.ndim == 3
76
+ num_tokens, num_layers, topk = routed_experts.shape
77
+ assert num_tokens > 0
78
+ assert num_layers == NUM_HIDDEN_LAYERS
79
+ assert topk == NUM_EXPERTS_PER_TOK
80
+ assert (routed_experts >= 0).all()
81
+ assert (routed_experts < NUM_LOCAL_EXPERTS).all()
vllm/entrypoints/openai/chat_completion/protocol.py
+10
@@ -92,6 +92,16 @@ class ChatCompletionResponseChoice(OpenAIBaseModel):
92
# not part of the OpenAI spec but is useful for tracing the tokens
93
# in agent scenarios
94
token_ids: list[int] | None = None
95
+ # Per-token expert routing decisions, base64-encoded ``.npy`` bytes
96
+ # (numpy serialization). Shape after decode:
97
+ # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16
98
+ # ``num_tokens - 1`` because the last sampled token has not been
99
+ # forwarded yet and therefore has no routing data.
100
+ # Decode:
101
+ # np.load(io.BytesIO(base64.b64decode(s)))
102
+ # ``None`` if (a) the request was aborted before any forward pass,
103
+ # or (b) ``enable_return_routed_experts`` is off server-side.
104
+ routed_experts: str | None = None
105
106
107
class ChatCompletionResponse(OpenAIBaseModel):
vllm/entrypoints/openai/chat_completion/serving.py
+27
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
+import io
6
import json
7
import time
8
from collections.abc import AsyncGenerator, AsyncIterator
@@ -9,6 +10,8 @@ from collections.abc import Sequence as GenericSequence
10
from http import HTTPStatus
11
from typing import TYPE_CHECKING, Any, Final
12
13
+import numpy as np
14
+import pybase64 as base64
15
from fastapi import Request
16
17
from vllm.engine.protocol import EngineClient
@@ -1086,6 +1089,18 @@ class OpenAIServingChat(OpenAIServing):
1089
content=content,
1090
)
1091
1092
+ # Encode routed_experts for transport. JSON can't carry raw
1093
+ # bytes, so we write the ndarray as a ``.npy`` byte stream
1094
+ # and base64-encode it. ``pybase64`` is ~3x faster than the
1095
+ # stdlib ``base64`` on large payloads thanks to SIMD.
1096
+ routed_experts_b64 = None
1097
+ if output.routed_experts is not None:
1098
+ buf = io.BytesIO()
1099
+ np.save(buf, output.routed_experts)
1100
+ routed_experts_b64 = base64.b64encode(buf.getvalue()).decode(
1101
+ "ascii"
1102
+ )
1103
+
1104
choice_data = ChatCompletionResponseChoice(
1105
index=output.index,
1106
message=message,
@@ -1101,6 +1116,7 @@ class OpenAIServingChat(OpenAIServing):
1116
token_ids=(
1117
as_list(output.token_ids) if request.return_token_ids else None
1118
),
1119
+ routed_experts=routed_experts_b64,
1120
)
1121
choices.append(choice_data)
1122
continue
@@ -1309,6 +1325,16 @@ class OpenAIServingChat(OpenAIServing):
1325
and output.finish_reason == "stop"
1326
)
1327
1328
+ # Encode routed_experts for transport. JSON can't carry raw
1329
+ # bytes, so we write the ndarray as a ``.npy`` byte stream
1330
+ # and base64-encode it. ``pybase64`` is ~3x faster than the
1331
+ # stdlib ``base64`` on large payloads thanks to SIMD.
1332
+ routed_experts_b64 = None
1333
+ if output.routed_experts is not None:
1334
+ buf = io.BytesIO()
1335
+ np.save(buf, output.routed_experts)
1336
+ routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii")
1337
+
1338
choice_data = ChatCompletionResponseChoice(
1339
index=output.index,
1340
message=message,
@@ -1322,6 +1348,7 @@ class OpenAIServingChat(OpenAIServing):
1348
token_ids=(
1349
as_list(output.token_ids) if request.return_token_ids else None
1350
),
1351
+ routed_experts=routed_experts_b64,
1352
)
1353
choice_data = maybe_filter_parallel_tool_calls(choice_data, request)
1354
vllm/entrypoints/openai/completion/protocol.py
+10
@@ -477,6 +477,16 @@ class CompletionResponseChoice(OpenAIBaseModel):
477
token_ids: list[int] | None = None # For response
478
prompt_logprobs: list[dict[int, Logprob] | None] | None = None
479
prompt_token_ids: list[int] | None = None # For prompt
480
+ # Per-token expert routing decisions, base64-encoded ``.npy`` bytes
481
+ # (numpy serialization). Shape after decode:
482
+ # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16
483
+ # ``num_tokens - 1`` because the last sampled token has not been
484
+ # forwarded yet and therefore has no routing data.
485
+ # Decode:
486
+ # np.load(io.BytesIO(base64.b64decode(s)))
487
+ # ``None`` if (a) the request was aborted before any forward pass,
488
+ # or (b) ``enable_return_routed_experts`` is off server-side.
489
+ routed_experts: str | None = None
490
491
492
class CompletionResponse(OpenAIBaseModel):
vllm/entrypoints/openai/completion/serving.py
+16
@@ -2,11 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
+import io
6
import time
7
from collections.abc import AsyncGenerator, AsyncIterator
8
from collections.abc import Sequence as GenericSequence
9
from typing import TYPE_CHECKING, cast
10
11
+import numpy as np
12
+import pybase64 as base64
13
from fastapi import Request
14
15
from vllm.engine.protocol import EngineClient
@@ -538,6 +541,18 @@ class OpenAIServingCompletion(OpenAIServing):
541
else:
542
logprobs = None
543
544
+ # Encode routed_experts for transport. JSON can't carry raw
545
+ # bytes, so we write the ndarray as a ``.npy`` byte stream
546
+ # and base64-encode it. ``pybase64`` is ~3x faster than the
547
+ # stdlib ``base64`` on large payloads thanks to SIMD.
548
+ routed_experts_b64 = None
549
+ if output.routed_experts is not None:
550
+ buf = io.BytesIO()
551
+ np.save(buf, output.routed_experts)
552
+ routed_experts_b64 = base64.b64encode(buf.getvalue()).decode(
553
+ "ascii"
554
+ )
555
+
556
choice_data = CompletionResponseChoice(
557
index=len(choices),
558
text=output_text,
@@ -551,6 +566,7 @@ class OpenAIServingCompletion(OpenAIServing):
566
token_ids=(
567
as_list(output.token_ids) if request.return_token_ids else None
568
),
569
+ routed_experts=routed_experts_b64,
570
)
571
choices.append(choice_data)
572