offload prompt_embeds decode in render_prompts_async to avoid blocking (#43792)
Signed-off-by: Gagan Dhakrey <gagandhakrey@gmail.com>
This commit is contained in:
+21
-1
@@ -104,6 +104,9 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
self._process_multimodal_async = make_async(
|
||||
self._process_multimodal, executor=self._mm_executor
|
||||
)
|
||||
self._safe_load_prompt_embeds_async = make_async(
|
||||
safe_load_prompt_embeds, executor=self._executor
|
||||
)
|
||||
if mm_registry.supports_multimodal_inputs(config.model_config):
|
||||
mm_processor_cache = mm_registry.processor_cache_from_config(config)
|
||||
|
||||
@@ -376,11 +379,28 @@ class BaseRenderer(ABC, Generic[_T]):
|
||||
|
||||
return [self.render_prompt(prompt) for prompt in prompts]
|
||||
|
||||
async def _render_prompt_async(
|
||||
self,
|
||||
prompt: DictPrompt | bytes,
|
||||
) -> DictPrompt:
|
||||
if isinstance(prompt, bytes):
|
||||
embeds = await self._safe_load_prompt_embeds_async(
|
||||
self.model_config, prompt
|
||||
)
|
||||
return EmbedsPrompt(prompt_embeds=embeds)
|
||||
|
||||
return prompt
|
||||
|
||||
async def render_prompts_async(
|
||||
self,
|
||||
prompts: Sequence[DictPrompt | bytes],
|
||||
) -> list[DictPrompt]:
|
||||
return self.render_prompts(prompts)
|
||||
if len(prompts) == 0:
|
||||
raise ValueError("You must pass at least one prompt")
|
||||
|
||||
return await asyncio.gather(
|
||||
*(self._render_prompt_async(prompt) for prompt in prompts)
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
def render_messages(
|
||||
|
||||
Reference in New Issue
Block a user