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 6, 23:00 UTC
[Model Runner V2] Use unpadded num_tokens for PW CUDA graph attn metadata (#36626)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Mar 10, 2026 at 09:30 UTC
f088a831dd6c35d995c4232cc2462c024c61925b
4 files changed
+14
-3
vllm/v1/worker/gpu/cudagraph_utils.py
+1
@@ -384,6 +384,7 @@ def prepare_inputs_to_capture(
384
385
attn_metadata = model_state.prepare_attn(
386
input_batch,
387
+ CUDAGraphMode.NONE,
388
input_block_tables,
389
slot_mappings,
390
attn_groups,
vllm/v1/worker/gpu/model_runner.py
+1
@@ -936,6 +936,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
936
assert block_tables is not None
937
attn_metadata = self.model_state.prepare_attn(
938
input_batch,
939
+ batch_desc.cg_mode,
940
block_tables,
941
slot_mappings,
942
self.attn_groups,
vllm/v1/worker/gpu/model_states/default.py
+10
-3
@@ -6,6 +6,7 @@ import torch
6
import torch.nn as nn
7
8
from vllm.config import VllmConfig
9
+from vllm.config.compilation import CUDAGraphMode
10
from vllm.v1.core.sched.output import NewRequestData
11
from vllm.v1.kv_cache_interface import KVCacheConfig
12
from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
@@ -140,14 +141,20 @@ class DefaultModelState(ModelState):
141
def prepare_attn(
142
self,
143
input_batch: InputBatch,
144
+ cudagraph_mode: CUDAGraphMode,
145
block_tables: tuple[torch.Tensor, ...],
146
slot_mappings: torch.Tensor,
147
attn_groups: list[list[AttentionGroup]],
148
kv_cache_config: KVCacheConfig,
149
) -> dict[str, Any]:
148
- # Use padded sizes - padding is handled by model_runner.prepare_attn.
149
- num_reqs = input_batch.num_reqs_after_padding
150
- num_tokens = input_batch.num_tokens_after_padding
150
+ if cudagraph_mode == CUDAGraphMode.FULL:
151
+ # Use padded sizes - padding is handled by model_runner.prepare_attn.
152
+ num_reqs = input_batch.num_reqs_after_padding
153
+ num_tokens = input_batch.num_tokens_after_padding
154
+ else:
155
+ # For piecewise cudagraphs and eager, use unpadded sizes.
156
+ num_reqs = input_batch.num_reqs
157
+ num_tokens = input_batch.num_tokens
158
query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
159
max_query_len = input_batch.num_scheduled_tokens.max().item()
160
attn_metadata = build_attn_metadata(
vllm/v1/worker/gpu/model_states/interface.py
+2
@@ -7,6 +7,7 @@ import torch
7
import torch.nn as nn
8
9
from vllm.config import VllmConfig
10
+from vllm.config.compilation import CUDAGraphMode
11
from vllm.v1.core.sched.output import NewRequestData
12
from vllm.v1.kv_cache_interface import KVCacheConfig
13
from vllm.v1.worker.gpu.input_batch import InputBatch
@@ -59,6 +60,7 @@ class ModelState(ABC):
60
def prepare_attn(
61
self,
62
input_batch: InputBatch,
63
+ cudagraph_mode: CUDAGraphMode,
64
block_tables: tuple[torch.Tensor, ...],
65
slot_mappings: torch.Tensor,
66
attn_groups: list[list[AttentionGroup]],