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

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Model Runner V2][Bugfix] Fix MRV2 LoRA warmup (#35536)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>

Jee Jee Li committed Jun 16, 2026 at 04:17 UTC eacff17c8d574daea685387216b6bb23959ab2b1

5 files changed +227 -63

tests/lora/test_qwen3_with_multi_loras.py

+14 -4

@@ -6,6 +6,8 @@ This script contains:

6 2. test multi loras request

7 """

9 +import os

10 +

11 import pytest

13 from tests.utils import multi_gpu_test

@@ -39,6 +41,18 @@ def format_chatml_messages(

41 ]

44 +@pytest.fixture(autouse=True)

45 +def set_mrv2_env():

46 + original = os.environ.get("VLLM_USE_V2_MODEL_RUNNER", "0")

47 + os.environ["VLLM_USE_V2_MODEL_RUNNER"] = "1"

48 + yield

49 +

50 + if original is None:

51 + os.environ.pop("VLLM_USE_V2_MODEL_RUNNER", None)

52 + else:

53 + os.environ["VLLM_USE_V2_MODEL_RUNNER"] = original

54 +

55 +

56 def make_add_lora_request(name: str, path: str):

57 global INCREASE_LORA_ID, LORA_NAME_ID_MAP

@@ -61,7 +75,6 @@ def test_multi_loras_with_tp_sync():

75 max_lora_rank=LORA_RANK,

76 max_model_len=512,

77 gpu_memory_utilization=0.5,

64 - enforce_eager=True,

78 tensor_parallel_size=2, # ensure tp >= 2

79 max_cpu_loras=4, # ensure max_cpu_loras >= 2

80 )

@@ -167,7 +180,6 @@ def test_multiple_lora_requests():

180 max_lora_rank=LORA_RANK,

181 max_model_len=512,

182 gpu_memory_utilization=0.5,

170 - enforce_eager=True,

183 )

184 PROMPTS = ["Hello, my name is"] * 2

185 LORA_NAME = "Alice"

@@ -203,7 +215,6 @@ def test_load_inplace_offline_reload(

215 max_lora_rank=LORA_RANK,

216 max_model_len=512,

217 gpu_memory_utilization=0.5,

206 - enforce_eager=True,

218 )

219 adapter_id = 1

220 messages = format_chatml_messages(

@@ -254,7 +265,6 @@ def test_load_inplace_false_no_reload(

265 max_lora_rank=LORA_RANK,

266 max_model_len=512,

267 gpu_memory_utilization=0.5,

257 - enforce_eager=True,

268 )

269 adapter_id = 2

270 messages = format_chatml_messages(

vllm/v1/worker/gpu/cudagraph_utils.py

+92 -20

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

3 from collections import defaultdict

4 from collections.abc import Callable

5 from dataclasses import dataclass

6 +from itertools import product

7 from typing import Any, NamedTuple, Protocol

9 import torch

@@ -56,6 +57,7 @@ class BatchExecutionDescriptor:

57 num_tokens: int

58 num_reqs: int | None # None means no request padding is needed (PIECEWISE graphs)

59 uniform_token_count: int | None = None

60 + num_active_loras: int = 0

63 class CreateForwardFn(Protocol):

@@ -75,6 +77,7 @@ def _is_compatible(

77 num_reqs: int,

78 num_tokens: int,

79 uniform_token_count: int | None,

80 + num_active_loras: int,

81 ) -> bool:

82 # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count

83 # desc.num_reqs=None means no request padding needed (PIECEWISE)

@@ -85,6 +88,7 @@ def _is_compatible(

88 )

89 and (desc.num_reqs is None or desc.num_reqs >= num_reqs)

90 and desc.num_tokens >= num_tokens

91 + and desc.num_active_loras == num_active_loras

92 )

@@ -111,6 +115,7 @@ class CudaGraphManager:

115 device: torch.device,

116 cudagraph_mode: CUDAGraphMode,

117 decode_query_len: int,

118 + lora_capture_cases: list[int] | None = None,

119 ):

120 self.vllm_config = vllm_config

121 self.device = device

@@ -124,12 +129,17 @@ class CudaGraphManager:

129 self.tp_size = vllm_config.parallel_config.tensor_parallel_size

130 self.is_first_pp_rank = get_pp_group().is_first_rank

131 self.is_last_pp_rank = get_pp_group().is_last_rank

132 + self.lora_capture_cases = lora_capture_cases or [0]

133 + # Precompute actual num_active_loras -> captured case mapping so that

134 + # dispatch() is a plain dict lookup instead of a per-call bisect.

135 + self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map()

136

137 self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}

138 self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None

139

140 self._graphs_captured = False

132 - self._candidates: list[list[BatchExecutionDescriptor]] = []

141 +

142 + self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}

143 self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}

144 # adjust the cudagraph sizes to be a multiple of the uniform decode query length

145 self.compilation_config.adjust_cudagraph_sizes_for_spec_decode(

@@ -144,6 +154,32 @@ class CudaGraphManager:

154 )

155 self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None

156

157 + def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]:

158 + """Precompute actual num_active_loras -> effective captured case.

159 +

160 + Mirrors the num_tokens candidate expansion in ``_init_candidates``:

161 + every possible active-LoRA count is mapped ahead of time to the

162 + smallest captured case that can serve it, so ``dispatch`` is a plain

163 + dict lookup instead of a per-call bisect.

164 + """

165 + captured_with_lora = sorted(c for c in self.lora_capture_cases if c > 0)

166 + if not captured_with_lora:

167 + return {}, 0

168 + dispatch_map: dict[int, int] = {}

169 + case_idx = 0

170 + for n in range(1, captured_with_lora[-1] + 1):

171 + while captured_with_lora[case_idx] < n:

172 + case_idx += 1

173 + dispatch_map[n] = captured_with_lora[case_idx]

174 + return dispatch_map, captured_with_lora[-1]

175 +

176 + def _resolve_effective_loras(self, num_active_loras: int) -> int:

177 + """Map an actual active-LoRA count to its captured graph case."""

178 + if num_active_loras <= 0 or not self._lora_dispatch_map:

179 + return num_active_loras

180 + # Counts above the largest captured case clamp to it.

181 + return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case)

182 +

183 def _init_candidates(self) -> None:

184 """Build priority-ordered candidate lists for each token count."""

185 capture_sizes = self.compilation_config.cudagraph_capture_sizes

@@ -156,10 +192,14 @@ class CudaGraphManager:

192 mixed_mode = self.cudagraph_mode.mixed_mode()

193 separate_decode_routine = self.cudagraph_mode.separate_routine()

194

159 - descs_by_token_count = defaultdict(list)

195 + descs_by_token_lora: dict[tuple[int, int], list[BatchExecutionDescriptor]] = (

196 + defaultdict(list)

197 + )

198 descs_by_mode = defaultdict(list)

199

162 - for num_tokens in capture_sizes:

200 + for num_tokens, num_active_loras in product(

201 + capture_sizes, self.lora_capture_cases

202 + ):

203 # Capture uniform decode specfifc graphs if required

204 # (i.e. separate decode routine)

205 if (

@@ -172,9 +212,10 @@ class CudaGraphManager:

212 num_tokens=num_tokens,

213 num_reqs=num_tokens // self.decode_query_len,

214 uniform_token_count=self.decode_query_len,

215 + num_active_loras=num_active_loras,

216 )

217 descs_by_mode[decode_mode].append(desc)

177 - descs_by_token_count[num_tokens].append(desc)

218 + descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)

219

220 if mixed_mode:

221 # for PIECEWISE graphs there is no limit on requests when replaying

@@ -189,21 +230,25 @@ class CudaGraphManager:

230 cg_mode=mixed_mode,

231 num_tokens=num_tokens,

232 num_reqs=num_reqs,

233 + num_active_loras=num_active_loras,

234 )

235 descs_by_mode[mixed_mode].append(desc)

194 - descs_by_token_count[num_tokens].append(desc)

236 + descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)

237

196 - if not descs_by_token_count:

238 + if not descs_by_token_lora:

239 return

240

199 - sorted_padded = sorted(descs_by_token_count.keys())

200 - self._candidates = [[] for _ in range(sorted_padded[-1] + 1)]

201 -

241 + all_token_counts = sorted({k[0] for k in descs_by_token_lora})

242 current_range_start = 0

203 - for cg_size in sorted_padded:

204 - for i in range(current_range_start, cg_size + 1):

205 - self._candidates[i] = descs_by_token_count[cg_size]

206 - current_range_start = cg_size + 1

243 + for token_cg_size in all_token_counts:

244 + for i in range(current_range_start, token_cg_size + 1):

245 + for num_active_loras in self.lora_capture_cases:

246 + staging_key = (token_cg_size, num_active_loras)

247 + if staging_key in descs_by_token_lora:

248 + self._candidates[(i, num_active_loras)] = descs_by_token_lora[

249 + staging_key

250 + ]

251 + current_range_start = token_cg_size + 1

252

253 for mode, descs in descs_by_mode.items():

254 descs.sort(key=lambda d: d.num_tokens, reverse=True)

@@ -289,14 +334,27 @@ class CudaGraphManager:

334 num_reqs: int,

335 num_tokens: int,

336 uniform_token_count: int | None,

337 + num_active_loras: int,

338 ) -> BatchExecutionDescriptor:

339 """Find matching cudagraph descriptor from priority-ordered candidates."""

294 - if self._graphs_captured and 0 < num_tokens < len(self._candidates):

295 - for desc in self._candidates[num_tokens]:

296 - if _is_compatible(desc, num_reqs, num_tokens, uniform_token_count):

340 +

341 + effective_loras = self._resolve_effective_loras(num_active_loras)

342 + key = (num_tokens, effective_loras)

343 + if self._graphs_captured and num_tokens > 0 and key in self._candidates:

344 + for desc in self._candidates[key]:

345 + if _is_compatible(

346 + desc,

347 + num_reqs,

348 + num_tokens,

349 + uniform_token_count,

350 + effective_loras,

351 + ):

352 return desc

353 return BatchExecutionDescriptor(

299 - cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens, num_reqs=num_reqs

354 + cg_mode=CUDAGraphMode.NONE,

355 + num_tokens=num_tokens,

356 + num_reqs=num_reqs,

357 + num_active_loras=effective_loras,

358 )

359

360 def run_fullgraph(self, desc: BatchExecutionDescriptor):

@@ -337,9 +395,15 @@ class ModelCudaGraphManager(CudaGraphManager):

395 device: torch.device,

396 cudagraph_mode: CUDAGraphMode,

397 decode_query_len: int,

398 + lora_capture_cases: list[int] | None = None,

399 ):

341 - super().__init__(vllm_config, device, cudagraph_mode, decode_query_len)

342 - # Used for FULL CUDA graphs. PW CUDA graphs do not use these.

400 + super().__init__(

401 + vllm_config,

402 + device,

403 + cudagraph_mode,

404 + decode_query_len,

405 + lora_capture_cases=lora_capture_cases,

406 + )

407 self.hidden_states: torch.Tensor | None = None

408 self.aux_hidden_states: list[torch.Tensor] = []

409 self.use_aux_hidden_state_outputs = False

@@ -356,6 +420,7 @@ class ModelCudaGraphManager(CudaGraphManager):

420 kv_cache_config: KVCacheConfig,

421 has_lora: bool = False,

422 use_aux_hidden_state_outputs: bool = False,

423 + lora_capture_hook: Callable[[int, int, int], None] | None = None,

424 progress_bar_desc: str = "Capturing CUDA graphs",

425 ) -> dict[BatchExecutionDescriptor, AttentionStatePair]:

426 """Capture CUDA graphs for model forward pass."""

@@ -372,6 +437,11 @@ class ModelCudaGraphManager(CudaGraphManager):

437 ]:

438 num_tokens = desc.num_tokens

439 num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)

440 +

441 + # Set LoRA state before capture so kernels see correct adapters.

442 + if lora_capture_hook is not None:

443 + lora_capture_hook(desc.num_active_loras, num_reqs, num_tokens)

444 +

445 num_tokens_across_dp = (

446 torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")

447 if self.dp_size > 1

@@ -406,7 +476,9 @@ class ModelCudaGraphManager(CudaGraphManager):

476 if cg_mode == CUDAGraphMode.PIECEWISE:

477 assert attn_metadata is None

478 batch_descriptor = BatchDescriptor(

409 - num_tokens=num_tokens, has_lora=has_lora

479 + num_tokens=num_tokens,

480 + has_lora=has_lora,

481 + num_active_loras=desc.num_active_loras,

482 )

483 with set_forward_context(

484 attn_metadata,

vllm/v1/worker/gpu/dp_utils.py

+15 -3

@@ -21,6 +21,7 @@ def sync_cudagraph_and_dp_padding(

21 uniform_token_count: int | None,

22 dp_size: int,

23 dp_rank: int,

24 + num_active_loras: int = 0,

25 ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:

26 """

27 Coordinates the batch descriptor and DP padding across all ranks.

@@ -53,6 +54,7 @@ def sync_cudagraph_and_dp_padding(

54 cg_mode=CUDAGraphMode.NONE,

55 num_tokens=num_tokens,

56 num_reqs=num_reqs,

57 + num_active_loras=desired_batch_desc.num_active_loras,

58 ), num_tokens_across_dp

60 assert cudagraph_manager is not None, (

@@ -68,9 +70,13 @@ def sync_cudagraph_and_dp_padding(

70 synced_uniform_token_count = None

72 # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs

71 - # so we don't perform request padding for PIECEWISE graphs

73 + # so we don't perform request padding for PIECEWISE graphs.

74 + # num_active_loras is per-rank and doesn't need cross-rank agreement.

75 synced_desc = cudagraph_manager.dispatch(

73 - num_reqs, synced_num_tokens, synced_uniform_token_count

76 + num_reqs,

77 + synced_num_tokens,

78 + synced_uniform_token_count,

79 + num_active_loras=num_active_loras,

80 )

82 # Update num_tokens_across_dp to reflect padded size.

@@ -87,12 +93,14 @@ def dispatch_cg_and_sync_dp(

93 dp_size: int,

94 dp_rank: int,

95 need_eager: bool = False,

96 + num_active_loras: int = 0,

97 ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:

98 if need_eager:

99 batch_desc = BatchExecutionDescriptor(

100 cg_mode=CUDAGraphMode.NONE,

101 num_tokens=num_tokens,

102 num_reqs=num_reqs,

103 + num_active_loras=num_active_loras,

104 )

105 else:

106 assert cudagraph_manager is not None, (

@@ -100,7 +108,10 @@ def dispatch_cg_and_sync_dp(

108 "where need_eager must be True"

109 )

110 batch_desc = cudagraph_manager.dispatch(

103 - num_reqs, num_tokens, uniform_token_count

111 + num_reqs,

112 + num_tokens,

113 + uniform_token_count,

114 + num_active_loras=num_active_loras,

115 )

116

117 if dp_size == 1:

@@ -114,4 +125,5 @@ def dispatch_cg_and_sync_dp(

125 uniform_token_count,

126 dp_size,

127 dp_rank,

128 + num_active_loras=num_active_loras,

129 )

vllm/v1/worker/gpu/lora_utils.py

+66 -1

@@ -1,12 +1,74 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +"""LoRA utilities for the Model Runner V2 and cudagraph."""

4 +

5 +from collections.abc import Callable

6 +from typing import TYPE_CHECKING, Any

7 +

8 import numpy as np

10 from vllm.lora.request import LoRARequest

11 +from vllm.lora.utils import get_captured_lora_counts

12 +

13 +if TYPE_CHECKING:

14 + from vllm.config.compilation import CompilationConfig

15 + from vllm.config.lora import LoRAConfig

17 NO_LORA_ID = 0

20 +def get_lora_capture_cases(

21 + lora_config: "LoRAConfig | None",

22 + compilation_config: "CompilationConfig",

23 +) -> list[int]:

24 + """

25 + Return num_active_loras values for cudagraph capture.

26 +

27 + When cudagraph_specialize_lora=True: powers of 2 up to max_loras, plus

28 + max_loras+1. When False: [0, max_loras+1]. When LoRA disabled: [0].

29 + """

30 + if lora_config is None:

31 + return [0]

32 + if compilation_config.cudagraph_specialize_lora:

33 + specialize = getattr(lora_config, "specialize_active_lora", False)

34 + captured = get_captured_lora_counts(lora_config.max_loras, specialize)

35 + return [0] + [c for c in captured if c > 0]

36 + return [0, lora_config.max_loras + 1]

37 +

38 +

39 +def get_num_active_loras_for_dispatch(

40 + lora_config: "LoRAConfig | None",

41 + lora_state: "LoraState",

42 + req_ids: list[str],

43 + dummy_run: bool,

44 +) -> int:

45 + """Compute num_active_loras for cudagraph dispatch."""

46 + if lora_config and not dummy_run:

47 + return len(lora_state.get_activate_loras(req_ids))

48 + if dummy_run and lora_config:

49 + return lora_config.max_loras + 1

50 + return 0

51 +

52 +

53 +def create_lora_capture_hook(

54 + lora_config: "LoRAConfig | None",

55 + runner: Any,

56 +) -> Callable[[int, int, int], None] | None:

57 + """Create a hook to set up LoRA state before each cudagraph capture."""

58 + if lora_config is None:

59 + return None

60 +

61 + def hook(num_active_loras: int, num_reqs: int, num_tokens: int) -> None:

62 + num_scheduled = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)

63 + num_scheduled[-1] += num_tokens % num_reqs

64 + with runner.maybe_select_dummy_loras(

65 + lora_config, num_scheduled, num_active_loras=num_active_loras

66 + ):

67 + pass

68 +

69 + return hook

70 +

71 +

72 class LoraState:

73 def __init__(self, max_num_reqs: int):

74 self.lora_ids = np.zeros(max_num_reqs, dtype=np.int32)

@@ -35,10 +97,13 @@ class LoraState:

97 lora_ids = self.lora_ids[idx_mapping]

98 prompt_lora_mapping = tuple(lora_ids)

99 token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens))

100 + active_lora_requests: set[LoRARequest] = self.get_activate_loras(req_ids)

101 + return prompt_lora_mapping, token_lora_mapping, active_lora_requests

102

103 + def get_activate_loras(self, req_ids: list[str]) -> set[LoRARequest]:

104 active_lora_requests: set[LoRARequest] = set()

105 for req_id in req_ids:

106 lora_request = self.lora_requests.get(req_id)

107 if lora_request is not None:

108 active_lora_requests.add(lora_request)

44 - return prompt_lora_mapping, token_lora_mapping, active_lora_requests

109 + return active_lora_requests

vllm/v1/worker/gpu/model_runner.py

+40 -35

@@ -37,7 +37,6 @@ from vllm.distributed.parallel_state import (

37 )

38 from vllm.forward_context import BatchDescriptor, set_forward_context

39 from vllm.logger import init_logger

40 -from vllm.lora.layers import LoRAMapping

40 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (

41 initialize_mamba_ssu_backend,

42 )

@@ -88,7 +87,12 @@ from vllm.v1.worker.gpu.kv_connector import (

87 KVConnector,

88 get_kv_connector,

89 )

91 -from vllm.v1.worker.gpu.lora_utils import LoraState

90 +from vllm.v1.worker.gpu.lora_utils import (

91 + LoraState,

92 + create_lora_capture_hook,

93 + get_lora_capture_cases,

94 + get_num_active_loras_for_dispatch,

95 +)

96 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache

97 from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras

98 from vllm.v1.worker.gpu.model_states import init_model_state

@@ -234,8 +238,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):

238 self.prompt_logprobs_worker: PromptLogprobsWorker | None = None

239 self.structured_outputs_worker: StructuredOutputsWorker | None = None

240 self.cudagraph_manager: ModelCudaGraphManager | None = None

241 +

242 # LoRA-related workers.

243 self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)

244 + self.lora_capture_cases = [0]

245 + if self.lora_config:

246 + self.lora_capture_cases = get_lora_capture_cases(

247 + self.lora_config, self.compilation_config

248 + )

249 +

250 # KV Connector if configured.

251 self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR

252

@@ -458,6 +469,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

469 self.device,

470 cudagraph_mode,

471 decode_query_len=self.decode_query_len,

472 + lora_capture_cases=self.lora_capture_cases,

473 )

474 if self.speculator is not None:

475 self.speculator.init_cudagraph_manager(cudagraph_mode)

@@ -540,14 +552,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):

552 assert self.intermediate_tensors is not None

553 intermediate_tensors = self.intermediate_tensors[:num_tokens]

554

543 - # Execute the model.

544 - self.execute_model(

545 - dummy_scheduler_output,

546 - intermediate_tensors=intermediate_tensors,

547 - dummy_run=True,

548 - skip_attn_for_dummy_run=skip_attn,

549 - is_profile=is_profile,

550 - )

555 + max_loras = self.lora_config.max_loras if self.lora_config is not None else 0

556 + with self.maybe_dummy_run_with_lora(

557 + self.lora_config,

558 + num_scheduled_tokens=np.array(num_tokens_per_request, dtype=np.int32),

559 + num_sampled_tokens=None,

560 + remove_lora=True,

561 + num_active_loras=max_loras,

562 + ):

563 + # Execute the model.

564 + self.execute_model(

565 + dummy_scheduler_output,

566 + intermediate_tensors=intermediate_tensors,

567 + dummy_run=True,

568 + skip_attn_for_dummy_run=skip_attn,

569 + is_profile=is_profile,

570 + )

571 self.kv_connector.set_disabled(False)

572

573 # Non-last PP ranks don't produce output for sampling.

@@ -694,6 +714,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

714 self.kv_cache_config,

715 has_lora=self.lora_config is not None,

716 use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,

717 + lora_capture_hook=create_lora_capture_hook(self.lora_config, self),

718 )

719 if self.speculator is not None:

720 self.speculator.capture(attn_states)

@@ -1105,6 +1126,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1126 max_query_len = max(scheduler_output.num_scheduled_tokens.values())

1127 uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len)

1128

1129 + num_active_loras = 0

1130 + if self.lora_config:

1131 + req_ids = list(scheduler_output.num_scheduled_tokens.keys())

1132 + num_active_loras = get_num_active_loras_for_dispatch(

1133 + self.lora_config, self.lora_state, req_ids, dummy_run

1134 + )

1135 +

1136 skip_compiled = False

1137 if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:

1138 # Encoder-decoder models such as Whisper should run eager/non-compiled

@@ -1120,6 +1148,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1148 self.dp_size,

1149 self.dp_rank,

1150 need_eager=is_profile or skip_compiled,

1151 + num_active_loras=num_active_loras,

1152 )

1153

1154 if batch_desc.num_tokens == 0:

@@ -1157,31 +1186,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1186 )

1187 block_tables = None

1188 slot_mappings = None

1160 - if self.lora_config:

1161 - # program a no-LoRA mapping here so kernels early-exit instead of

1162 - # reading uninitialized metadata during dummy runs.

1163 - # FIXME: Replace this with LoRA warmup:

1164 - # https://github.com/vllm-project/vllm/pull/35536

1165 - assert hasattr(self, "lora_manager")

1166 - adapter_manager = self.lora_manager._adapter_manager

1167 - adapter_manager.set_adapter_mapping(

1168 - LoRAMapping(

1169 - index_mapping=(0,) * input_batch.num_tokens_after_padding,

1170 - prompt_mapping=(0,) * input_batch.num_reqs,

1171 - is_prefill=True,

1172 - )

1173 - )

1174 - seen_wrappers: set[int] = set()

1175 - for punica_wrapper in adapter_manager.punica_wrapper_mapping.values():

1176 - if id(punica_wrapper) in seen_wrappers:

1177 - continue

1178 - seen_wrappers.add(id(punica_wrapper))

1179 - for kernel_meta in (

1180 - punica_wrapper.token_mapping_meta, # type: ignore[attr-defined]

1181 - punica_wrapper.prompt_mapping_meta, # type: ignore[attr-defined]

1182 - ):

1183 - kernel_meta.no_lora_flag_cpu[0] = False

1184 - kernel_meta.num_active_loras_cpu[0] = 1

1189

1190 attn_metadata = None

1191 slot_mappings_by_layer = None

@@ -1258,6 +1262,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1262 batch_descriptor = BatchDescriptor(

1263 num_tokens=input_batch.num_tokens_after_padding,

1264 has_lora=self.lora_config is not None,

1265 + num_active_loras=batch_desc.num_active_loras,

1266 )

1267

1268 with set_forward_context(