@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 321331b8ae

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 15, 08:30 UTC

[Core] Add Lora Support to Beam Search (#18346)

Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com>

Alex Brooks committed May 28, 2025 at 09:58 UTC 321331b8ae41f13e519a63f99a0c427dc3907126

7 files changed +150 -16

tests/entrypoints/openai/test_lora_adapters.py

+34

@@ -313,3 +313,37 @@ async def test_loading_invalid_adapters_does_not_break_others(

313 prompt=["Hello there", "Foo bar bazz buzz"],

314 max_tokens=5,

315 )

316 +

317 +

318 +@pytest.mark.asyncio

319 +async def test_beam_search_with_lora_adapters(

320 + client: openai.AsyncOpenAI,

321 + tmp_path,

322 + zephyr_lora_files,

323 +):

324 + """Validate that async beam search can be used with lora."""

325 +

326 + async def load_and_run_adapter(adapter_name: str):

327 + await client.post("load_lora_adapter",

328 + cast_to=str,

329 + body={

330 + "lora_name": adapter_name,

331 + "lora_path": str(zephyr_lora_files)

332 + })

333 + for _ in range(3):

334 + await client.completions.create(

335 + model=adapter_name,

336 + prompt=["Hello there", "Foo bar bazz buzz"],

337 + max_tokens=5,

338 + extra_body=dict(use_beam_search=True),

339 + )

340 +

341 + lora_tasks = []

342 + for i in range(3):

343 + lora_tasks.append(

344 + asyncio.create_task(load_and_run_adapter(f"adapter_{i}")))

345 +

346 + results, _ = await asyncio.wait(lora_tasks)

347 +

348 + for r in results:

349 + assert not isinstance(r, Exception), f"Got exception {r}"

tests/lora/test_qwen2vl.py

+60 -2

@@ -10,6 +10,7 @@ import vllm

10 from vllm.assets.image import ImageAsset

11 from vllm.lora.request import LoRARequest

12 from vllm.platforms import current_platform

13 +from vllm.sampling_params import BeamSearchParams

16 @pytest.fixture(autouse=not current_platform.is_cpu())

@@ -69,7 +70,7 @@ class Qwen2VLTester:

70 expected_outputs: list[str],

71 lora_id: Optional[int] = None,

72 temperature: float = 0,

72 - max_tokens: int = 5) -> list[str]:

73 + max_tokens: int = 5):

75 sampling_params = vllm.SamplingParams(

76 temperature=temperature,

@@ -97,7 +98,35 @@ class Qwen2VLTester:

98 generated), f"Generated text {generated} doesn't "

99 f"match expected pattern {expected}"

100

100 - return generated_texts

101 + def run_beam_search_test(self,

102 + images: list[ImageAsset],

103 + expected_outputs: list[list[str]],

104 + lora_id: Optional[int] = None,

105 + temperature: float = 0,

106 + beam_width: int = 2,

107 + max_tokens: int = 5):

108 +

109 + beam_search_params = BeamSearchParams(beam_width=beam_width,

110 + max_tokens=max_tokens,

111 + temperature=temperature)

112 +

113 + inputs = [{

114 + "prompt": self.PROMPT_TEMPLATE,

115 + "multi_modal_data": {

116 + "image": asset.pil_image

117 + },

118 + } for asset in images]

119 +

120 + lora_request = LoRARequest(str(lora_id), lora_id,

121 + self.config.lora_path)

122 + outputs = self.llm.beam_search(inputs,

123 + beam_search_params,

124 + lora_request=lora_request)

125 +

126 + for output_obj, expected_outs in zip(outputs, expected_outputs):

127 + output_texts = [seq.text for seq in output_obj.sequences]

128 + assert output_texts == expected_outs, \

129 + f"Generated texts {output_texts} do not match expected {expected_outs}" # noqa: E501

130

131

132 TEST_IMAGES = [

@@ -110,6 +139,14 @@ EXPECTED_OUTPUTS = [

139 "A majestic skyscraper stands tall, partially obscured by a vibrant canopy of cherry blossoms, against a clear blue sky.", # noqa: E501

140 ]

141

142 +# NOTE - beam search .text contains the whole text

143 +EXPECTED_BEAM_SEARCH_OUTPUTS = [

144 + [

147 + ],

148 +]

149 +

150 QWEN2VL_MODEL_PATH = "Qwen/Qwen2-VL-2B-Instruct"

151 QWEN25VL_MODEL_PATH = "Qwen/Qwen2.5-VL-3B-Instruct"

152

@@ -130,6 +167,27 @@ def test_qwen2vl_lora(qwen2vl_lora_files):

167 lora_id=lora_id)

168

169

170 +@pytest.mark.xfail(

171 + current_platform.is_rocm(),

172 + reason="Qwen2-VL dependency xformers incompatible with ROCm")

173 +def test_qwen2vl_lora_beam_search(qwen2vl_lora_files):

174 + """Test Qwen 2.0 VL model with LoRA through beam search."""

175 + config = TestConfig(model_path=QWEN2VL_MODEL_PATH,

176 + lora_path=qwen2vl_lora_files)

177 + tester = Qwen2VLTester(config)

178 +

179 + # Test with different LoRA IDs

180 + for lora_id in [1, 2]:

181 + # NOTE currently, we only test cherry blossom since stop sign

182 + # output is slightly different for v1; - the root cause is likely

183 + # independent of the intent of this test, which is to ensure beam

184 + # search passes through lora through correctly.

185 + tester.run_beam_search_test(

186 + [ImageAsset("cherry_blossom")],

187 + expected_outputs=EXPECTED_BEAM_SEARCH_OUTPUTS,

188 + lora_id=lora_id)

189 +

190 +

191 @pytest.mark.xfail(

192 current_platform.is_rocm(),

193 reason="Qwen2.5-VL dependency xformers incompatible with ROCm",

vllm/beam_search.py

@@ -3,6 +3,7 @@

3 from dataclasses import dataclass

4 from typing import TYPE_CHECKING, Any, Optional, Union

6 +from vllm.lora.request import LoRARequest

7 from vllm.sequence import Logprob

9 if TYPE_CHECKING:

@@ -19,6 +20,7 @@ class BeamSearchSequence:

20 # The tokens includes the prompt.

21 tokens: list[int]

22 logprobs: list[dict[int, Logprob]]

23 + lora_request: Optional[LoRARequest] = None

24 cum_logprob: float = 0.0

25 text: Optional[str] = None

26 finish_reason: Optional[str] = None

@@ -41,6 +43,7 @@ class BeamSearchInstance:

43 def __init__(

44 self,

45 prompt_tokens: list[int],

46 + lora_request: Optional[LoRARequest] = None,

47 logprobs: Optional[list[dict[int, Logprob]]] = None,

48 **kwargs,

49 ):

@@ -48,6 +51,7 @@ class BeamSearchInstance:

51 BeamSearchSequence(

52 tokens=prompt_tokens,

53 logprobs=[] if logprobs is None else list(logprobs),

54 + lora_request=lora_request,

55 **kwargs,

56 )

57 ]

vllm/engine/protocol.py

+14 -8

@@ -65,6 +65,7 @@ class EngineClient(ABC):

65 prompt: PromptType,

66 request_id: str,

67 params: BeamSearchParams,

68 + lora_request: Optional[LoRARequest] = None,

69 ) -> AsyncGenerator[RequestOutput, None]:

71 beam_width = params.beam_width

@@ -106,27 +107,31 @@ class EngineClient(ABC):

107 cum_logprob=0,

108 logprobs=[],

109 multi_modal_data=multi_modal_data,

109 - mm_processor_kwargs=mm_processor_kwargs)

110 + mm_processor_kwargs=mm_processor_kwargs,

111 + lora_request=lora_request)

112 ]

113 completed = []

114

115 for _ in range(max_tokens):

114 - prompts_batch = [

116 + prompts_batch, lora_req_batch = zip(*[(

117 TokensPrompt(prompt_token_ids=beam.tokens,

118 multi_modal_data=beam.multi_modal_data,

117 - mm_processor_kwargs=beam.mm_processor_kwargs)

118 - for beam in all_beams

119 - ]

119 + mm_processor_kwargs=beam.mm_processor_kwargs),

120 + beam.lora_request,

121 + ) for beam in all_beams])

122

123 tasks = []

124

125 request_id = f"beam_search-{random_uuid()}"

124 - for i, individual_prompt in enumerate(prompts_batch):

126 + for i, (individual_prompt,

127 + lora_req) in enumerate(zip(prompts_batch, lora_req_batch)):

128 request_id_item = f"{request_id}-{i}"

129 task = asyncio.create_task(

130 collect_from_async_generator(

128 - self.generate(individual_prompt, beam_search_params,

129 - request_id_item)))

131 + self.generate(individual_prompt,

132 + beam_search_params,

133 + request_id_item,

134 + lora_request=lora_req)))

135 tasks.append(task)

136

137 output = await asyncio.gather(*tasks)

@@ -159,6 +164,7 @@ class EngineClient(ABC):

164 tokens=current_beam.tokens + [token_id],

165 logprobs=current_beam.logprobs +

166 [logprobs],

167 + lora_request=current_beam.lora_request,

168 cum_logprob=current_beam.cum_logprob +

169 logprob_obj.logprob,

170 multi_modal_data=current_beam.

vllm/entrypoints/llm.py

+36 -6

@@ -522,10 +522,28 @@ class LLM:

522 executor = self.llm_engine.model_executor

523 return executor.apply_model(func)

524

525 + def _get_beam_search_lora_requests(

526 + self,

527 + lora_request: Optional[Union[list[LoRARequest], LoRARequest]],

528 + prompts: list[Union[TokensPrompt, TextPrompt]],

529 + ) -> list[Optional[LoRARequest]]:

530 + """Get the optional lora request corresponding to each prompt."""

531 + if isinstance(lora_request,

532 + Sequence) and len(lora_request) != len(prompts):

533 + raise ValueError(

534 + "Lora request list should be the same length as the prompts")

535 + return lora_request

536 +

537 + if lora_request is None or isinstance(lora_request, LoRARequest):

538 + return [lora_request] * len(prompts)

539 +

540 + raise TypeError(f"Invalid lora_request type {type(lora_request)}")

541 +

542 def beam_search(

543 self,

544 prompts: list[Union[TokensPrompt, TextPrompt]],

545 params: BeamSearchParams,

546 + lora_request: Optional[Union[list[LoRARequest], LoRARequest]] = None,

547 ) -> list[BeamSearchOutput]:

548 """

549 Generate sequences using beam search.

@@ -534,6 +552,7 @@ class LLM:

552 prompts: A list of prompts. Each prompt can be a string or a list

553 of token IDs.

554 params: The beam search parameters.

555 + lora_request: LoRA request to use for generation, if any.

556 """

557 # TODO: how does beam search work together with length penalty,

558 # frequency, penalty, and stopping criteria, etc.?

@@ -543,6 +562,9 @@ class LLM:

562 ignore_eos = params.ignore_eos

563 length_penalty = params.length_penalty

564

565 + lora_requests = self._get_beam_search_lora_requests(

566 + lora_request, prompts)

567 +

568 def sort_beams_key(x: BeamSearchSequence) -> float:

569 return get_beam_search_score(x.tokens, x.cum_logprob,

570 tokenizer.eos_token_id,

@@ -570,7 +592,7 @@ class LLM:

592 temperature=temperature)

593 instances: list[BeamSearchInstance] = []

594

573 - for prompt in prompts:

595 + for lora_req, prompt in zip(lora_requests, prompts):

596 # Add multimodal processor kwargs & data

597 mm_kwargs = {}

598 if "multi_modal_data" in prompt:

@@ -586,7 +608,12 @@ class LLM:

608 prompt_tokens = tokenizer.encode(prompt["prompt"])

609

610 instances.append(

589 - BeamSearchInstance(prompt_tokens, logprobs=None, **mm_kwargs))

611 + BeamSearchInstance(

612 + prompt_tokens,

613 + lora_request=lora_req,

614 + logprobs=None,

615 + **mm_kwargs,

616 + ), )

617

618 for _ in range(max_tokens):

619 all_beams: list[BeamSearchSequence] = list(

@@ -600,15 +627,17 @@ class LLM:

627 if len(all_beams) == 0:

628 break

629

603 - prompts_batch = [

604 - create_tokens_prompt_from_beam(beam) for beam in all_beams

605 - ]

630 + # create the corresponding batch entries for prompt & optional lora

631 + prompts_batch, lora_req_batch = zip(

632 + *[(create_tokens_prompt_from_beam(beam), beam.lora_request)

633 + for beam in all_beams])

634

635 # only runs for one step

636 # we don't need to use tqdm here

637 output = self.generate(prompts_batch,

638 sampling_params=beam_search_params,

611 - use_tqdm=False)

639 + use_tqdm=False,

640 + lora_request=lora_req_batch)

641

642 for (start, end), instance in zip(instance_start_and_end,

643 instances):

@@ -626,6 +655,7 @@ class LLM:

655 new_beam = BeamSearchSequence(

656 tokens=current_beam.tokens + [token_id],

657 logprobs=current_beam.logprobs + [logprobs],

658 + lora_request=current_beam.lora_request,

659 cum_logprob=current_beam.cum_logprob +

660 logprob_obj.logprob,

661 multi_modal_data=current_beam.multi_modal_data,

vllm/entrypoints/openai/serving_chat.py

@@ -236,6 +236,7 @@ class OpenAIServingChat(OpenAIServing):

236 prompt=engine_prompt,

237 request_id=request_id,

238 params=sampling_params,

239 + lora_request=lora_request,

240 )

241 else:

242 generator = self.engine_client.generate(

vllm/entrypoints/openai/serving_completion.py

@@ -186,6 +186,7 @@ class OpenAIServingCompletion(OpenAIServing):

186 prompt=engine_prompt,

187 request_id=request_id,

188 params=sampling_params,

189 + lora_request=lora_request,

190 )

191 else:

192 generator = self.engine_client.generate(