A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 10:15 UTC
[Model Runner V2] Use ModelState.prepare_attn() for cuda graph capture [5/N] (#35774)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 2, 2026 at 20:06 UTC
a0a5178ab4c032c8031e11563e4f1b29a65b92a5
5 files changed
+38
-53
vllm/v1/worker/gpu/cudagraph_utils.py
+24
-45
@@ -3,7 +3,6 @@
3
from collections.abc import Callable
4
from typing import Any
5
6
-import numpy as np
6
import torch
7
import torch.nn as nn
8
from tqdm import tqdm
@@ -15,13 +14,11 @@ from vllm.forward_context import BatchDescriptor, set_forward_context
14
from vllm.model_executor.offloader.base import get_offloader
15
from vllm.utils.math_utils import cdiv
16
from vllm.v1.kv_cache_interface import KVCacheConfig
18
-from vllm.v1.worker.gpu.attn_utils import (
19
- build_attn_metadata,
20
- build_slot_mappings_by_layer,
21
-)
17
+from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer
18
from vllm.v1.worker.gpu.block_table import BlockTables
19
+from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens
20
from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
24
-from vllm.v1.worker.gpu.input_batch import InputBuffers
21
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
22
from vllm.v1.worker.gpu.model_states.interface import ModelState
23
from vllm.v1.worker.utils import AttentionGroup
24
@@ -123,14 +120,11 @@ class CudaGraphManager:
120
attn_metadata, slot_mappings = prepare_inputs_to_capture(
121
num_reqs,
122
num_tokens,
123
+ model_state,
124
input_buffers,
125
block_tables,
126
attn_groups,
129
- self.max_model_len,
127
kv_cache_config,
131
- uniform_decode_query_len=(
132
- self.uniform_decode_query_len if uniform_decode else 0
133
- ),
128
)
129
num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
130
@@ -393,51 +387,36 @@ def capture_graphs(
387
def prepare_inputs_to_capture(
388
num_reqs: int,
389
num_tokens: int,
390
+ model_state: ModelState,
391
input_buffers: InputBuffers,
392
block_tables: BlockTables,
393
attn_groups: list[list[AttentionGroup]],
399
- max_model_len: int,
394
kv_cache_config: KVCacheConfig,
401
- uniform_decode_query_len: int = 0,
395
) -> tuple[dict[str, Any], dict[str, torch.Tensor]]:
403
- if uniform_decode_query_len > 0:
404
- num_tokens_per_req = uniform_decode_query_len
405
- else:
406
- num_tokens_per_req = num_tokens // num_reqs
407
-
408
- query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) * num_tokens_per_req
409
- query_start_loc_np[-1] = num_tokens
410
- query_start_loc_cpu = torch.from_numpy(query_start_loc_np)
411
- input_buffers.query_start_loc[: num_reqs + 1] = query_start_loc_cpu
412
- input_buffers.query_start_loc[num_reqs + 1 :] = num_tokens
413
- query_start_loc = input_buffers.query_start_loc[: num_reqs + 1]
414
-
415
- # HACK(woosuk): For faster warmup, we set seq_lens (GPU) to num_tokens
416
- # rather than max_model_len.
417
- input_buffers.seq_lens[:num_reqs] = num_tokens
418
- input_buffers.seq_lens[num_reqs:] = 0
419
-
420
- input_buffers.dcp_local_seq_lens[:num_reqs] = num_tokens
421
- input_buffers.dcp_local_seq_lens[num_reqs:] = 0
422
-
396
+ input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
397
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
398
slot_mappings = block_tables.get_dummy_slot_mappings(num_tokens)
399
slot_mappings_by_layer = build_slot_mappings_by_layer(
400
slot_mappings, kv_cache_config
401
)
402
429
- attn_metadata = build_attn_metadata(
430
- attn_groups=attn_groups,
431
- num_reqs=num_reqs,
432
- num_tokens=num_tokens,
433
- query_start_loc_gpu=query_start_loc,
434
- query_start_loc_cpu=query_start_loc_cpu,
435
- max_query_len=num_tokens_per_req,
436
- seq_lens=input_buffers.seq_lens,
437
- max_seq_len=max_model_len,
438
- block_tables=input_block_tables,
439
- slot_mappings=slot_mappings,
440
- kv_cache_config=kv_cache_config,
441
- dcp_local_seq_lens=input_buffers.dcp_local_seq_lens,
403
+ # HACK(woosuk): Special handling for DCP.
404
+ if block_tables.cp_size > 1:
405
+ prepare_dcp_local_seq_lens(
406
+ input_buffers.dcp_local_seq_lens,
407
+ input_batch.seq_lens,
408
+ num_reqs,
409
+ block_tables.cp_size,
410
+ block_tables.cp_rank,
411
+ block_tables.cp_interleave,
412
+ )
413
+ input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs]
414
+
415
+ attn_metadata = model_state.prepare_attn(
416
+ input_batch,
417
+ input_block_tables,
418
+ slot_mappings,
419
+ attn_groups,
420
+ kv_cache_config,
421
)
422
return attn_metadata, slot_mappings_by_layer
vllm/v1/worker/gpu/input_batch.py
+3
-2
@@ -82,14 +82,16 @@ class InputBatch:
82
num_reqs: int,
83
num_tokens: int,
84
input_buffers: InputBuffers,
85
- device: torch.device,
85
) -> "InputBatch":
86
assert 0 < num_reqs <= num_tokens
87
+ device = input_buffers.device
88
+
89
req_ids = [f"req_{i}_{random_uuid()}" for i in range(num_reqs)]
90
idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
91
idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
92
expanded_idx_mapping = idx_mapping
93
expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)
94
+
95
num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
96
num_scheduled_tokens[-1] += num_tokens % num_reqs
97
assert int(num_scheduled_tokens.sum()) == num_tokens
@@ -115,7 +117,6 @@ class InputBatch:
117
input_ids = input_buffers.input_ids[:num_tokens].zero_()
118
positions = input_buffers.positions[:num_tokens].zero_()
119
118
- # attn_metadata = defaultdict(lambda: None)
120
logits_indices = query_start_loc[1:] - 1
121
cu_num_logits = torch.arange(num_reqs + 1, device=device, dtype=torch.int32)
122
cu_num_logits_np = np.arange(num_reqs + 1, dtype=np.int32)
vllm/v1/worker/gpu/model_runner.py
+2
-4
@@ -311,6 +311,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
311
if self.speculator is not None:
312
# HACK(woosuk)
313
self.speculator.set_attn(
314
+ self.model_state,
315
self.kv_cache_config,
316
self.attn_groups,
317
self.block_tables,
@@ -880,10 +881,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
881
# No actual tokens to run. A dummy run for DP or memory profiling.
882
num_reqs = min(num_tokens_after_padding, self.max_num_reqs)
883
input_batch = InputBatch.make_dummy(
883
- num_reqs=num_reqs,
884
- num_tokens=num_tokens_after_padding,
885
- input_buffers=self.input_buffers,
886
- device=self.device,
884
+ num_reqs, num_tokens_after_padding, self.input_buffers
885
)
886
if not skip_attn_for_dummy_run:
887
block_tables, slot_mappings = self.prepare_dummy_attn(input_batch)
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+5
-2
@@ -17,6 +17,7 @@ from vllm.v1.worker.gpu.cudagraph_utils import (
17
)
18
from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
19
from vllm.v1.worker.gpu.input_batch import InputBuffers
20
+from vllm.v1.worker.gpu.model_states.interface import ModelState
21
from vllm.v1.worker.utils import AttentionGroup
22
23
@@ -59,6 +60,7 @@ class EagleCudaGraphManager:
60
num_tokens: int,
61
capture_cg_mode: CUDAGraphMode,
62
generate_fn: Callable,
63
+ model_state: ModelState,
64
input_buffers: InputBuffers,
65
block_tables: BlockTables,
66
attn_groups: list[list[AttentionGroup]],
@@ -76,12 +78,11 @@ class EagleCudaGraphManager:
78
attn_metadata, slot_mappings = prepare_inputs_to_capture(
79
num_reqs,
80
num_tokens,
81
+ model_state,
82
input_buffers,
83
block_tables,
84
attn_groups,
82
- self.max_model_len,
85
kv_cache_config,
84
- uniform_decode_query_len=1,
86
)
87
num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
88
@@ -158,6 +159,7 @@ class EagleCudaGraphManager:
159
def capture(
160
self,
161
generate_fn: Callable,
162
+ model_state: ModelState,
163
input_buffers: InputBuffers,
164
block_tables: BlockTables,
165
attn_groups: list[list[AttentionGroup]],
@@ -173,6 +175,7 @@ class EagleCudaGraphManager:
175
capture_cudagraph_mode=self.cudagraph_mode,
176
desc=f"Capturing eagle CUDA graphs ({self.cudagraph_mode.name})",
177
generate_fn=generate_fn,
178
+ model_state=model_state,
179
input_buffers=input_buffers,
180
block_tables=block_tables,
181
attn_groups=attn_groups,
vllm/v1/worker/gpu/spec_decode/eagle/speculator.py
+4
@@ -17,6 +17,7 @@ from vllm.v1.worker.gpu.attn_utils import (
17
)
18
from vllm.v1.worker.gpu.block_table import BlockTables
19
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
20
+from vllm.v1.worker.gpu.model_states.interface import ModelState
21
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
22
from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import EagleCudaGraphManager
23
from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model
@@ -76,10 +77,12 @@ class EagleSpeculator:
77
78
def set_attn(
79
self,
80
+ model_state: ModelState,
81
kv_cache_config: KVCacheConfig,
82
attn_groups: list[list[AttentionGroup]],
83
block_tables: BlockTables,
84
) -> None:
85
+ self.model_state = model_state
86
self.kv_cache_config = kv_cache_config
87
self.attn_groups = attn_groups
88
self.block_tables = block_tables
@@ -171,6 +174,7 @@ class EagleSpeculator:
174
logger.info("Capturing model for Eagle speculator...")
175
self.cudagraph_manager.capture(
176
self.generate_draft,
177
+ self.model_state,
178
self.input_buffers,
179
self.block_tables,
180
self.attn_groups,