[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>

tests/lora/test_qwen3_with_multi_loras.py

@@ -6,6 +6,8 @@ This script contains:
 2. test multi loras request
 """
 
+import os
+
 import pytest
 
 from tests.utils import multi_gpu_test
@@ -39,6 +41,18 @@ def format_chatml_messages(
     ]
 
 
+@pytest.fixture(autouse=True)
+def set_mrv2_env():
+    original = os.environ.get("VLLM_USE_V2_MODEL_RUNNER", "0")
+    os.environ["VLLM_USE_V2_MODEL_RUNNER"] = "1"
+    yield
+
+    if original is None:
+        os.environ.pop("VLLM_USE_V2_MODEL_RUNNER", None)
+    else:
+        os.environ["VLLM_USE_V2_MODEL_RUNNER"] = original
+
+
 def make_add_lora_request(name: str, path: str):
     global INCREASE_LORA_ID, LORA_NAME_ID_MAP
 
@@ -61,7 +75,6 @@ def test_multi_loras_with_tp_sync():
         max_lora_rank=LORA_RANK,
         max_model_len=512,
         gpu_memory_utilization=0.5,
-        enforce_eager=True,
         tensor_parallel_size=2,  # ensure tp >= 2
         max_cpu_loras=4,  # ensure max_cpu_loras >= 2
     )
@@ -167,7 +180,6 @@ def test_multiple_lora_requests():
         max_lora_rank=LORA_RANK,
         max_model_len=512,
         gpu_memory_utilization=0.5,
-        enforce_eager=True,
     )
     PROMPTS = ["Hello, my name is"] * 2
     LORA_NAME = "Alice"
@@ -203,7 +215,6 @@ def test_load_inplace_offline_reload(
         max_lora_rank=LORA_RANK,
         max_model_len=512,
         gpu_memory_utilization=0.5,
-        enforce_eager=True,
     )
     adapter_id = 1
     messages = format_chatml_messages(
@@ -254,7 +265,6 @@ def test_load_inplace_false_no_reload(
         max_lora_rank=LORA_RANK,
         max_model_len=512,
         gpu_memory_utilization=0.5,
-        enforce_eager=True,
     )
     adapter_id = 2
     messages = format_chatml_messages(

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -3,6 +3,7 @@
 from collections import defaultdict
 from collections.abc import Callable
 from dataclasses import dataclass
+from itertools import product
 from typing import Any, NamedTuple, Protocol
 
 import torch
@@ -56,6 +57,7 @@ class BatchExecutionDescriptor:
     num_tokens: int
     num_reqs: int | None  # None means no request padding is needed (PIECEWISE graphs)
     uniform_token_count: int | None = None
+    num_active_loras: int = 0
 
 
 class CreateForwardFn(Protocol):
@@ -75,6 +77,7 @@ def _is_compatible(
     num_reqs: int,
     num_tokens: int,
     uniform_token_count: int | None,
+    num_active_loras: int,
 ) -> bool:
     # desc.uniform_token_count=None (PIECEWISE) can handle any uniform_token_count
     # desc.num_reqs=None means no request padding needed (PIECEWISE)
@@ -85,6 +88,7 @@ def _is_compatible(
         )
         and (desc.num_reqs is None or desc.num_reqs >= num_reqs)
         and desc.num_tokens >= num_tokens
+        and desc.num_active_loras == num_active_loras
     )
 
 
@@ -111,6 +115,7 @@ class CudaGraphManager:
         device: torch.device,
         cudagraph_mode: CUDAGraphMode,
         decode_query_len: int,
+        lora_capture_cases: list[int] | None = None,
     ):
         self.vllm_config = vllm_config
         self.device = device
@@ -124,12 +129,17 @@ class CudaGraphManager:
         self.tp_size = vllm_config.parallel_config.tensor_parallel_size
         self.is_first_pp_rank = get_pp_group().is_first_rank
         self.is_last_pp_rank = get_pp_group().is_last_rank
+        self.lora_capture_cases = lora_capture_cases or [0]
+        # Precompute actual num_active_loras -> captured case mapping so that
+        # dispatch() is a plain dict lookup instead of a per-call bisect.
+        self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map()
 
         self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}
         self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None
 
         self._graphs_captured = False
-        self._candidates: list[list[BatchExecutionDescriptor]] = []
+
+        self._candidates: dict[tuple[int, int], list[BatchExecutionDescriptor]] = {}
         self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}
         # adjust the cudagraph sizes to be a multiple of the uniform decode query length
         self.compilation_config.adjust_cudagraph_sizes_for_spec_decode(
@@ -144,6 +154,32 @@ class CudaGraphManager:
         )
         self.breakable_cg_runner: BreakableCUDAGraphWrapper | None = None
 
+    def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]:
+        """Precompute actual num_active_loras -> effective captured case.
+
+        Mirrors the num_tokens candidate expansion in ``_init_candidates``:
+        every possible active-LoRA count is mapped ahead of time to the
+        smallest captured case that can serve it, so ``dispatch`` is a plain
+        dict lookup instead of a per-call bisect.
+        """
+        captured_with_lora = sorted(c for c in self.lora_capture_cases if c > 0)
+        if not captured_with_lora:
+            return {}, 0
+        dispatch_map: dict[int, int] = {}
+        case_idx = 0
+        for n in range(1, captured_with_lora[-1] + 1):
+            while captured_with_lora[case_idx] < n:
+                case_idx += 1
+            dispatch_map[n] = captured_with_lora[case_idx]
+        return dispatch_map, captured_with_lora[-1]
+
+    def _resolve_effective_loras(self, num_active_loras: int) -> int:
+        """Map an actual active-LoRA count to its captured graph case."""
+        if num_active_loras <= 0 or not self._lora_dispatch_map:
+            return num_active_loras
+        # Counts above the largest captured case clamp to it.
+        return self._lora_dispatch_map.get(num_active_loras, self._max_lora_case)
+
     def _init_candidates(self) -> None:
         """Build priority-ordered candidate lists for each token count."""
         capture_sizes = self.compilation_config.cudagraph_capture_sizes
@@ -156,10 +192,14 @@ class CudaGraphManager:
         mixed_mode = self.cudagraph_mode.mixed_mode()
         separate_decode_routine = self.cudagraph_mode.separate_routine()
 
-        descs_by_token_count = defaultdict(list)
+        descs_by_token_lora: dict[tuple[int, int], list[BatchExecutionDescriptor]] = (
+            defaultdict(list)
+        )
         descs_by_mode = defaultdict(list)
 
-        for num_tokens in capture_sizes:
+        for num_tokens, num_active_loras in product(
+            capture_sizes, self.lora_capture_cases
+        ):
             # Capture uniform decode specfifc graphs if required
             #  (i.e. separate decode routine)
             if (
@@ -172,9 +212,10 @@ class CudaGraphManager:
                     num_tokens=num_tokens,
                     num_reqs=num_tokens // self.decode_query_len,
                     uniform_token_count=self.decode_query_len,
+                    num_active_loras=num_active_loras,
                 )
                 descs_by_mode[decode_mode].append(desc)
-                descs_by_token_count[num_tokens].append(desc)
+                descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)
 
             if mixed_mode:
                 # for PIECEWISE graphs there is no limit on requests when replaying
@@ -189,21 +230,25 @@ class CudaGraphManager:
                     cg_mode=mixed_mode,
                     num_tokens=num_tokens,
                     num_reqs=num_reqs,
+                    num_active_loras=num_active_loras,
                 )
                 descs_by_mode[mixed_mode].append(desc)
-                descs_by_token_count[num_tokens].append(desc)
+                descs_by_token_lora[(num_tokens, num_active_loras)].append(desc)
 
-        if not descs_by_token_count:
+        if not descs_by_token_lora:
             return
 
-        sorted_padded = sorted(descs_by_token_count.keys())
-        self._candidates = [[] for _ in range(sorted_padded[-1] + 1)]
-
+        all_token_counts = sorted({k[0] for k in descs_by_token_lora})
         current_range_start = 0
-        for cg_size in sorted_padded:
-            for i in range(current_range_start, cg_size + 1):
-                self._candidates[i] = descs_by_token_count[cg_size]
-            current_range_start = cg_size + 1
+        for token_cg_size in all_token_counts:
+            for i in range(current_range_start, token_cg_size + 1):
+                for num_active_loras in self.lora_capture_cases:
+                    staging_key = (token_cg_size, num_active_loras)
+                    if staging_key in descs_by_token_lora:
+                        self._candidates[(i, num_active_loras)] = descs_by_token_lora[
+                            staging_key
+                        ]
+            current_range_start = token_cg_size + 1
 
         for mode, descs in descs_by_mode.items():
             descs.sort(key=lambda d: d.num_tokens, reverse=True)
@@ -289,14 +334,27 @@ class CudaGraphManager:
         num_reqs: int,
         num_tokens: int,
         uniform_token_count: int | None,
+        num_active_loras: int,
     ) -> BatchExecutionDescriptor:
         """Find matching cudagraph descriptor from priority-ordered candidates."""
-        if self._graphs_captured and 0 < num_tokens < len(self._candidates):
-            for desc in self._candidates[num_tokens]:
-                if _is_compatible(desc, num_reqs, num_tokens, uniform_token_count):
+
+        effective_loras = self._resolve_effective_loras(num_active_loras)
+        key = (num_tokens, effective_loras)
+        if self._graphs_captured and num_tokens > 0 and key in self._candidates:
+            for desc in self._candidates[key]:
+                if _is_compatible(
+                    desc,
+                    num_reqs,
+                    num_tokens,
+                    uniform_token_count,
+                    effective_loras,
+                ):
                     return desc
         return BatchExecutionDescriptor(
-            cg_mode=CUDAGraphMode.NONE, num_tokens=num_tokens, num_reqs=num_reqs
+            cg_mode=CUDAGraphMode.NONE,
+            num_tokens=num_tokens,
+            num_reqs=num_reqs,
+            num_active_loras=effective_loras,
         )
 
     def run_fullgraph(self, desc: BatchExecutionDescriptor):
@@ -337,9 +395,15 @@ class ModelCudaGraphManager(CudaGraphManager):
         device: torch.device,
         cudagraph_mode: CUDAGraphMode,
         decode_query_len: int,
+        lora_capture_cases: list[int] | None = None,
     ):
-        super().__init__(vllm_config, device, cudagraph_mode, decode_query_len)
-        # Used for FULL CUDA graphs. PW CUDA graphs do not use these.
+        super().__init__(
+            vllm_config,
+            device,
+            cudagraph_mode,
+            decode_query_len,
+            lora_capture_cases=lora_capture_cases,
+        )
         self.hidden_states: torch.Tensor | None = None
         self.aux_hidden_states: list[torch.Tensor] = []
         self.use_aux_hidden_state_outputs = False
@@ -356,6 +420,7 @@ class ModelCudaGraphManager(CudaGraphManager):
         kv_cache_config: KVCacheConfig,
         has_lora: bool = False,
         use_aux_hidden_state_outputs: bool = False,
+        lora_capture_hook: Callable[[int, int, int], None] | None = None,
         progress_bar_desc: str = "Capturing CUDA graphs",
     ) -> dict[BatchExecutionDescriptor, AttentionStatePair]:
         """Capture CUDA graphs for model forward pass."""
@@ -372,6 +437,11 @@ class ModelCudaGraphManager(CudaGraphManager):
         ]:
             num_tokens = desc.num_tokens
             num_reqs = desc.num_reqs or min(num_tokens, self.max_num_reqs)
+
+            # Set LoRA state before capture so kernels see correct adapters.
+            if lora_capture_hook is not None:
+                lora_capture_hook(desc.num_active_loras, num_reqs, num_tokens)
+
             num_tokens_across_dp = (
                 torch.full((self.dp_size,), num_tokens, dtype=torch.int32, device="cpu")
                 if self.dp_size > 1
@@ -406,7 +476,9 @@ class ModelCudaGraphManager(CudaGraphManager):
                 if cg_mode == CUDAGraphMode.PIECEWISE:
                     assert attn_metadata is None
                     batch_descriptor = BatchDescriptor(
-                        num_tokens=num_tokens, has_lora=has_lora
+                        num_tokens=num_tokens,
+                        has_lora=has_lora,
+                        num_active_loras=desc.num_active_loras,
                     )
                 with set_forward_context(
                     attn_metadata,

vllm/v1/worker/gpu/dp_utils.py

@@ -21,6 +21,7 @@ def sync_cudagraph_and_dp_padding(
     uniform_token_count: int | None,
     dp_size: int,
     dp_rank: int,
+    num_active_loras: int = 0,
 ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:
     """
     Coordinates the batch descriptor and DP padding across all ranks.
@@ -53,6 +54,7 @@ def sync_cudagraph_and_dp_padding(
             cg_mode=CUDAGraphMode.NONE,
             num_tokens=num_tokens,
             num_reqs=num_reqs,
+            num_active_loras=desired_batch_desc.num_active_loras,
         ), num_tokens_across_dp
 
     assert cudagraph_manager is not None, (
@@ -68,9 +70,13 @@ def sync_cudagraph_and_dp_padding(
         synced_uniform_token_count = None
 
     # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs
-    # so we don't perform request padding for PIECEWISE graphs
+    # so we don't perform request padding for PIECEWISE graphs.
+    # num_active_loras is per-rank and doesn't need cross-rank agreement.
     synced_desc = cudagraph_manager.dispatch(
-        num_reqs, synced_num_tokens, synced_uniform_token_count
+        num_reqs,
+        synced_num_tokens,
+        synced_uniform_token_count,
+        num_active_loras=num_active_loras,
     )
 
     # Update num_tokens_across_dp to reflect padded size.
@@ -87,12 +93,14 @@ def dispatch_cg_and_sync_dp(
     dp_size: int,
     dp_rank: int,
     need_eager: bool = False,
+    num_active_loras: int = 0,
 ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]:
     if need_eager:
         batch_desc = BatchExecutionDescriptor(
             cg_mode=CUDAGraphMode.NONE,
             num_tokens=num_tokens,
             num_reqs=num_reqs,
+            num_active_loras=num_active_loras,
         )
     else:
         assert cudagraph_manager is not None, (
@@ -100,7 +108,10 @@ def dispatch_cg_and_sync_dp(
             "where need_eager must be True"
         )
         batch_desc = cudagraph_manager.dispatch(
-            num_reqs, num_tokens, uniform_token_count
+            num_reqs,
+            num_tokens,
+            uniform_token_count,
+            num_active_loras=num_active_loras,
         )
 
     if dp_size == 1:
@@ -114,4 +125,5 @@ def dispatch_cg_and_sync_dp(
         uniform_token_count,
         dp_size,
         dp_rank,
+        num_active_loras=num_active_loras,
     )

vllm/v1/worker/gpu/lora_utils.py

@@ -1,12 +1,74 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""LoRA utilities for the Model Runner V2 and cudagraph."""
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
 import numpy as np
 
 from vllm.lora.request import LoRARequest
+from vllm.lora.utils import get_captured_lora_counts
+
+if TYPE_CHECKING:
+    from vllm.config.compilation import CompilationConfig
+    from vllm.config.lora import LoRAConfig
 
 NO_LORA_ID = 0
 
 
+def get_lora_capture_cases(
+    lora_config: "LoRAConfig | None",
+    compilation_config: "CompilationConfig",
+) -> list[int]:
+    """
+    Return num_active_loras values for cudagraph capture.
+
+    When cudagraph_specialize_lora=True: powers of 2 up to max_loras, plus
+    max_loras+1. When False: [0, max_loras+1]. When LoRA disabled: [0].
+    """
+    if lora_config is None:
+        return [0]
+    if compilation_config.cudagraph_specialize_lora:
+        specialize = getattr(lora_config, "specialize_active_lora", False)
+        captured = get_captured_lora_counts(lora_config.max_loras, specialize)
+        return [0] + [c for c in captured if c > 0]
+    return [0, lora_config.max_loras + 1]
+
+
+def get_num_active_loras_for_dispatch(
+    lora_config: "LoRAConfig | None",
+    lora_state: "LoraState",
+    req_ids: list[str],
+    dummy_run: bool,
+) -> int:
+    """Compute num_active_loras for cudagraph dispatch."""
+    if lora_config and not dummy_run:
+        return len(lora_state.get_activate_loras(req_ids))
+    if dummy_run and lora_config:
+        return lora_config.max_loras + 1
+    return 0
+
+
+def create_lora_capture_hook(
+    lora_config: "LoRAConfig | None",
+    runner: Any,
+) -> Callable[[int, int, int], None] | None:
+    """Create a hook to set up LoRA state before each cudagraph capture."""
+    if lora_config is None:
+        return None
+
+    def hook(num_active_loras: int, num_reqs: int, num_tokens: int) -> None:
+        num_scheduled = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
+        num_scheduled[-1] += num_tokens % num_reqs
+        with runner.maybe_select_dummy_loras(
+            lora_config, num_scheduled, num_active_loras=num_active_loras
+        ):
+            pass
+
+    return hook
+
+
 class LoraState:
     def __init__(self, max_num_reqs: int):
         self.lora_ids = np.zeros(max_num_reqs, dtype=np.int32)
@@ -35,10 +97,13 @@ class LoraState:
         lora_ids = self.lora_ids[idx_mapping]
         prompt_lora_mapping = tuple(lora_ids)
         token_lora_mapping = tuple(lora_ids.repeat(num_scheduled_tokens))
+        active_lora_requests: set[LoRARequest] = self.get_activate_loras(req_ids)
+        return prompt_lora_mapping, token_lora_mapping, active_lora_requests
 
+    def get_activate_loras(self, req_ids: list[str]) -> set[LoRARequest]:
         active_lora_requests: set[LoRARequest] = set()
         for req_id in req_ids:
             lora_request = self.lora_requests.get(req_id)
             if lora_request is not None:
                 active_lora_requests.add(lora_request)
-        return prompt_lora_mapping, token_lora_mapping, active_lora_requests
+        return active_lora_requests

vllm/v1/worker/gpu/model_runner.py

@@ -37,7 +37,6 @@ from vllm.distributed.parallel_state import (
 )
 from vllm.forward_context import BatchDescriptor, set_forward_context
 from vllm.logger import init_logger
-from vllm.lora.layers import LoRAMapping
 from vllm.model_executor.layers.mamba.ops.ssu_dispatch import (
     initialize_mamba_ssu_backend,
 )
@@ -88,7 +87,12 @@ from vllm.v1.worker.gpu.kv_connector import (
     KVConnector,
     get_kv_connector,
 )
-from vllm.v1.worker.gpu.lora_utils import LoraState
+from vllm.v1.worker.gpu.lora_utils import (
+    LoraState,
+    create_lora_capture_hook,
+    get_lora_capture_cases,
+    get_num_active_loras_for_dispatch,
+)
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.mm.lora import set_active_mm_loras
 from vllm.v1.worker.gpu.model_states import init_model_state
@@ -234,8 +238,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         self.prompt_logprobs_worker: PromptLogprobsWorker | None = None
         self.structured_outputs_worker: StructuredOutputsWorker | None = None
         self.cudagraph_manager: ModelCudaGraphManager | None = None
+
         # LoRA-related workers.
         self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
+        self.lora_capture_cases = [0]
+        if self.lora_config:
+            self.lora_capture_cases = get_lora_capture_cases(
+                self.lora_config, self.compilation_config
+            )
+
         # KV Connector if configured.
         self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
 
@@ -458,6 +469,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.device,
             cudagraph_mode,
             decode_query_len=self.decode_query_len,
+            lora_capture_cases=self.lora_capture_cases,
         )
         if self.speculator is not None:
             self.speculator.init_cudagraph_manager(cudagraph_mode)
@@ -540,14 +552,22 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             assert self.intermediate_tensors is not None
             intermediate_tensors = self.intermediate_tensors[:num_tokens]
 
-        # Execute the model.
-        self.execute_model(
-            dummy_scheduler_output,
-            intermediate_tensors=intermediate_tensors,
-            dummy_run=True,
-            skip_attn_for_dummy_run=skip_attn,
-            is_profile=is_profile,
-        )
+        max_loras = self.lora_config.max_loras if self.lora_config is not None else 0
+        with self.maybe_dummy_run_with_lora(
+            self.lora_config,
+            num_scheduled_tokens=np.array(num_tokens_per_request, dtype=np.int32),
+            num_sampled_tokens=None,
+            remove_lora=True,
+            num_active_loras=max_loras,
+        ):
+            # Execute the model.
+            self.execute_model(
+                dummy_scheduler_output,
+                intermediate_tensors=intermediate_tensors,
+                dummy_run=True,
+                skip_attn_for_dummy_run=skip_attn,
+                is_profile=is_profile,
+            )
         self.kv_connector.set_disabled(False)
 
         # Non-last PP ranks don't produce output for sampling.
@@ -694,6 +714,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 self.kv_cache_config,
                 has_lora=self.lora_config is not None,
                 use_aux_hidden_state_outputs=self.use_aux_hidden_state_outputs,
+                lora_capture_hook=create_lora_capture_hook(self.lora_config, self),
             )
             if self.speculator is not None:
                 self.speculator.capture(attn_states)
@@ -1105,6 +1126,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         max_query_len = max(scheduler_output.num_scheduled_tokens.values())
         uniform_tok_count = get_uniform_token_count(num_reqs, num_toks, max_query_len)
 
+        num_active_loras = 0
+        if self.lora_config:
+            req_ids = list(scheduler_output.num_scheduled_tokens.keys())
+            num_active_loras = get_num_active_loras_for_dispatch(
+                self.lora_config, self.lora_state, req_ids, dummy_run
+            )
+
         skip_compiled = False
         if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
             # Encoder-decoder models such as Whisper should run eager/non-compiled
@@ -1120,6 +1148,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.dp_size,
             self.dp_rank,
             need_eager=is_profile or skip_compiled,
+            num_active_loras=num_active_loras,
         )
 
         if batch_desc.num_tokens == 0:
@@ -1157,31 +1186,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 )
                 block_tables = None
                 slot_mappings = None
-            if self.lora_config:
-                # program a no-LoRA mapping here so kernels early-exit instead of
-                # reading uninitialized metadata during dummy runs.
-                # FIXME: Replace this with LoRA warmup:
-                # https://github.com/vllm-project/vllm/pull/35536
-                assert hasattr(self, "lora_manager")
-                adapter_manager = self.lora_manager._adapter_manager
-                adapter_manager.set_adapter_mapping(
-                    LoRAMapping(
-                        index_mapping=(0,) * input_batch.num_tokens_after_padding,
-                        prompt_mapping=(0,) * input_batch.num_reqs,
-                        is_prefill=True,
-                    )
-                )
-                seen_wrappers: set[int] = set()
-                for punica_wrapper in adapter_manager.punica_wrapper_mapping.values():
-                    if id(punica_wrapper) in seen_wrappers:
-                        continue
-                    seen_wrappers.add(id(punica_wrapper))
-                    for kernel_meta in (
-                        punica_wrapper.token_mapping_meta,  # type: ignore[attr-defined]
-                        punica_wrapper.prompt_mapping_meta,  # type: ignore[attr-defined]
-                    ):
-                        kernel_meta.no_lora_flag_cpu[0] = False
-                        kernel_meta.num_active_loras_cpu[0] = 1
 
         attn_metadata = None
         slot_mappings_by_layer = None
@@ -1258,6 +1262,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             batch_descriptor = BatchDescriptor(
                 num_tokens=input_batch.num_tokens_after_padding,
                 has_lora=self.lora_config is not None,
+                num_active_loras=batch_desc.num_active_loras,
             )
 
             with set_forward_context(