[Frontend] Add routed-experts prompt offset (#51906)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

tests/entrypoints/openai/completion/test_completion.py

@@ -822,3 +822,19 @@ def test_completion_request_bad_words_default_empty():
         default_sampling_params={},
     )
     assert sampling_params.bad_words == []
+
+
+def test_completion_request_forwards_routed_experts_prompt_start():
+    request = CompletionRequest(
+        model="test-model",
+        prompt="Hello",
+        max_tokens=10,
+        routed_experts_prompt_start=3,
+    )
+
+    sampling_params = request.to_sampling_params(
+        max_tokens=10,
+        default_sampling_params={},
+    )
+
+    assert sampling_params.routed_experts_prompt_start == 3

tests/entrypoints/openai/test_return_routed_experts.py

@@ -20,6 +20,18 @@ NUM_EXPERTS_PER_TOK = 2
 NUM_HIDDEN_LAYERS = 2
 
 
+def assert_valid_routed_experts(encoded: str | None) -> None:
+    assert encoded is not None
+    routed_experts = np.load(io.BytesIO(base64.b64decode(encoded)))
+    assert routed_experts.ndim == 3
+    num_tokens, num_layers, topk = routed_experts.shape
+    assert num_tokens > 0
+    assert num_layers == NUM_HIDDEN_LAYERS
+    assert topk == NUM_EXPERTS_PER_TOK
+    assert (routed_experts >= 0).all()
+    assert (routed_experts < NUM_LOCAL_EXPERTS).all()
+
+
 @pytest.fixture(scope="module")
 def server():
     args = [
@@ -50,15 +62,5 @@ async def test_routed_experts(server):
 
         choice = result.model_dump()["choices"][0]
 
-        assert choice["routed_experts"] is not None
         assert choice["token_ids"] is not None
-
-        # routed_experts is base64-encoded .npy bytes; decode to ndarray.
-        routed_experts = np.load(io.BytesIO(base64.b64decode(choice["routed_experts"])))
-        assert routed_experts.ndim == 3
-        num_tokens, num_layers, topk = routed_experts.shape
-        assert num_tokens > 0
-        assert num_layers == NUM_HIDDEN_LAYERS
-        assert topk == NUM_EXPERTS_PER_TOK
-        assert (routed_experts >= 0).all()
-        assert (routed_experts < NUM_LOCAL_EXPERTS).all()
+        assert_valid_routed_experts(choice["routed_experts"])

tests/entrypoints/openai/test_stop_token_ids.py

@@ -87,6 +87,20 @@ class TestChatCompletionStopTokenIds:
 
         assert set(sampling_params.stop_token_ids) == {42, 43}
 
+    def test_routed_experts_prompt_start_is_forwarded(self):
+        request = ChatCompletionRequest(
+            model="test-model",
+            messages=[{"role": "user", "content": "hello"}],
+            routed_experts_prompt_start=3,
+        )
+
+        sampling_params = request.to_sampling_params(
+            max_tokens=100,
+            default_sampling_params={},
+        )
+
+        assert sampling_params.routed_experts_prompt_start == 3
+
     def test_duplicate_stop_token_ids_deduplicated(self):
         """Overlapping stop_token_ids between client and server are deduplicated."""
         request = ChatCompletionRequest(

tests/utils_/test_serial_utils.py

@@ -1,5 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import io
+
+import numpy as np
+import pybase64
 import pytest
 import torch
 
@@ -12,6 +16,7 @@ from vllm.utils.serial_utils import (
     Endianness,
     MmMetadataDType,
     binary2tensor,
+    numpy2base64,
     tensor2binary,
 )
 
@@ -19,6 +24,17 @@ FLOAT_EMBED_DTYPES = tuple(EMBED_DTYPES.keys())
 INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys())
 
 
+def test_numpy2base64_round_trip():
+    array = np.arange(24, dtype=np.uint8).reshape(2, 3, 4)
+
+    decoded = np.load(
+        io.BytesIO(pybase64.b64decode(numpy2base64(array))),
+        allow_pickle=False,
+    )
+
+    np.testing.assert_array_equal(decoded, array)
+
+
 def _build_integer_tensor(
     embed_dtype: MmMetadataDType, shape: tuple[int, ...]
 ) -> torch.Tensor:

vllm/config/vllm.py

@@ -1074,6 +1074,14 @@ class VllmConfig:
                     "--enable-return-routed-experts is incompatible with "
                     "pipeline parallelism (PP > 1)."
                 )
+            if (
+                self.parallel_config.decode_context_parallel_size > 1
+                or self.parallel_config.prefill_context_parallel_size > 1
+            ):
+                raise ValueError(
+                    "--enable-return-routed-experts is incompatible with context "
+                    "parallelism (DCP > 1 or PCP > 1)."
+                )
 
             # Incompatible with any KV connector — covers both PD disaggregation
             # (kv_producer/kv_consumer: routing captured on P can't reach D) and

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -114,7 +114,7 @@ class ChatCompletionResponseChoice(OpenAIBaseModel):
     token_ids: list[int] | None = None
     # Per-token expert routing decisions, base64-encoded ``.npy`` bytes
     # (numpy serialization). Shape after decode:
-    #   (num_tokens - 1, num_layers, num_experts_per_tok)  dtype uint8/uint16
+    #   (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32
     # ``num_tokens - 1`` because the last sampled token has not been
     # forwarded yet and therefore has no routing data.
     # Decode:
@@ -422,6 +422,11 @@ class ChatCompletionRequest(OpenAIBaseModel):
             "need to map generated text back to input tokens."
         ),
     )
+    routed_experts_prompt_start: int = Field(
+        default=0,
+        ge=0,
+        description="Skip the first N prompt tokens from returned routed-expert data.",
+    )
     return_token_offsets: bool | None = Field(
         default=False,
         description=(
@@ -741,6 +746,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
             extra_args=extra_args or None,
             skip_clone=True,  # Created fresh per request, safe to skip clone
             repetition_detection=self.repetition_detection,
+            routed_experts_prompt_start=self.routed_experts_prompt_start,
         )
 
     @model_validator(mode="before")

vllm/entrypoints/openai/chat_completion/serving.py

@@ -2,15 +2,12 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import asyncio
-import io
 import time
 from collections.abc import AsyncGenerator, AsyncIterator
 from collections.abc import Sequence as GenericSequence
 from http import HTTPStatus
 from typing import Any, Final, cast
 
-import numpy as np
-import pybase64 as base64
 from fastapi import Request
 
 from vllm.engine.protocol import EngineClient
@@ -64,6 +61,7 @@ from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.sampling_params import BeamSearchParams, SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.utils.collection_utils import as_list
+from vllm.utils.serial_utils import numpy2base64
 
 logger = init_logger(__name__)
 
@@ -1007,15 +1005,11 @@ class OpenAIServingChat(GenerateBaseServing):
                 and output.finish_reason == "stop"
             )
 
-            # Encode routed_experts for transport. JSON can't carry raw
-            # bytes, so we write the ndarray as a ``.npy`` byte stream
-            # and base64-encode it. ``pybase64`` is ~3x faster than the
-            # stdlib ``base64`` on large payloads thanks to SIMD.
-            routed_experts_b64 = None
-            if output.routed_experts is not None:
-                buf = io.BytesIO()
-                np.save(buf, output.routed_experts)
-                routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+            routed_experts_b64 = (
+                numpy2base64(output.routed_experts)
+                if output.routed_experts is not None
+                else None
+            )
 
             choice_data = ChatCompletionResponseChoice(
                 index=output.index,

vllm/entrypoints/openai/completion/protocol.py

@@ -175,6 +175,11 @@ class CompletionRequest(OpenAIBaseModel):
             "need to map generated text back to input tokens."
         ),
     )
+    routed_experts_prompt_start: int = Field(
+        default=0,
+        ge=0,
+        description="Skip the first N prompt tokens from returned routed-expert data.",
+    )
     return_token_offsets: bool | None = Field(
         default=False,
         description=(
@@ -394,6 +399,7 @@ class CompletionRequest(OpenAIBaseModel):
             skip_clone=True,  # Created fresh per request, safe to skip clone
             repetition_detection=self.repetition_detection,
             thinking_token_budget=self.thinking_token_budget,
+            routed_experts_prompt_start=self.routed_experts_prompt_start,
         )
 
     @model_validator(mode="before")
@@ -611,7 +617,7 @@ class CompletionResponseChoice(OpenAIBaseModel):
     prompt_token_ids: list[int] | None = None  # For prompt
     # Per-token expert routing decisions, base64-encoded ``.npy`` bytes
     # (numpy serialization). Shape after decode:
-    #   (num_tokens - 1, num_layers, num_experts_per_tok)  dtype uint8/uint16
+    #   (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32
     # ``num_tokens - 1`` because the last sampled token has not been
     # forwarded yet and therefore has no routing data.
     # Decode:

vllm/entrypoints/openai/completion/serving.py

@@ -2,14 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import asyncio
-import io
 import time
 from collections.abc import AsyncGenerator, AsyncIterator
 from collections.abc import Sequence as GenericSequence
 from typing import cast
 
-import numpy as np
-import pybase64 as base64
 from fastapi import Request
 
 from vllm.engine.protocol import EngineClient
@@ -48,6 +45,7 @@ from vllm.sampling_params import BeamSearchParams, SamplingParams
 from vllm.tokenizers import TokenizerLike
 from vllm.utils.async_utils import merge_async_iterators
 from vllm.utils.collection_utils import as_list
+from vllm.utils.serial_utils import numpy2base64
 
 logger = init_logger(__name__)
 
@@ -570,17 +568,11 @@ class OpenAIServingCompletion(GenerateBaseServing):
                 else:
                     logprobs = None
 
-                # Encode routed_experts for transport. JSON can't carry raw
-                # bytes, so we write the ndarray as a ``.npy`` byte stream
-                # and base64-encode it. ``pybase64`` is ~3x faster than the
-                # stdlib ``base64`` on large payloads thanks to SIMD.
-                routed_experts_b64 = None
-                if output.routed_experts is not None:
-                    buf = io.BytesIO()
-                    np.save(buf, output.routed_experts)
-                    routed_experts_b64 = base64.b64encode(buf.getvalue()).decode(
-                        "ascii"
-                    )
+                routed_experts_b64 = (
+                    numpy2base64(output.routed_experts)
+                    if output.routed_experts is not None
+                    else None
+                )
 
                 choice_data = CompletionResponseChoice(
                     index=len(choices),

vllm/entrypoints/scale_out/token_in_token_out/protocol.py

@@ -202,7 +202,7 @@ class GenerateResponseChoice(BaseModel):
     token_ids: list[int] | None = None
     # Per-token expert routing decisions, base64-encoded ``.npy`` bytes
     # (numpy serialization). Shape after decode:
-    #   (num_tokens - 1, num_layers, num_experts_per_tok)  dtype uint8/uint16
+    #   (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32
     # ``num_tokens - 1`` because the last sampled token has not been
     # forwarded yet and therefore has no routing data.
     # Decode:

vllm/entrypoints/scale_out/token_in_token_out/serving.py

@@ -3,14 +3,11 @@
 
 
 import asyncio
-import io
 import time
 from collections.abc import AsyncGenerator
 from collections.abc import Sequence as GenericSequence
 
 import msgspec
-import numpy as np
-import pybase64 as base64
 from fastapi import Request
 
 from vllm.engine.protocol import EngineClient
@@ -46,6 +43,7 @@ from vllm.outputs import RequestOutput
 from vllm.renderers.online_renderer import OnlineRenderer
 from vllm.sampling_params import RequestOutputKind, SamplingParams
 from vllm.utils.collection_utils import as_list
+from vllm.utils.serial_utils import numpy2base64
 
 from .mm_serde import decode_mm_kwargs_item
 from .protocol import (
@@ -303,17 +301,11 @@ class ServingTokens(GenerateBaseServing):
             else:
                 logprobs = None
 
-            # Encode routed_experts for transport. JSON can't carry raw
-            # bytes, so we write the ndarray as a ``.npy`` byte stream
-            # and base64-encode it. ``pybase64`` is ~3x faster than the
-            # stdlib ``base64`` on large payloads thanks to SIMD.
-            # This is the only base64 hop in the pipeline -- the
-            # engine<->API-server link is binary msgpack + zmq.
-            routed_experts_b64 = None
-            if output.routed_experts is not None:
-                buf = io.BytesIO()
-                np.save(buf, output.routed_experts)
-                routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii")
+            routed_experts_b64 = (
+                numpy2base64(output.routed_experts)
+                if output.routed_experts is not None
+                else None
+            )
 
             sampling_mask = None
             if output.sampling_mask is not None:
@@ -435,13 +427,11 @@ class ServingTokens(GenerateBaseServing):
                     else:
                         logprobs = None
 
-                    routed_experts_b64 = None
-                    if output.routed_experts is not None:
-                        buf = io.BytesIO()
-                        np.save(buf, output.routed_experts)
-                        routed_experts_b64 = base64.b64encode(buf.getvalue()).decode(
-                            "ascii"
-                        )
+                    routed_experts_b64 = (
+                        numpy2base64(output.routed_experts)
+                        if output.routed_experts is not None
+                        else None
+                    )
 
                     chunk = GenerateStreamResponse(
                         request_id=request_id,

vllm/sampling_params.py

@@ -406,6 +406,7 @@ class SamplingParams(
         skip_clone: bool = False,
         repetition_detection: RepetitionDetectionParams | None = None,
         logprob_token_ids: list[int] | None = None,
+        routed_experts_prompt_start: int = 0,
     ) -> "SamplingParams":
         if logit_bias is not None:
             # Fast path uses a dict comprehension; on failure we iterate once
@@ -468,6 +469,7 @@ class SamplingParams(
             extra_args=extra_args,
             skip_clone=skip_clone,
             repetition_detection=repetition_detection,
+            routed_experts_prompt_start=routed_experts_prompt_start,
         )
 
     def __post_init__(self) -> None:

vllm/utils/serial_utils.py

@@ -64,6 +64,13 @@ def tensor2base64(x: torch.Tensor) -> str:
     return pybase64.b64encode(binary_data).decode("utf-8")
 
 
+def numpy2base64(array: np.ndarray) -> str:
+    """Encode a NumPy array using its `.npy` representation."""
+    with io.BytesIO() as buffer:
+        np.save(buffer, array, allow_pickle=False)
+        return pybase64.b64encode(buffer.getbuffer()).decode("ascii")
+
+
 def tensor2binary(
     tensor: torch.Tensor,
     embed_dtype: "EmbedDType | MmMetadataDType",