@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 999dfc1622

Fork 0

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 5, 22:15 UTC

[Bugfix] Offload blocking tokenizer ops to shared thread pool to unblock event loop (#34789)

Signed-off-by: Bvicii <yizhanhuang2002@gmail.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

Bvicii committed Mar 26, 2026 at 22:17 UTC 999dfc1622d042dffbe4e273056fa4554453c8e7

15 files changed +195 -28

tests/entrypoints/openai/chat_completion/test_chat_error.py

@@ -55,6 +55,7 @@ class MockModelConfig:

55 skip_tokenizer_init = False

56 is_encoder_decoder: bool = False

57 is_multimodal_model: bool = False

58 + renderer_num_workers: int = 1

60 def get_diff_sampling_param(self):

61 return self.diff_sampling_param or {}

tests/entrypoints/openai/chat_completion/test_serving_chat.py

@@ -536,6 +536,7 @@ class MockModelConfig:

536 skip_tokenizer_init: bool = False

537 is_encoder_decoder: bool = False

538 is_multimodal_model: bool = False

539 + renderer_num_workers: int = 1

540

541 def get_diff_sampling_param(self):

542 return self.diff_sampling_param or {}

tests/entrypoints/openai/completion/test_completion_error.py

@@ -54,6 +54,7 @@ class MockModelConfig:

54 skip_tokenizer_init = False

55 is_encoder_decoder: bool = False

56 is_multimodal_model: bool = False

57 + renderer_num_workers: int = 1

59 def get_diff_sampling_param(self):

60 return self.diff_sampling_param or {}

tests/entrypoints/openai/completion/test_lora_resolvers.py

@@ -54,6 +54,7 @@ class MockModelConfig:

54 skip_tokenizer_init: bool = False

55 is_encoder_decoder: bool = False

56 is_multimodal_model: bool = False

57 + renderer_num_workers: int = 1

59 def get_diff_sampling_param(self):

60 return self.diff_sampling_param or {}

tests/renderers/test_completions.py

@@ -38,6 +38,7 @@ class MockModelConfig:

38 skip_tokenizer_init: bool = False

39 is_encoder_decoder: bool = False

40 is_multimodal_model: bool = False

41 + renderer_num_workers: int = 1

44 @dataclass

tests/renderers/test_mistral.py

@@ -37,6 +37,7 @@ class MockModelConfig:

37 skip_tokenizer_init: bool = False

38 is_encoder_decoder: bool = False

39 is_multimodal_model: bool = False

40 + renderer_num_workers: int = 1

43 @dataclass

vllm/config/model.py

@@ -291,6 +291,10 @@ class ModelConfig:

291 definitions"""

292 io_processor_plugin: str | None = None

293 """IOProcessor plugin name to load at model startup"""

294 + renderer_num_workers: int = 1

295 + """Number of worker threads in the renderer thread pool. This pool

296 + handles async tokenization, chat template rendering, and multimodal

297 + preprocessing."""

298

299 # Pooler config

300 pooler_config: PoolerConfig | None = None

vllm/engine/arg_utils.py

@@ -508,6 +508,7 @@ class EngineArgs:

508 MultiModalConfig.mm_encoder_attn_backend

509 )

510 io_processor_plugin: str | None = None

511 + renderer_num_workers: int = 1

512 skip_mm_profiling: bool = MultiModalConfig.skip_mm_profiling

513 video_pruning_rate: float | None = MultiModalConfig.video_pruning_rate

514 mm_tensor_ipc: MMTensorIPC = MultiModalConfig.mm_tensor_ipc

@@ -767,6 +768,10 @@ class EngineArgs:

768 model_group.add_argument(

769 "--io-processor-plugin", **model_kwargs["io_processor_plugin"]

770 )

771 + model_group.add_argument(

772 + "--renderer-num-workers",

773 + **model_kwargs["renderer_num_workers"],

774 + )

775

776 # Model loading arguments

777 load_kwargs = get_kwargs(LoadConfig)

@@ -1438,6 +1443,7 @@ class EngineArgs:

1443 video_pruning_rate=self.video_pruning_rate,

1444 mm_tensor_ipc=self.mm_tensor_ipc,

1445 io_processor_plugin=self.io_processor_plugin,

1446 + renderer_num_workers=self.renderer_num_workers,

1447 )

1448

1449 def validate_tensorizer_args(self):

vllm/renderers/base.py

+118 -6

@@ -5,6 +5,7 @@ import copy

5 import time

6 from abc import ABC, abstractmethod

7 from collections.abc import Mapping, Sequence

8 +from concurrent.futures import Executor, ThreadPoolExecutor

9 from functools import cached_property

10 from typing import TYPE_CHECKING, Any, Generic, overload

@@ -38,7 +39,10 @@ from vllm.multimodal.processing import BaseMultiModalProcessor

39 from vllm.multimodal.processing import ProcessorInputs as MMProcessorInputs

40 from vllm.multimodal.registry import MultiModalTimingRegistry

41 from vllm.tokenizers import TokenizerLike

41 -from vllm.utils.async_utils import AsyncMicrobatchTokenizer

42 +from vllm.utils.async_utils import (

43 + AsyncMicrobatchTokenizer,

44 + make_async,

45 +)

46 from vllm.utils.counter import AtomicCounter

47 from vllm.utils.torch_utils import set_default_torch_num_threads

48 from vllm.v1.metrics.stats import MultiModalCacheStats

@@ -78,11 +82,28 @@ class BaseRenderer(ABC, Generic[_T]):

83 self.tokenizer = tokenizer

85 + # Shared thread pool executor for blocking tokenizer and

86 + # multimodal preprocessing operations. The multimodal processor

87 + # receives a deep-copied tokenizer (see #36557) so it is safe to

88 + # run tokenization and MM preprocessing concurrently.

89 + pool_workers = config.model_config.renderer_num_workers

90 + self._executor = ThreadPoolExecutor(max_workers=pool_workers)

91 +

92 + # Multimodal preprocessing is always offloaded to the thread pool

93 + # to keep the asyncio event loop responsive under concurrent load.

94 + self._mm_executor: Executor = self._executor

95 +

96 # Lazy initialization since offline LLM doesn't use async

97 self._async_tokenizer: AsyncMicrobatchTokenizer | None = None

99 self.mm_processor: BaseMultiModalProcessor | None = None

100 self._mm_cache_stats: MultiModalCacheStats | None = None

101 + self._clear_mm_cache_async = make_async(

102 + self.clear_mm_cache, executor=self._executor

103 + )

104 + self._process_multimodal_async = make_async(

105 + self._process_multimodal, executor=self._mm_executor

106 + )

107 if config.model_config.is_multimodal_model:

108 mm_processor_cache = mm_registry.processor_cache_from_config(config)

109

@@ -119,7 +140,9 @@ class BaseRenderer(ABC, Generic[_T]):

140

141 def get_async_tokenizer(self) -> AsyncMicrobatchTokenizer:

142 if self._async_tokenizer is None:

122 - self._async_tokenizer = AsyncMicrobatchTokenizer(self.get_tokenizer())

143 + self._async_tokenizer = AsyncMicrobatchTokenizer(

144 + self.get_tokenizer(), executor=self._executor

145 + )

146

147 return self._async_tokenizer

148

@@ -211,11 +234,24 @@ class BaseRenderer(ABC, Generic[_T]):

234 finally:

235 self.clear_mm_cache()

236

237 + async def clear_mm_cache_async(self) -> None:

238 + """Serialize clear_mm_cache through the shared executor to avoid

239 + races with concurrent process_inputs on the mm_processor_cache."""

240 + await self._clear_mm_cache_async()

241 +

242 def shutdown(self) -> None:

243 mm_processor_cache = self.mm_processor_cache

244 if mm_processor_cache is not None:

245 mm_processor_cache.close()

246

247 + if executor := getattr(self, "_executor", None):

248 + executor.shutdown(wait=False)

249 +

250 + if (

251 + mm_executor := getattr(self, "_mm_executor", None)

252 + ) is not None and mm_executor is not executor:

253 + mm_executor.shutdown(wait=False)

254 +

255 def get_bos_token_id(self) -> int | None:

256 if self.tokenizer is None:

257 logger.warning_once(

@@ -621,6 +657,9 @@ class BaseRenderer(ABC, Generic[_T]):

657 self,

658 prompt: TokensPrompt,

659 ) -> TokensInput | MultiModalInput:

660 + """Process token inputs, with multimodal preprocessing offloaded

661 + to the shared thread pool in the async variant.

662 + """

663 prompt_token_ids = prompt["prompt_token_ids"]

664

665 engine_input: TokensInput | MultiModalInput

@@ -670,12 +709,46 @@ class BaseRenderer(ABC, Generic[_T]):

709 cache_salt=prompt.get("cache_salt"),

710 )

711

712 + async def _process_tokens_async(

713 + self,

714 + prompt: TokensPrompt,

715 + ) -> TokensInput | MultiModalInput:

716 + prompt_token_ids = prompt["prompt_token_ids"]

717 +

718 + engine_input: TokensInput | MultiModalInput

719 + if multi_modal_data := prompt.get("multi_modal_data"):

720 + engine_input = await self._process_multimodal_async(

721 + prompt_token_ids,

722 + multi_modal_data,

723 + mm_processor_kwargs=prompt.get("mm_processor_kwargs"),

724 + tokenization_kwargs=None,

725 + mm_uuids=prompt.get("multi_modal_uuids"),

726 + )

727 + else:

728 + engine_input = tokens_input(prompt_token_ids)

729 +

730 + if prompt_text := prompt.get("prompt"):

731 + engine_input["prompt"] = prompt_text

732 + if cache_salt := prompt.get("cache_salt"):

733 + engine_input["cache_salt"] = cache_salt

734 +

735 + return engine_input

736 +

737 def _process_singleton(self, prompt: SingletonTokPrompt) -> SingletonInput:

738 if "prompt_embeds" in prompt:

739 return self._process_embeds(prompt) # type: ignore[arg-type]

740

741 return self._process_tokens(prompt) # type: ignore[arg-type]

742

743 + async def _process_singleton_async(

744 + self,

745 + prompt: SingletonTokPrompt,

746 + ) -> SingletonInput:

747 + if "prompt_embeds" in prompt:

748 + return self._process_embeds(prompt) # type: ignore[arg-type]

749 +

750 + return await self._process_tokens_async(prompt) # type: ignore[arg-type]

751 +

752 def _process_enc_dec(

753 self,

754 prompt: EncoderDecoderTokPrompt,

@@ -699,6 +772,28 @@ class BaseRenderer(ABC, Generic[_T]):

772 skip_decoder_start_token=skip_decoder_start_token,

773 )

774

775 + async def _process_enc_dec_async(

776 + self,

777 + prompt: EncoderDecoderTokPrompt,

778 + ) -> EncoderDecoderInput:

779 + enc_prompt = prompt["encoder_prompt"]

780 + dec_prompt = prompt["decoder_prompt"]

781 +

782 + encoder_input, decoder_input = await asyncio.gather(

783 + self._process_singleton_async(enc_prompt),

784 + (

785 + asyncio.sleep(0)

786 + if dec_prompt is None

787 + else self._process_singleton_async(dec_prompt)

788 + ),

789 + )

790 +

791 + return build_enc_dec_input(

792 + encoder_input=encoder_input,

793 + decoder_input=decoder_input,

794 + decoder_start_token_id=self.get_dec_start_token_id(),

795 + )

796 +

797 def process_for_engine(self, prompt: TokPrompt, arrival_time: float) -> EngineInput:

798 engine_input: EngineInput

799 if "encoder_prompt" in prompt:

@@ -710,6 +805,21 @@ class BaseRenderer(ABC, Generic[_T]):

805

806 return engine_input

807

808 + async def process_for_engine_async(

809 + self, prompt: TokPrompt, arrival_time: float

810 + ) -> EngineInput:

811 + engine_input: EngineInput

812 + if "encoder_prompt" in prompt:

813 + engine_input = await self._process_enc_dec_async(

814 + prompt # type: ignore[arg-type]

815 + )

816 + else:

817 + engine_input = await self._process_singleton_async(prompt)

818 +

819 + engine_input["arrival_time"] = arrival_time

820 +

821 + return engine_input

822 +

823 # Top-level methods

824 def render_cmpl(

825 self,

@@ -747,7 +857,9 @@ class BaseRenderer(ABC, Generic[_T]):

857

858 self._apply_prompt_extras(tok_prompts, prompt_extras)

859

750 - return [self.process_for_engine(prompt, arrival_time) for prompt in tok_prompts]

860 + return await asyncio.gather(

861 + *(self.process_for_engine_async(p, arrival_time) for p in tok_prompts)

862 + )

863

864 def render_chat(

865 self,

@@ -811,8 +923,8 @@ class BaseRenderer(ABC, Generic[_T]):

923

924 self._apply_prompt_extras(tok_prompts, prompt_extras)

925

814 - eng_prompts = [

815 - self.process_for_engine(prompt, arrival_time) for prompt in tok_prompts

816 - ]

926 + eng_prompts = await asyncio.gather(

927 + *(self.process_for_engine_async(p, arrival_time) for p in tok_prompts)

928 + )

929

930 return out_conversations, eng_prompts

vllm/renderers/deepseek_v32.py

+18 -4

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 +from vllm.config import VllmConfig

5 from vllm.entrypoints.chat_utils import (

6 ChatCompletionMessageParam,

7 ConversationMessage,

@@ -9,6 +10,7 @@ from vllm.entrypoints.chat_utils import (

10 )

11 from vllm.logger import init_logger

12 from vllm.tokenizers.deepseek_v32 import DeepseekV32Tokenizer

13 +from vllm.utils.async_utils import make_async

15 from .base import BaseRenderer

16 from .inputs import DictPrompt

@@ -19,12 +21,25 @@ logger = init_logger(__name__)

23 class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]):

24 + def __init__(

25 + self,

26 + config: VllmConfig,

27 + tokenizer: DeepseekV32Tokenizer | None,

28 + ) -> None:

29 + super().__init__(config, tokenizer)

30 +

31 + self._apply_chat_template_async = make_async(

32 + self._apply_chat_template, executor=self._executor

33 + )

34 +

35 + def _apply_chat_template(self, *args, **kwargs):

36 + return self.get_tokenizer().apply_chat_template(*args, **kwargs)

37 +

38 def render_messages(

39 self,

40 messages: list[ChatCompletionMessageParam],

41 params: ChatParams,

42 ) -> tuple[list[ConversationMessage], DictPrompt]:

27 - tokenizer = self.get_tokenizer()

43 conversation, mm_data, mm_uuids = parse_chat_messages(

44 messages,

45 self.model_config,

@@ -33,7 +48,7 @@ class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]):

48 mm_processor_kwargs=params.mm_processor_kwargs,

49 )

36 - prompt_raw = tokenizer.apply_chat_template(

51 + prompt_raw = self._apply_chat_template(

52 conversation=conversation,

53 messages=messages,

54 **params.get_apply_chat_template_kwargs(),

@@ -52,7 +67,6 @@ class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]):

67 messages: list[ChatCompletionMessageParam],

68 params: ChatParams,

69 ) -> tuple[list[ConversationMessage], DictPrompt]:

55 - tokenizer = self.get_tokenizer()

70 conversation, mm_data, mm_uuids = await parse_chat_messages_async(

71 messages,

72 self.model_config,

@@ -61,7 +75,7 @@ class DeepseekV32Renderer(BaseRenderer[DeepseekV32Tokenizer]):

75 mm_processor_kwargs=params.mm_processor_kwargs,

76 )

64 - prompt_raw = tokenizer.apply_chat_template(

78 + prompt_raw = await self._apply_chat_template_async(

79 conversation=conversation,

80 messages=messages,

81 **params.get_apply_chat_template_kwargs(),

vllm/renderers/grok2.py

+18 -4

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 +from vllm.config import VllmConfig

5 from vllm.entrypoints.chat_utils import (

6 ChatCompletionMessageParam,

7 ConversationMessage,

@@ -9,6 +10,7 @@ from vllm.entrypoints.chat_utils import (

10 )

11 from vllm.logger import init_logger

12 from vllm.tokenizers.grok2 import Grok2Tokenizer

13 +from vllm.utils.async_utils import make_async

15 from .base import BaseRenderer

16 from .inputs import DictPrompt

@@ -19,12 +21,25 @@ logger = init_logger(__name__)

23 class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):

24 + def __init__(

25 + self,

26 + config: VllmConfig,

27 + tokenizer: Grok2Tokenizer | None,

28 + ) -> None:

29 + super().__init__(config, tokenizer)

30 +

31 + self._apply_chat_template_async = make_async(

32 + self._apply_chat_template, executor=self._executor

33 + )

34 +

35 + def _apply_chat_template(self, *args, **kwargs):

36 + return self.get_tokenizer().apply_chat_template(*args, **kwargs)

37 +

38 def render_messages(

39 self,

40 messages: list[ChatCompletionMessageParam],

41 params: ChatParams,

42 ) -> tuple[list[ConversationMessage], DictPrompt]:

27 - tokenizer = self.get_tokenizer()

43 conversation, mm_data, mm_uuids = parse_chat_messages(

44 messages,

45 self.model_config,

@@ -33,7 +48,7 @@ class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):

48 mm_processor_kwargs=params.mm_processor_kwargs,

49 )

36 - prompt_raw = tokenizer.apply_chat_template(

51 + prompt_raw = self._apply_chat_template(

52 conversation=conversation,

53 messages=messages,

54 **params.get_apply_chat_template_kwargs(),

@@ -52,7 +67,6 @@ class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):

67 messages: list[ChatCompletionMessageParam],

68 params: ChatParams,

69 ) -> tuple[list[ConversationMessage], DictPrompt]:

55 - tokenizer = self.get_tokenizer()

70 conversation, mm_data, mm_uuids = await parse_chat_messages_async(

71 messages,

72 self.model_config,

@@ -61,7 +75,7 @@ class Grok2Renderer(BaseRenderer[Grok2Tokenizer]):

75 mm_processor_kwargs=params.mm_processor_kwargs,

76 )

64 - prompt_raw = tokenizer.apply_chat_template(

78 + prompt_raw = await self._apply_chat_template_async(

79 conversation=conversation,

80 messages=messages,

81 **params.get_apply_chat_template_kwargs(),

vllm/renderers/hf.py

+20 -9

@@ -30,6 +30,7 @@ from vllm.logger import init_logger

30 from vllm.tokenizers.hf import HfTokenizer

31 from vllm.transformers_utils.chat_templates import get_chat_template_fallback_path

32 from vllm.transformers_utils.processor import cached_get_processor

33 +from vllm.utils.async_utils import make_async

34 from vllm.utils.func_utils import supports_kw

36 from .base import BaseRenderer

@@ -614,6 +615,10 @@ class HfRenderer(BaseRenderer[HfTokenizer]):

615 config.model_config.hf_config, "use_unified_vision_chunk", False

616 )

617

618 + self._apply_chat_template_async = make_async(

619 + safe_apply_chat_template, executor=self._executor

620 + )

621 +

622 def render_messages(

623 self,

624 messages: list[ChatCompletionMessageParam],

@@ -656,10 +661,13 @@ class HfRenderer(BaseRenderer[HfTokenizer]):

661 video_placeholder = getattr(

662 model_config.hf_config, "video_placeholder", None

663 )

659 - prompt_raw = replace_vision_chunk_video_placeholder(

660 - prompt_raw,

661 - mm_data,

662 - video_placeholder,

664 + prompt_raw = cast(

665 + list[int],

666 + replace_vision_chunk_video_placeholder(

667 + prompt_raw,

668 + mm_data,

669 + video_placeholder,

670 + ),

671 )

672

673 prompt = parse_dec_only_prompt(prompt_raw)

@@ -692,7 +700,7 @@ class HfRenderer(BaseRenderer[HfTokenizer]):

700 mm_processor_kwargs=params.mm_processor_kwargs,

701 )

702

695 - prompt_raw = safe_apply_chat_template(

703 + prompt_raw = await self._apply_chat_template_async(

704 model_config,

705 tokenizer,

706 conversation,

@@ -710,10 +718,13 @@ class HfRenderer(BaseRenderer[HfTokenizer]):

718 video_placeholder = getattr(

719 model_config.hf_config, "video_placeholder", None

720 )

713 - prompt_raw = replace_vision_chunk_video_placeholder(

714 - prompt_raw,

715 - mm_data,

716 - video_placeholder,

721 + prompt_raw = cast(

722 + list[int],

723 + replace_vision_chunk_video_placeholder(

724 + prompt_raw,

725 + mm_data,

726 + video_placeholder,

727 + ),

728 )

729

730 prompt = parse_dec_only_prompt(prompt_raw)

vllm/renderers/mistral.py

+1 -3

@@ -1,6 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from concurrent.futures import ThreadPoolExecutor

4 from vllm.config import VllmConfig

5 from vllm.entrypoints.chat_utils import (

@@ -56,9 +55,8 @@ class MistralRenderer(BaseRenderer[MistralTokenizer]):

55 ) -> None:

56 super().__init__(config, tokenizer)

59 - self._apply_chat_template_executor = ThreadPoolExecutor(max_workers=1)

58 self._apply_chat_template_async = make_async(

61 - safe_apply_chat_template, executor=self._apply_chat_template_executor

59 + safe_apply_chat_template, executor=self._executor

60 )

62 def render_messages(

vllm/utils/async_utils.py

+3 -1

@@ -34,6 +34,7 @@ class AsyncMicrobatchTokenizer:

34 tokenizer,

35 max_batch_size: int = 32,

36 batch_wait_timeout_s: float = 0.002,

37 + executor: ThreadPoolExecutor | None = None,

38 ) -> None:

39 self.tokenizer = tokenizer

40 self.max_batch_size = max_batch_size

@@ -47,7 +48,8 @@ class AsyncMicrobatchTokenizer:

48 self._batcher_tasks: list[Task] = []

50 # Single-thread executor for blocking tokenizer calls.

50 - self._executor = ThreadPoolExecutor(max_workers=1)

51 + # Accept an external executor to serialize with other tokenizer users.

52 + self._executor = executor or ThreadPoolExecutor(max_workers=1)

54 # === Public async API ===

55 async def __call__(self, prompt, **kwargs) -> BatchEncoding:

vllm/v1/engine/async_llm.py

+1 -1

@@ -889,7 +889,7 @@ class AsyncLLM(EngineClient):

889 await asyncio.gather(*coros)

890

891 async def reset_mm_cache(self) -> None:

892 - self.renderer.clear_mm_cache()

892 + await self.renderer.clear_mm_cache_async()

893 await self.engine_core.reset_mm_cache_async()

894

895 async def reset_prefix_cache(