[Bugfix][MRV2] Match fast-prefill padding to active LoRA batches (#56456)

Signed-off-by: waizuichougou <2082431897@qq.com>
Signed-off-by: waizuichougou <105572299+waizuichougou@users.noreply.github.com>

tests/v1/worker/test_attn_utils.py

@@ -9,11 +9,13 @@ never addressed by the logical view.
 
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 import torch
 
 import vllm.v1.hisparse.binding as attn_utils_module
 from tests.v1.attention.utils import dense_kv_cache_views
+from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf
 from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
 from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches
@@ -30,9 +32,11 @@ from vllm.v1.kv_cache_interface import (
 )
 from vllm.v1.worker.gpu import attn_utils
 from vllm.v1.worker.gpu.attn_utils import (
+    FastPrefillHelper,
     get_attn_cg_support,
     get_query_lens_mismatch_unsupported_backend,
 )
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
 from vllm.v1.worker.utils import (
     AttentionGroup,
     allocate_kv_cache,
@@ -240,6 +244,63 @@ def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyP
     assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set()
 
 
+@pytest.mark.parametrize(
+    "num_tokens", [1, 2, 3, 4, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128]
+)
+@pytest.mark.parametrize("num_active_loras", [1, 2, 4])
+def test_fast_prefill_dispatch_preserves_active_lora_count(
+    num_tokens: int, num_active_loras: int
+):
+    """Fast-prefill padding must match the main dispatch's LoRA variant."""
+
+    class FakeCudaGraphManager:
+        device = "cpu"
+
+        def __init__(self):
+            self.dispatch_calls = []
+
+        def dispatch(self, **kwargs):
+            self.dispatch_calls.append(kwargs)
+            # A captured no-LoRA graph is padded to 8 tokens, while the
+            # active-LoRA path stays eager at the unpadded token count.
+            if kwargs["num_active_loras"] == 0:
+                num_tokens = 8
+                mode = CUDAGraphMode.PIECEWISE
+            else:
+                num_tokens = kwargs["num_tokens"]
+                mode = CUDAGraphMode.NONE
+            return BatchExecutionDescriptor(
+                cg_mode=mode,
+                num_tokens=num_tokens,
+                num_reqs=kwargs["num_reqs"],
+                num_active_loras=kwargs["num_active_loras"],
+                num_ubatches=1,
+            )
+
+    manager = FakeCudaGraphManager()
+    helper = FastPrefillHelper(manager, max_num_tokens=max(32, num_tokens))
+    metadata = helper.prepare(
+        torch.arange(num_tokens, dtype=torch.int32),
+        num_reqs=1,
+        cu_num_logits_np=np.array([0, num_tokens], dtype=np.int32),
+        has_prefill=True,
+        batch_desc=BatchExecutionDescriptor(
+            cg_mode=CUDAGraphMode.NONE,
+            num_tokens=num_tokens,
+            num_reqs=1,
+            num_active_loras=num_active_loras,
+            num_ubatches=1,
+        ),
+        num_active_loras=num_active_loras,
+    )
+
+    assert metadata is not None
+    assert metadata.num_logits_indices == num_tokens
+    assert metadata.logits_indices_padded.shape[0] == num_tokens
+    assert metadata.max_logits_per_req == num_tokens
+    assert manager.dispatch_calls[-1]["num_active_loras"] == num_active_loras
+
+
 class _FakeSharedHostRegion:
     def __init__(self) -> None:
         self.cleanup_calls = 0

vllm/v1/worker/gpu/attn_utils.py

@@ -95,6 +95,7 @@ class FastPrefillHelper:
         cu_num_logits_np: np.ndarray,
         has_prefill: bool,
         batch_desc: "BatchExecutionDescriptor",
+        num_active_loras: int = 0,
     ) -> FastPrefillBatchMetadata | None:
         if (
             not has_prefill
@@ -116,7 +117,7 @@ class FastPrefillHelper:
             num_reqs=num_reqs,
             num_tokens=num_logits,
             uniform_token_count=None,
-            num_active_loras=0,
+            num_active_loras=num_active_loras,
         )
         num_logits_padded = min(desc.num_tokens, self.max_num_tokens)
         return FastPrefillBatchMetadata(

vllm/v1/worker/gpu/model_runner.py

@@ -1271,6 +1271,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         scheduler_output: SchedulerOutput,
         batch_req_state: "BatchReqState",
         batch_desc: BatchExecutionDescriptor,
+        num_active_loras: int,
     ) -> InputBatch:
         num_tokens = batch_req_state.num_tokens
         num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
@@ -1405,6 +1406,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 cu_num_logits_np,
                 batch_req_state.has_prefill,
                 batch_desc,
+                num_active_loras=num_active_loras,
             )
 
         # CPU upper bound on seq_lens; padded entries left at zero.
@@ -1710,7 +1712,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 cudagraph_stats = make_cudagraph_stats(batch_desc, num_toks)
             assert batch_req_state is not None
             input_batch = self.prepare_inputs(
-                scheduler_output, batch_req_state, batch_desc
+                scheduler_output, batch_req_state, batch_desc, num_active_loras
             )
             block_tables, slot_mappings = self.prepare_attn(input_batch)
             # Mamba "align" pre-copy: migrate recurrent state across block

vllm/v1/worker/mm_encoder_model_runner.py

@@ -149,7 +149,9 @@ class MMEncoderModelRunner(GPUModelRunner):
                 num_tokens=batch_req_state.num_tokens,
                 num_reqs=None,
             )
-            self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
+            self.prepare_inputs(
+                scheduler_output, batch_req_state, batch_desc, num_active_loras=0
+            )
 
         scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
         if self.lora_config is not None: