[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>
waizuichougou committed
Sep 19, 2026 at 06:13 UTC
a1bf8ac12d9f1537ff2d233f5ab3d1346fd8bd44
4 files changed
+69
-3
tests/v1/worker/test_attn_utils.py
+61
@@ -9,11 +9,13 @@ never addressed by the logical view.
9
10
from types import SimpleNamespace
11
12
+import numpy as np
13
import pytest
14
import torch
15
16
import vllm.v1.hisparse.binding as attn_utils_module
17
from tests.v1.attention.utils import dense_kv_cache_views
18
+from vllm.config.compilation import CUDAGraphMode
19
from vllm.v1.attention.backend import AttentionBackend, AttentionCGSupport, MultipleOf
20
from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
21
from vllm.v1.hisparse.binding import allocate_hisparse_kv_caches
@@ -30,9 +32,11 @@ from vllm.v1.kv_cache_interface import (
32
)
33
from vllm.v1.worker.gpu import attn_utils
34
from vllm.v1.worker.gpu.attn_utils import (
35
+ FastPrefillHelper,
36
get_attn_cg_support,
37
get_query_lens_mismatch_unsupported_backend,
38
)
39
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
40
from vllm.v1.worker.utils import (
41
AttentionGroup,
42
allocate_kv_cache,
@@ -240,6 +244,63 @@ def test_get_kv_sharing_fast_prefill_eligible_layers(monkeypatch: pytest.MonkeyP
244
assert attn_utils.get_kv_sharing_fast_prefill_eligible_layers(vllm_config) == set()
245
246
247
+@pytest.mark.parametrize(
248
+ "num_tokens", [1, 2, 3, 4, 7, 8, 15, 16, 31, 32, 63, 64, 127, 128]
249
+)
250
+@pytest.mark.parametrize("num_active_loras", [1, 2, 4])
251
+def test_fast_prefill_dispatch_preserves_active_lora_count(
252
+ num_tokens: int, num_active_loras: int
253
+):
254
+ """Fast-prefill padding must match the main dispatch's LoRA variant."""
255
+
256
+ class FakeCudaGraphManager:
257
+ device = "cpu"
258
+
259
+ def __init__(self):
260
+ self.dispatch_calls = []
261
+
262
+ def dispatch(self, **kwargs):
263
+ self.dispatch_calls.append(kwargs)
264
+ # A captured no-LoRA graph is padded to 8 tokens, while the
265
+ # active-LoRA path stays eager at the unpadded token count.
266
+ if kwargs["num_active_loras"] == 0:
267
+ num_tokens = 8
268
+ mode = CUDAGraphMode.PIECEWISE
269
+ else:
270
+ num_tokens = kwargs["num_tokens"]
271
+ mode = CUDAGraphMode.NONE
272
+ return BatchExecutionDescriptor(
273
+ cg_mode=mode,
274
+ num_tokens=num_tokens,
275
+ num_reqs=kwargs["num_reqs"],
276
+ num_active_loras=kwargs["num_active_loras"],
277
+ num_ubatches=1,
278
+ )
279
+
280
+ manager = FakeCudaGraphManager()
281
+ helper = FastPrefillHelper(manager, max_num_tokens=max(32, num_tokens))
282
+ metadata = helper.prepare(
283
+ torch.arange(num_tokens, dtype=torch.int32),
284
+ num_reqs=1,
285
+ cu_num_logits_np=np.array([0, num_tokens], dtype=np.int32),
286
+ has_prefill=True,
287
+ batch_desc=BatchExecutionDescriptor(
288
+ cg_mode=CUDAGraphMode.NONE,
289
+ num_tokens=num_tokens,
290
+ num_reqs=1,
291
+ num_active_loras=num_active_loras,
292
+ num_ubatches=1,
293
+ ),
294
+ num_active_loras=num_active_loras,
295
+ )
296
+
297
+ assert metadata is not None
298
+ assert metadata.num_logits_indices == num_tokens
299
+ assert metadata.logits_indices_padded.shape[0] == num_tokens
300
+ assert metadata.max_logits_per_req == num_tokens
301
+ assert manager.dispatch_calls[-1]["num_active_loras"] == num_active_loras
302
+
303
+
304
class _FakeSharedHostRegion:
305
def __init__(self) -> None:
306
self.cleanup_calls = 0
vllm/v1/worker/gpu/attn_utils.py
+2
-1
@@ -95,6 +95,7 @@ class FastPrefillHelper:
95
cu_num_logits_np: np.ndarray,
96
has_prefill: bool,
97
batch_desc: "BatchExecutionDescriptor",
98
+ num_active_loras: int = 0,
99
) -> FastPrefillBatchMetadata | None:
100
if (
101
not has_prefill
@@ -116,7 +117,7 @@ class FastPrefillHelper:
117
num_reqs=num_reqs,
118
num_tokens=num_logits,
119
uniform_token_count=None,
119
- num_active_loras=0,
120
+ num_active_loras=num_active_loras,
121
)
122
num_logits_padded = min(desc.num_tokens, self.max_num_tokens)
123
return FastPrefillBatchMetadata(
vllm/v1/worker/gpu/model_runner.py
+3
-1
@@ -1271,6 +1271,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1271
scheduler_output: SchedulerOutput,
1272
batch_req_state: "BatchReqState",
1273
batch_desc: BatchExecutionDescriptor,
1274
+ num_active_loras: int,
1275
) -> InputBatch:
1276
num_tokens = batch_req_state.num_tokens
1277
num_tokens_after_padding = max(num_tokens, batch_desc.num_tokens)
@@ -1405,6 +1406,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1406
cu_num_logits_np,
1407
batch_req_state.has_prefill,
1408
batch_desc,
1409
+ num_active_loras=num_active_loras,
1410
)
1411
1412
# CPU upper bound on seq_lens; padded entries left at zero.
@@ -1710,7 +1712,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1712
cudagraph_stats = make_cudagraph_stats(batch_desc, num_toks)
1713
assert batch_req_state is not None
1714
input_batch = self.prepare_inputs(
1713
- scheduler_output, batch_req_state, batch_desc
1715
+ scheduler_output, batch_req_state, batch_desc, num_active_loras
1716
)
1717
block_tables, slot_mappings = self.prepare_attn(input_batch)
1718
# Mamba "align" pre-copy: migrate recurrent state across block
vllm/v1/worker/mm_encoder_model_runner.py
+3
-1
@@ -149,7 +149,9 @@ class MMEncoderModelRunner(GPUModelRunner):
149
num_tokens=batch_req_state.num_tokens,
150
num_reqs=None,
151
)
152
- self.prepare_inputs(scheduler_output, batch_req_state, batch_desc)
152
+ self.prepare_inputs(
153
+ scheduler_output, batch_req_state, batch_desc, num_active_loras=0
154
+ )
155
156
scheduled_encoder_inputs = scheduler_output.scheduled_encoder_inputs
157
if self.lora_config is not None: