7553aac77b · tduncan/Nvidia-Optimization-Stack-VLLM

tests/entrypoints/openai/completion/test_completion.py

+16

@@ -822,3 +822,19 @@ def test_completion_request_bad_words_default_empty():

822 default_sampling_params={},

823 )

824 assert sampling_params.bad_words == []

825 +

826 +

827 +def test_completion_request_forwards_routed_experts_prompt_start():

828 + request = CompletionRequest(

829 + model="test-model",

830 + prompt="Hello",

831 + max_tokens=10,

832 + routed_experts_prompt_start=3,

833 + )

834 +

835 + sampling_params = request.to_sampling_params(

836 + max_tokens=10,

837 + default_sampling_params={},

838 + )

839 +

840 + assert sampling_params.routed_experts_prompt_start == 3

tests/entrypoints/openai/test_return_routed_experts.py

+13 -11

@@ -20,6 +20,18 @@ NUM_EXPERTS_PER_TOK = 2

20 NUM_HIDDEN_LAYERS = 2

21

22

23 +def assert_valid_routed_experts(encoded: str | None) -> None:

24 + assert encoded is not None

25 + routed_experts = np.load(io.BytesIO(base64.b64decode(encoded)))

26 + assert routed_experts.ndim == 3

27 + num_tokens, num_layers, topk = routed_experts.shape

28 + assert num_tokens > 0

29 + assert num_layers == NUM_HIDDEN_LAYERS

30 + assert topk == NUM_EXPERTS_PER_TOK

31 + assert (routed_experts >= 0).all()

32 + assert (routed_experts < NUM_LOCAL_EXPERTS).all()

33 +

34 +

35 @pytest.fixture(scope="module")

36 def server():

37 args = [

@@ -50,15 +62,5 @@ async def test_routed_experts(server):

62

63 choice = result.model_dump()["choices"][0]

64

53 - assert choice["routed_experts"] is not None

65 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()

66 + assert_valid_routed_experts(choice["routed_experts"])

tests/entrypoints/openai/test_stop_token_ids.py

+14

@@ -87,6 +87,20 @@ class TestChatCompletionStopTokenIds:

87

88 assert set(sampling_params.stop_token_ids) == {42, 43}

89

90 + def test_routed_experts_prompt_start_is_forwarded(self):

91 + request = ChatCompletionRequest(

92 + model="test-model",

93 + messages=[{"role": "user", "content": "hello"}],

94 + routed_experts_prompt_start=3,

95 + )

96 +

97 + sampling_params = request.to_sampling_params(

98 + max_tokens=100,

99 + default_sampling_params={},

100 + )

101 +

102 + assert sampling_params.routed_experts_prompt_start == 3

103 +

104 def test_duplicate_stop_token_ids_deduplicated(self):

105 """Overlapping stop_token_ids between client and server are deduplicated."""

106 request = ChatCompletionRequest(

tests/utils_/test_serial_utils.py

+16

@@ -1,5 +1,9 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import io

4 +

5 +import numpy as np

6 +import pybase64

7 import pytest

8 import torch

9

@@ -12,6 +16,7 @@ from vllm.utils.serial_utils import (

16 Endianness,

17 MmMetadataDType,

18 binary2tensor,

19 + numpy2base64,

20 tensor2binary,

21 )

22

@@ -19,6 +24,17 @@ FLOAT_EMBED_DTYPES = tuple(EMBED_DTYPES.keys())

24 INTEGER_EMBED_DTYPES = tuple(MM_METADATA_DTYPES.keys())

25

26

27 +def test_numpy2base64_round_trip():

28 + array = np.arange(24, dtype=np.uint8).reshape(2, 3, 4)

29 +

30 + decoded = np.load(

31 + io.BytesIO(pybase64.b64decode(numpy2base64(array))),

32 + allow_pickle=False,

33 + )

34 +

35 + np.testing.assert_array_equal(decoded, array)

36 +

37 +

38 def _build_integer_tensor(

39 embed_dtype: MmMetadataDType, shape: tuple[int, ...]

40 ) -> torch.Tensor:

vllm/config/vllm.py

+8

@@ -1074,6 +1074,14 @@ class VllmConfig:

1074 "--enable-return-routed-experts is incompatible with "

1075 "pipeline parallelism (PP > 1)."

1076 )

1077 + if (

1078 + self.parallel_config.decode_context_parallel_size > 1

1079 + or self.parallel_config.prefill_context_parallel_size > 1

1080 + ):

1081 + raise ValueError(

1082 + "--enable-return-routed-experts is incompatible with context "

1083 + "parallelism (DCP > 1 or PCP > 1)."

1084 + )

1085

1086 # Incompatible with any KV connector — covers both PD disaggregation

1087 # (kv_producer/kv_consumer: routing captured on P can't reach D) and

vllm/entrypoints/openai/chat_completion/protocol.py

+7 -1

@@ -114,7 +114,7 @@ class ChatCompletionResponseChoice(OpenAIBaseModel):

114 token_ids: list[int] | None = None

115 # Per-token expert routing decisions, base64-encoded ``.npy`` bytes

116 # (numpy serialization). Shape after decode:

117 - # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16

117 + # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32

118 # ``num_tokens - 1`` because the last sampled token has not been

119 # forwarded yet and therefore has no routing data.

120 # Decode:

@@ -422,6 +422,11 @@ class ChatCompletionRequest(OpenAIBaseModel):

422 "need to map generated text back to input tokens."

423 ),

424 )

425 + routed_experts_prompt_start: int = Field(

426 + default=0,

427 + ge=0,

428 + description="Skip the first N prompt tokens from returned routed-expert data.",

429 + )

430 return_token_offsets: bool | None = Field(

431 default=False,

432 description=(

@@ -741,6 +746,7 @@ class ChatCompletionRequest(OpenAIBaseModel):

746 extra_args=extra_args or None,

747 skip_clone=True, # Created fresh per request, safe to skip clone

748 repetition_detection=self.repetition_detection,

749 + routed_experts_prompt_start=self.routed_experts_prompt_start,

750 )

751

752 @model_validator(mode="before")

vllm/entrypoints/openai/chat_completion/serving.py

+6 -12

@@ -2,15 +2,12 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import asyncio

5 -import io

5 import time

6 from collections.abc import AsyncGenerator, AsyncIterator

7 from collections.abc import Sequence as GenericSequence

8 from http import HTTPStatus

9 from typing import Any, Final, cast

10

12 -import numpy as np

13 -import pybase64 as base64

11 from fastapi import Request

12

13 from vllm.engine.protocol import EngineClient

@@ -64,6 +61,7 @@ from vllm.renderers.online_renderer import OnlineRenderer

61 from vllm.sampling_params import BeamSearchParams, SamplingParams

62 from vllm.tokenizers import TokenizerLike

63 from vllm.utils.collection_utils import as_list

64 +from vllm.utils.serial_utils import numpy2base64

65

66 logger = init_logger(__name__)

67

@@ -1007,15 +1005,11 @@ class OpenAIServingChat(GenerateBaseServing):

1005 and output.finish_reason == "stop"

1006 )

1007

1010 - # Encode routed_experts for transport. JSON can't carry raw

1011 - # bytes, so we write the ndarray as a ``.npy`` byte stream

1012 - # and base64-encode it. ``pybase64`` is ~3x faster than the

1013 - # stdlib ``base64`` on large payloads thanks to SIMD.

1014 - routed_experts_b64 = None

1015 - if output.routed_experts is not None:

1016 - buf = io.BytesIO()

1017 - np.save(buf, output.routed_experts)

1018 - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii")

1008 + routed_experts_b64 = (

1009 + numpy2base64(output.routed_experts)

1010 + if output.routed_experts is not None

1011 + else None

1012 + )

1013

1014 choice_data = ChatCompletionResponseChoice(

1015 index=output.index,

vllm/entrypoints/openai/completion/protocol.py

+7 -1

@@ -175,6 +175,11 @@ class CompletionRequest(OpenAIBaseModel):

175 "need to map generated text back to input tokens."

176 ),

177 )

178 + routed_experts_prompt_start: int = Field(

179 + default=0,

180 + ge=0,

181 + description="Skip the first N prompt tokens from returned routed-expert data.",

182 + )

183 return_token_offsets: bool | None = Field(

184 default=False,

185 description=(

@@ -394,6 +399,7 @@ class CompletionRequest(OpenAIBaseModel):

399 skip_clone=True, # Created fresh per request, safe to skip clone

400 repetition_detection=self.repetition_detection,

401 thinking_token_budget=self.thinking_token_budget,

402 + routed_experts_prompt_start=self.routed_experts_prompt_start,

403 )

404

405 @model_validator(mode="before")

@@ -611,7 +617,7 @@ class CompletionResponseChoice(OpenAIBaseModel):

617 prompt_token_ids: list[int] | None = None # For prompt

618 # Per-token expert routing decisions, base64-encoded ``.npy`` bytes

619 # (numpy serialization). Shape after decode:

614 - # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16

620 + # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32

621 # ``num_tokens - 1`` because the last sampled token has not been

622 # forwarded yet and therefore has no routing data.

623 # Decode:

vllm/entrypoints/openai/completion/serving.py

+6 -14

@@ -2,14 +2,11 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 import asyncio

5 -import io

5 import time

6 from collections.abc import AsyncGenerator, AsyncIterator

7 from collections.abc import Sequence as GenericSequence

8 from typing import cast

9

11 -import numpy as np

12 -import pybase64 as base64

10 from fastapi import Request

11

12 from vllm.engine.protocol import EngineClient

@@ -48,6 +45,7 @@ from vllm.sampling_params import BeamSearchParams, SamplingParams

45 from vllm.tokenizers import TokenizerLike

46 from vllm.utils.async_utils import merge_async_iterators

47 from vllm.utils.collection_utils import as_list

48 +from vllm.utils.serial_utils import numpy2base64

49

50 logger = init_logger(__name__)

51

@@ -570,17 +568,11 @@ class OpenAIServingCompletion(GenerateBaseServing):

568 else:

569 logprobs = None

570

573 - # Encode routed_experts for transport. JSON can't carry raw

574 - # bytes, so we write the ndarray as a ``.npy`` byte stream

575 - # and base64-encode it. ``pybase64`` is ~3x faster than the

576 - # stdlib ``base64`` on large payloads thanks to SIMD.

577 - routed_experts_b64 = None

578 - if output.routed_experts is not None:

579 - buf = io.BytesIO()

580 - np.save(buf, output.routed_experts)

581 - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode(

582 - "ascii"

583 - )

571 + routed_experts_b64 = (

572 + numpy2base64(output.routed_experts)

573 + if output.routed_experts is not None

574 + else None

575 + )

576

577 choice_data = CompletionResponseChoice(

578 index=len(choices),

vllm/entrypoints/scale_out/token_in_token_out/protocol.py

+1 -1

@@ -202,7 +202,7 @@ class GenerateResponseChoice(BaseModel):

202 token_ids: list[int] | None = None

203 # Per-token expert routing decisions, base64-encoded ``.npy`` bytes

204 # (numpy serialization). Shape after decode:

205 - # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16

205 + # (num_tokens - 1, num_layers, num_experts_per_tok) dtype uint8/uint16/int32

206 # ``num_tokens - 1`` because the last sampled token has not been

207 # forwarded yet and therefore has no routing data.

208 # Decode:

vllm/entrypoints/scale_out/token_in_token_out/serving.py

+11 -21

@@ -3,14 +3,11 @@

3

4

5 import asyncio

6 -import io

6 import time

7 from collections.abc import AsyncGenerator

8 from collections.abc import Sequence as GenericSequence

9

10 import msgspec

12 -import numpy as np

13 -import pybase64 as base64

11 from fastapi import Request

12

13 from vllm.engine.protocol import EngineClient

@@ -46,6 +43,7 @@ from vllm.outputs import RequestOutput

43 from vllm.renderers.online_renderer import OnlineRenderer

44 from vllm.sampling_params import RequestOutputKind, SamplingParams

45 from vllm.utils.collection_utils import as_list

46 +from vllm.utils.serial_utils import numpy2base64

47

48 from .mm_serde import decode_mm_kwargs_item

49 from .protocol import (

@@ -303,17 +301,11 @@ class ServingTokens(GenerateBaseServing):

301 else:

302 logprobs = None

303

306 - # Encode routed_experts for transport. JSON can't carry raw

307 - # bytes, so we write the ndarray as a ``.npy`` byte stream

308 - # and base64-encode it. ``pybase64`` is ~3x faster than the

309 - # stdlib ``base64`` on large payloads thanks to SIMD.

310 - # This is the only base64 hop in the pipeline -- the

311 - # engine<->API-server link is binary msgpack + zmq.

312 - routed_experts_b64 = None

313 - if output.routed_experts is not None:

314 - buf = io.BytesIO()

315 - np.save(buf, output.routed_experts)

316 - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode("ascii")

304 + routed_experts_b64 = (

305 + numpy2base64(output.routed_experts)

306 + if output.routed_experts is not None

307 + else None

308 + )

309

310 sampling_mask = None

311 if output.sampling_mask is not None:

@@ -435,13 +427,11 @@ class ServingTokens(GenerateBaseServing):

427 else:

428 logprobs = None

429

438 - routed_experts_b64 = None

439 - if output.routed_experts is not None:

440 - buf = io.BytesIO()

441 - np.save(buf, output.routed_experts)

442 - routed_experts_b64 = base64.b64encode(buf.getvalue()).decode(

443 - "ascii"

444 - )

430 + routed_experts_b64 = (

431 + numpy2base64(output.routed_experts)

432 + if output.routed_experts is not None

433 + else None

434 + )

435

436 chunk = GenerateStreamResponse(

437 request_id=request_id,

vllm/sampling_params.py

+2

@@ -406,6 +406,7 @@ class SamplingParams(

406 skip_clone: bool = False,

407 repetition_detection: RepetitionDetectionParams | None = None,

408 logprob_token_ids: list[int] | None = None,

409 + routed_experts_prompt_start: int = 0,

410 ) -> "SamplingParams":

411 if logit_bias is not None:

412 # Fast path uses a dict comprehension; on failure we iterate once

@@ -468,6 +469,7 @@ class SamplingParams(

469 extra_args=extra_args,

470 skip_clone=skip_clone,

471 repetition_detection=repetition_detection,

472 + routed_experts_prompt_start=routed_experts_prompt_start,

473 )

474

475 def __post_init__(self) -> None:

vllm/utils/serial_utils.py

+7

@@ -64,6 +64,13 @@ def tensor2base64(x: torch.Tensor) -> str:

64 return pybase64.b64encode(binary_data).decode("utf-8")

65

66

67 +def numpy2base64(array: np.ndarray) -> str:

68 + """Encode a NumPy array using its `.npy` representation."""

69 + with io.BytesIO() as buffer:

70 + np.save(buffer, array, allow_pickle=False)

71 + return pybase64.b64encode(buffer.getbuffer()).decode("ascii")

72 +

73 +

74 def tensor2binary(

75 tensor: torch.Tensor,

76 embed_dtype: "EmbedDType | MmMetadataDType",

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