From a1bf8ac12d9f1537ff2d233f5ab3d1346fd8bd44 Mon Sep 17 00:00:00 2001 From: waizuichougou <105572299+waizuichougou@users.noreply.github.com> Date: Sat, 19 Sep 2026 06:13:23 +0800 Subject: [PATCH] [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 | 61 +++++++++++++++++++++++ vllm/v1/worker/gpu/attn_utils.py | 3 +- vllm/v1/worker/gpu/model_runner.py | 4 +- vllm/v1/worker/mm_encoder_model_runner.py | 4 +- 4 files changed, 69 insertions(+), 3 deletions(-) diff --git a/tests/v1/worker/test_attn_utils.py b/tests/v1/worker/test_attn_utils.py index 2dadc8c980..99481a1293 100644 --- a/tests/v1/worker/test_attn_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 89f54d4d11..e8b4da8f36 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/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( diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 903e5abd67..a05658e7ce 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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 diff --git a/vllm/v1/worker/mm_encoder_model_runner.py b/vllm/v1/worker/mm_encoder_model_runner.py index 4dc0bf77dd..f0d19fa12e 100644 --- a/vllm/v1/worker/mm_encoder_model_runner.py +++ b/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: