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 4, 03:45 UTC
[Model Runner V2] support qwen35 / mamba hybrid model (#35520)
Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>
zhrrr committed
May 8, 2026 at 00:31 UTC
7a08b34fbf17a491ccf1f7e2cff3b8184a9e4454
14 files changed
+440
-107
vllm/config/kernel.py
+1
-1
@@ -142,7 +142,7 @@ class KernelConfig:
142
"""Backend for MoE expert computation kernels. Available options:
143
144
- "auto": Automatically select the best backend based on model and hardware
145
- - "triton": Use Triton-based fused MoE kernels
145
+ - "triton": Use Triton-based fused MoE kernels
146
- "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)
147
- "deep_gemm_mega_moe": Use DeepGEMM mega MoE kernels
148
- "cutlass": Use vLLM CUTLASS kernels
vllm/config/vllm.py
+4
@@ -1892,6 +1892,10 @@ class VllmConfig:
1892
"to schedule a multiple of block_size tokens even if they are "
1893
"in the middle of a mm input"
1894
)
1895
+ # TODO: support align mamba cache mode for model runner v2
1896
+ assert not envs.VLLM_USE_V2_MODEL_RUNNER, (
1897
+ "Model Runner V2 has not yet supported mamba_cache_mode='align'. "
1898
+ )
1899
1900
@model_validator(mode="after")
1901
def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig":
vllm/multimodal/registry.py
+14
-2
@@ -111,7 +111,15 @@ class MultiModalRegistry:
111
return False
112
113
mm_config = model_config.get_multimodal_config()
114
- info = self._create_processing_info(model_config, tokenizer=None)
114
+ try:
115
+ info = self._create_processing_info(model_config, tokenizer=None)
116
+ except ValueError:
117
+ logger.warning_once(
118
+ "Model %s is treated as multimodal but has no registered "
119
+ "multimodal processor; running in text-only mode.",
120
+ model_config.model,
121
+ )
122
+ return False
123
124
# Check if all supported modalities have limit == 0
125
if all(
@@ -170,7 +178,11 @@ class MultiModalRegistry:
178
from vllm.model_executor.model_loader import get_model_architecture
179
180
model_cls, _ = get_model_architecture(model_config)
173
- assert hasattr(model_cls, "_processor_factory")
181
+ if not hasattr(model_cls, "_processor_factory"):
182
+ raise ValueError(
183
+ f"Model class {model_cls.__name__} has no registered "
184
+ "multimodal processor"
185
+ )
186
return cast("SupportsMultiModal", model_cls)
187
188
def _create_processing_ctx(
vllm/v1/worker/gpu/attn_utils.py
+138
-61
@@ -4,7 +4,6 @@ from collections.abc import Sequence
4
from dataclasses import dataclass
5
from typing import Any, cast
6
7
-import numpy as np
7
import torch
8
9
from vllm.config import VllmConfig, get_layers_from_vllm_config
@@ -19,8 +18,10 @@ from vllm.v1.kv_cache_interface import (
18
AttentionSpec,
19
KVCacheConfig,
20
KVCacheSpec,
21
+ MambaSpec,
22
UniformTypeKVCacheSpecs,
23
)
24
+from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
25
from vllm.v1.worker.utils import AttentionGroup, bind_kv_cache
26
27
@@ -147,68 +148,126 @@ def _reshape_kv_cache(
148
kv_cache_raw_tensors: dict[str, torch.Tensor],
149
attn_backends: dict[str, type[AttentionBackend]],
150
cache_dtype: str,
150
-) -> dict[str, torch.Tensor]:
151
- kv_caches: dict[str, torch.Tensor] = {}
151
+) -> dict[str, Any]:
152
+ kv_caches: dict[str, Any] = {}
153
+ has_attn, has_mamba = False, False
154
for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
155
for layer_name in kv_cache_group_spec.layer_names:
156
kv_cache_spec = kv_cache_group_spec.kv_cache_spec
157
if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
158
kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name]
157
- assert isinstance(kv_cache_spec, AttentionSpec)
158
-
159
- raw_tensor = kv_cache_raw_tensors[layer_name]
160
- assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
161
- num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
162
-
163
- attn_backend = attn_backends[layer_name]
164
- kv_cache_shape = attn_backend.get_kv_cache_shape(
165
- num_blocks,
166
- kv_cache_spec.storage_block_size,
167
- kv_cache_spec.num_kv_heads,
168
- kv_cache_spec.head_size,
169
- cache_dtype,
170
- )
159
172
- # FIXME(woosuk): Add kv_cache_stride_order to all attention backends.
173
- try:
174
- kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
175
- assert len(kv_cache_stride_order) == len(kv_cache_shape)
176
- except (AttributeError, NotImplementedError):
177
- kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
178
-
179
- kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
180
- inv_order = [
181
- kv_cache_stride_order.index(i)
182
- for i in range(len(kv_cache_stride_order))
183
- ]
184
-
185
- dtype = kv_cache_spec.dtype
186
- raw_tensor = raw_tensor.view(dtype)
187
- if kv_cache_spec.page_size_padded is not None:
188
- # Use strided view to handle page_size_bytes that
189
- # include padding. This follows the same pattern as
190
- # MambaSpec handling in gpu_model_runner.py.
191
- # NOTE: This assumes kv_cache_shape[0] == num_blocks
192
- # (i.e. the first physical dimension is the block
193
- # index), which holds for MLA backends but NOT for
194
- # standard attention backends whose shape starts with
195
- # a K/V dimension of size 2.
196
- dtype_size = get_dtype_size(dtype)
197
- page_stride = kv_cache_spec.page_size_bytes // dtype_size
198
- strides = list(torch.empty(kv_cache_shape).stride())
199
- strides[inv_order[0]] = page_stride
200
- kv_cache = torch.as_strided(
201
- raw_tensor,
202
- size=kv_cache_shape,
203
- stride=tuple(strides),
160
+ kv_raw_tensor = kv_cache_raw_tensors[layer_name]
161
+ assert kv_raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
162
+ num_blocks = kv_raw_tensor.numel() // kv_cache_spec.page_size_bytes
163
+
164
+ if isinstance(kv_cache_spec, AttentionSpec):
165
+ has_attn = True
166
+ attn_backend = attn_backends[layer_name]
167
+ kv_cache_shape = attn_backend.get_kv_cache_shape(
168
+ num_blocks,
169
+ kv_cache_spec.storage_block_size,
170
+ kv_cache_spec.num_kv_heads,
171
+ kv_cache_spec.head_size,
172
+ cache_dtype_str=cache_dtype,
173
)
174
+
175
+ # FIXME(woosuk): Add kv_cache_stride_order to all attention backends.
176
+ try:
177
+ kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
178
+ assert len(kv_cache_stride_order) == len(kv_cache_shape)
179
+ except (AttributeError, NotImplementedError):
180
+ kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
181
+
182
+ kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
183
+ inv_order = [
184
+ kv_cache_stride_order.index(i)
185
+ for i in range(len(kv_cache_stride_order))
186
+ ]
187
+
188
+ dtype = kv_cache_spec.dtype
189
+ kv_tensor = kv_raw_tensor.view(dtype)
190
+ if kv_cache_spec.page_size_padded is not None:
191
+ # Use strided view to handle page_size_bytes that
192
+ # include padding. This follows the same pattern as
193
+ # MambaSpec handling in gpu_model_runner.py.
194
+ # NOTE: This assumes kv_cache_shape[0] == num_blocks
195
+ # (i.e. the first physical dimension is the block
196
+ # index), which holds for MLA backends but NOT for
197
+ # standard attention backends whose shape starts with
198
+ # a K/V dimension of size 2.
199
+ dtype_size = get_dtype_size(dtype)
200
+ page_stride = kv_cache_spec.page_size_bytes // dtype_size
201
+ strides = list(torch.empty(kv_cache_shape).stride())
202
+ strides[inv_order[0]] = page_stride
203
+ kv_cache = torch.as_strided(
204
+ kv_tensor,
205
+ size=kv_cache_shape,
206
+ stride=tuple(strides),
207
+ )
208
+ else:
209
+ # No padding — safe to use a contiguous view.
210
+ kv_cache = kv_tensor.view(kv_cache_shape)
211
+ kv_caches[layer_name] = kv_cache.permute(*inv_order)
212
+
213
+ elif isinstance(kv_cache_spec, MambaSpec):
214
+ has_mamba = True
215
+ state_tensors = []
216
+ storage_offset_bytes = 0
217
+ for shape, dtype in zip(kv_cache_spec.shapes, kv_cache_spec.dtypes):
218
+ dtype_size = get_dtype_size(dtype)
219
+ num_element_per_page = kv_cache_spec.page_size_bytes // dtype_size
220
+ target_shape = (num_blocks, *shape)
221
+ stride = torch.empty(target_shape).stride()
222
+ target_stride = (num_element_per_page, *stride[1:])
223
+ assert storage_offset_bytes % dtype_size == 0
224
+ tensor = torch.as_strided(
225
+ kv_raw_tensor.view(dtype),
226
+ size=target_shape,
227
+ stride=target_stride,
228
+ storage_offset=storage_offset_bytes // dtype_size,
229
+ )
230
+ state_tensors.append(tensor)
231
+ storage_offset_bytes += stride[0] * dtype_size
232
+ kv_caches[layer_name] = state_tensors
233
else:
206
- # No padding — safe to use a contiguous view.
207
- kv_cache = raw_tensor.view(kv_cache_shape)
208
- kv_caches[layer_name] = kv_cache.permute(*inv_order)
234
+ raise NotImplementedError(
235
+ f"Unsupported KV cache spec type: {type(kv_cache_spec)}"
236
+ )
237
+
238
+ if has_attn and has_mamba:
239
+ _update_hybrid_attention_layout(kv_caches, kv_cache_config)
240
+
241
return kv_caches
242
243
244
+def _update_hybrid_attention_layout(
245
+ kv_caches: dict[str, Any],
246
+ kv_cache_config: KVCacheConfig,
247
+) -> None:
248
+ for kv_cache_group_spec in kv_cache_config.kv_cache_groups:
249
+ for layer_name in kv_cache_group_spec.layer_names:
250
+ kv_cache_spec = kv_cache_group_spec.kv_cache_spec
251
+ if isinstance(kv_cache_spec, UniformTypeKVCacheSpecs):
252
+ kv_cache_spec = kv_cache_spec.kv_cache_specs[layer_name]
253
+ if not isinstance(kv_cache_spec, AttentionSpec):
254
+ continue
255
+ kv_cache = kv_caches[layer_name]
256
+ if kv_cache.shape[0] == 2:
257
+ assert kv_cache.shape[1] != 2, (
258
+ f"Cannot determine layout for tensor of shape {kv_cache.shape}"
259
+ )
260
+ hidden_size = kv_cache.shape[2:].numel()
261
+ kv_cache.as_strided_(
262
+ size=kv_cache.shape,
263
+ stride=(
264
+ hidden_size,
265
+ 2 * hidden_size,
266
+ *kv_cache.stride()[2:],
267
+ ),
268
+ )
269
+
270
+
271
def init_kv_cache(
272
runner_kv_caches: list[torch.Tensor],
273
forward_context: dict[str, Any],
@@ -216,7 +275,7 @@ def init_kv_cache(
275
attn_backends: dict[str, type[AttentionBackend]],
276
device: torch.device,
277
cache_dtype: str,
219
-) -> dict[str, torch.Tensor]:
278
+) -> dict[str, Any]:
279
kv_cache_raw_tensors = _allocate_kv_cache(kv_cache_config, device)
280
kv_caches = _reshape_kv_cache(
281
kv_cache_config, kv_cache_raw_tensors, attn_backends, cache_dtype
@@ -250,8 +309,9 @@ def build_attn_metadata(
309
kv_cache_config: KVCacheConfig,
310
seq_lens_cpu_upper_bound: torch.Tensor | None = None,
311
dcp_local_seq_lens: torch.Tensor | None = None,
253
- encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]] | None = None,
312
positions: torch.Tensor | None = None,
313
+ model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
314
+ for_cudagraph_capture: bool = False,
315
) -> dict[str, Any]:
316
seq_lens = seq_lens[:num_reqs]
317
if dcp_local_seq_lens is not None:
@@ -265,6 +325,11 @@ def build_attn_metadata(
325
block_table = block_tables[i]
326
slot_mapping = slot_mappings[i]
327
328
+ common_attn_metadata_extra_kwargs = (
329
+ model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs)
330
+ if model_specific_attn_metadata is not None
331
+ else {}
332
+ )
333
common_attn_metadata = CommonAttentionMetadata(
334
query_start_loc=query_start_loc_gpu,
335
query_start_loc_cpu=query_start_loc_cpu,
@@ -279,17 +344,29 @@ def build_attn_metadata(
344
causal=True,
345
dcp_local_seq_lens=dcp_local_seq_lens,
346
positions=positions,
347
+ **common_attn_metadata_extra_kwargs,
348
)
283
- if encoder_seq_lens and i in encoder_seq_lens:
284
- encoder_seq_lens_gpu, encoder_seq_lens_cpu = encoder_seq_lens[i]
285
- common_attn_metadata.encoder_seq_lens = encoder_seq_lens_gpu
286
- common_attn_metadata.encoder_seq_lens_cpu = encoder_seq_lens_cpu
349
350
for attn_group in attn_groups[i]:
351
attn_metadata_builder = attn_group.get_metadata_builder(0)
290
- metadata = attn_metadata_builder.build(
291
- common_prefix_len=0, common_attn_metadata=common_attn_metadata
292
- )
352
+ if for_cudagraph_capture:
353
+ metadata = attn_metadata_builder.build_for_cudagraph_capture(
354
+ common_attn_metadata
355
+ )
356
+ else:
357
+ attn_metadata_extra_kwargs = (
358
+ model_specific_attn_metadata.get_extra_attn_kwargs(
359
+ attn_metadata_builder,
360
+ num_reqs,
361
+ )
362
+ if model_specific_attn_metadata is not None
363
+ else {}
364
+ )
365
+ metadata = attn_metadata_builder.build(
366
+ common_prefix_len=0,
367
+ common_attn_metadata=common_attn_metadata,
368
+ **attn_metadata_extra_kwargs,
369
+ )
370
for layer_name in attn_group.layer_names:
371
attn_metadata[layer_name] = metadata
372
return attn_metadata
vllm/v1/worker/gpu/block_table.py
+3
-13
@@ -5,7 +5,6 @@ from collections.abc import Iterable
5
import torch
6
7
from vllm.triton_utils import tl, triton
8
-from vllm.utils.math_utils import cdiv
8
from vllm.v1.attention.backends.utils import PAD_SLOT_ID
9
from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
10
@@ -16,7 +15,7 @@ class BlockTables:
15
block_sizes: list[int],
16
max_num_reqs: int,
17
max_num_batched_tokens: int,
19
- max_model_len: int,
18
+ max_num_blocks_per_group: list[int],
19
device: torch.device,
20
cp_size: int = 1,
21
cp_rank: int = 0,
@@ -25,7 +24,6 @@ class BlockTables:
24
self.block_sizes = block_sizes
25
self.max_num_reqs = max_num_reqs
26
self.max_num_batched_tokens = max_num_batched_tokens
28
- self.max_model_len = max_model_len
27
self.device = device
28
29
self.cp_size = cp_size
@@ -33,19 +31,11 @@ class BlockTables:
31
self.cp_interleave = cp_interleave
32
33
self.num_kv_cache_groups = len(self.block_sizes)
34
+ assert len(max_num_blocks_per_group) == self.num_kv_cache_groups
35
# num_kv_cache_groups x [max_num_reqs, max_num_blocks]
36
self.block_tables: list[StagedWriteTensor] = []
37
for i in range(self.num_kv_cache_groups):
39
- block_size = self.block_sizes[i]
40
- # When using DCP, each request's KV cache is sharded among different ranks.
41
- # As a result, one block on the current rank covers `block_size * cp_size`
42
- # tokens in the full, global (unsharded) sequence.
43
- max_num_blocks = cdiv(self.max_model_len, block_size * self.cp_size)
44
- # Align to a multiple of (128 / block_size) as required
45
- # by some attention backends such as TRTLLM (#39324)
46
- if block_size <= 128:
47
- alignment = 128 // block_size
48
- max_num_blocks = cdiv(max_num_blocks, alignment) * alignment
38
+ max_num_blocks = max_num_blocks_per_group[i]
39
block_table = StagedWriteTensor(
40
(self.max_num_reqs, max_num_blocks),
41
dtype=torch.int32,
vllm/v1/worker/gpu/cudagraph_utils.py
+24
-16
@@ -34,7 +34,7 @@ logger = init_logger(__name__)
34
35
36
class CapturedAttentionState(NamedTuple):
37
- attn_metadata: dict[str, Any]
37
+ attn_metadata: dict[str, Any] | None
38
slot_mappings: dict[str, torch.Tensor]
39
40
@@ -100,6 +100,7 @@ class CudaGraphManager:
100
self.decode_query_len = decode_query_len
101
102
self.dp_size = vllm_config.parallel_config.data_parallel_size
103
+ self.tp_size = vllm_config.parallel_config.tensor_parallel_size
104
self.is_first_pp_rank = get_pp_group().is_first_rank
105
self.is_last_pp_rank = get_pp_group().is_last_rank
106
@@ -109,6 +110,10 @@ class CudaGraphManager:
110
self._graphs_captured = False
111
self._candidates: list[list[BatchExecutionDescriptor]] = []
112
self._capture_descs: dict[CUDAGraphMode, list[BatchExecutionDescriptor]] = {}
113
+ # adjust the cudagraph sizes to be a multiple of the uniform decode query length
114
+ self.compilation_config.adjust_cudagraph_sizes_for_spec_decode(
115
+ self.decode_query_len, self.tp_size
116
+ )
117
self._init_candidates()
118
119
def _init_candidates(self) -> None:
@@ -340,16 +345,16 @@ class ModelCudaGraphManager(CudaGraphManager):
345
block_tables,
346
attn_groups,
347
kv_cache_config,
348
+ skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
349
)
350
351
def forward_fn(cg_mode: CUDAGraphMode) -> None:
346
- batch_descriptor = (
347
- BatchDescriptor(num_tokens=num_tokens)
348
- if cg_mode == CUDAGraphMode.PIECEWISE
349
- else None
350
- )
352
+ batch_descriptor = None
353
+ if cg_mode == CUDAGraphMode.PIECEWISE:
354
+ assert attn_metadata is None
355
+ batch_descriptor = BatchDescriptor(num_tokens=num_tokens)
356
with set_forward_context(
352
- attn_metadata if cg_mode != CUDAGraphMode.PIECEWISE else None,
357
+ attn_metadata,
358
self.vllm_config,
359
num_tokens=num_tokens,
360
cudagraph_runtime_mode=cg_mode,
@@ -419,6 +424,7 @@ def prepare_inputs_to_capture(
424
block_tables: BlockTables,
425
attn_groups: list[list[AttentionGroup]],
426
kv_cache_config: KVCacheConfig,
427
+ skip_attn: bool = False,
428
) -> CapturedAttentionState:
429
input_batch = InputBatch.make_dummy(num_reqs, num_tokens, input_buffers)
430
input_block_tables = block_tables.get_dummy_block_tables(num_reqs)
@@ -439,13 +445,15 @@ def prepare_inputs_to_capture(
445
)
446
input_batch.dcp_local_seq_lens = input_buffers.dcp_local_seq_lens[:num_reqs]
447
442
- attn_metadata = model_state.prepare_attn(
443
- input_batch,
444
- CUDAGraphMode.NONE,
445
- input_block_tables,
446
- slot_mappings,
447
- attn_groups,
448
- kv_cache_config,
449
- for_capture=True,
450
- )
448
+ attn_metadata = None
449
+ if not skip_attn:
450
+ attn_metadata = model_state.prepare_attn(
451
+ input_batch,
452
+ CUDAGraphMode.NONE,
453
+ input_block_tables,
454
+ slot_mappings,
455
+ attn_groups,
456
+ kv_cache_config,
457
+ for_capture=True,
458
+ )
459
return CapturedAttentionState(attn_metadata, slot_mappings_by_layer)
vllm/v1/worker/gpu/input_batch.py
+7
@@ -53,7 +53,10 @@ class InputBatch:
53
# sum(num_scheduled_tokens)
54
num_tokens: int
55
num_tokens_after_padding: int
56
+ # Sum of draft tokens scheduled across requests.
57
num_draft_tokens: int
58
+ # [num_reqs] number of draft tokens scheduled for each request, if any.
59
+ num_draft_tokens_per_req: np.ndarray | None
60
61
# [num_reqs + 1]
62
query_start_loc: torch.Tensor
@@ -64,6 +67,8 @@ class InputBatch:
67
seq_lens_cpu_upper_bound: torch.Tensor
68
# [num_reqs]
69
dcp_local_seq_lens: torch.Tensor | None
70
+ # [num_reqs] CPU bool array.
71
+ is_prefilling_np: np.ndarray
72
73
# [num_tokens_after_padding]
74
input_ids: torch.Tensor
@@ -137,11 +142,13 @@ class InputBatch:
142
num_tokens=num_tokens,
143
num_tokens_after_padding=num_tokens,
144
num_draft_tokens=0,
145
+ num_draft_tokens_per_req=None,
146
query_start_loc=query_start_loc,
147
query_start_loc_np=query_start_loc_np,
148
seq_lens=seq_lens,
149
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
150
dcp_local_seq_lens=None,
151
+ is_prefilling_np=np.zeros(num_reqs, dtype=np.bool_),
152
input_ids=input_ids,
153
positions=positions,
154
logits_indices=logits_indices,
vllm/v1/worker/gpu/model_runner.py
+39
-10
@@ -44,10 +44,11 @@ from vllm.model_executor.model_loader import get_model_loader
44
from vllm.multimodal import MULTIMODAL_REGISTRY
45
from vllm.sequence import IntermediateTensors
46
from vllm.tasks import SupportedTask
47
+from vllm.utils.math_utils import cdiv
48
from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
49
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
50
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
50
-from vllm.v1.kv_cache_interface import KVCacheConfig
51
+from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
52
from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
53
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
54
from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
@@ -338,10 +339,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
339
def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
340
kv_cache_config = deepcopy(kv_cache_config)
341
self.kv_cache_config = kv_cache_config
341
- block_sizes = [
342
- kv_cache_group.kv_cache_spec.block_size
343
- for kv_cache_group in kv_cache_config.kv_cache_groups
344
- ]
342
343
block_table_max_model_len = self.max_model_len
344
if self.is_encoder_decoder:
@@ -352,11 +349,34 @@ class GPUModelRunner(LoRAModelRunnerMixin):
349
getattr(self.model_config.hf_config, "max_source_positions", 0),
350
)
351
352
+ block_sizes = []
353
+ max_num_blocks_per_group = []
354
+ for kv_cache_group in kv_cache_config.kv_cache_groups:
355
+ spec = kv_cache_group.kv_cache_spec
356
+ block_sizes.append(spec.block_size)
357
+ # When using DCP, each request's KV cache is sharded among different ranks.
358
+ # As a result, one block on the current rank covers `block_size * cp_size`
359
+ # tokens in the full, global (unsharded) sequence.
360
+ max_num_blocks = cdiv(
361
+ block_table_max_model_len, spec.block_size * self.dcp_size
362
+ )
363
+ # Align to a multiple of (128 / block_size) as required by some attention
364
+ # backends such as TRTLLM (#39324)
365
+ if spec.block_size <= 128:
366
+ alignment = 128 // spec.block_size
367
+ max_num_blocks = cdiv(max_num_blocks, alignment) * alignment
368
+ # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens
369
+ if isinstance(spec, MambaSpec):
370
+ max_num_blocks = (
371
+ max_num_blocks if self.cache_config.enable_prefix_caching else 1
372
+ ) + spec.num_speculative_blocks
373
+ max_num_blocks_per_group.append(max_num_blocks)
374
+
375
self.block_tables = BlockTables(
376
block_sizes=block_sizes,
377
max_num_reqs=self.max_num_reqs,
378
max_num_batched_tokens=self.max_num_tokens,
359
- max_model_len=block_table_max_model_len,
379
+ max_num_blocks_per_group=max_num_blocks_per_group,
380
device=self.device,
381
cp_size=self.dcp_size,
382
cp_rank=self.dcp_rank,
@@ -722,6 +742,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
742
743
# Get the number of draft tokens for each request.
744
draft_tokens = scheduler_output.scheduled_spec_decode_tokens
745
+ num_draft_tokens_per_req = None
746
if not draft_tokens:
747
# No draft token scheduled (common case).
748
total_num_draft_tokens = 0
@@ -735,15 +756,15 @@ class GPUModelRunner(LoRAModelRunnerMixin):
756
num_reqs, dtype=torch.int32, device=self.device
757
)
758
else:
738
- num_draft_tokens = np.fromiter(
759
+ num_draft_tokens_per_req = np.fromiter(
760
(len(draft_tokens.get(req_id, ())) for req_id in req_ids),
761
dtype=np.int32,
762
count=num_reqs,
763
)
743
- total_num_draft_tokens = int(num_draft_tokens.sum())
764
+ total_num_draft_tokens = int(num_draft_tokens_per_req.sum())
765
total_num_logits = num_reqs + total_num_draft_tokens
766
746
- num_logits = num_draft_tokens + 1
767
+ num_logits = num_draft_tokens_per_req + 1
768
cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32)
769
cu_num_logits_np[0] = 0
770
np.cumsum(num_logits, out=cu_num_logits_np[1:])
@@ -766,9 +787,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
787
async_copy_to_gpu(query_start_loc_np, out=self.input_buffers.query_start_loc)
788
query_start_loc_np = query_start_loc_np[: num_reqs_padded + 1]
789
query_start_loc = self.input_buffers.query_start_loc[: num_reqs_padded + 1]
790
+ is_prefilling_np = (
791
+ self.req_states.num_computed_prefill_tokens[idx_mapping_np]
792
+ < self.req_states.prefill_len.np[idx_mapping_np]
793
+ )
794
795
# Get prefill tokens if any.
771
- if self.req_states.any_prefills(idx_mapping_np):
796
+ if np.any(is_prefilling_np):
797
prepare_prefill_inputs(
798
self.input_buffers.input_ids,
799
self.req_states.next_prefill_tokens,
@@ -836,11 +861,13 @@ class GPUModelRunner(LoRAModelRunnerMixin):
861
num_tokens=num_tokens,
862
num_tokens_after_padding=num_tokens_after_padding,
863
num_draft_tokens=total_num_draft_tokens,
864
+ num_draft_tokens_per_req=num_draft_tokens_per_req,
865
query_start_loc=query_start_loc,
866
query_start_loc_np=query_start_loc_np,
867
seq_lens=seq_lens,
868
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
869
dcp_local_seq_lens=dcp_local_seq_lens,
870
+ is_prefilling_np=is_prefilling_np,
871
input_ids=self.input_buffers.input_ids[:num_tokens_after_padding],
872
positions=self.input_buffers.positions[:num_tokens_after_padding],
873
logits_indices=logits_indices,
@@ -958,6 +985,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
985
input_batch.num_scheduled_tokens
986
)
987
988
+ self.model_state.postprocess_state(input_batch, num_sampled)
989
+
990
@torch.inference_mode()
991
def execute_model(
992
self,
vllm/v1/worker/gpu/model_states/__init__.py
+5
@@ -18,6 +18,11 @@ def init_model_state(
18
19
return WhisperModelState(vllm_config, model, encoder_cache, device)
20
21
+ if vllm_config.model_config.is_hybrid:
22
+ from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
23
+
24
+ return MambaHybridModelState(vllm_config, model, encoder_cache, device)
25
+
26
from vllm.v1.worker.gpu.model_states.default import DefaultModelState
27
28
return DefaultModelState(vllm_config, model, encoder_cache, device)
vllm/v1/worker/gpu/model_states/default.py
+1
@@ -194,5 +194,6 @@ class DefaultModelState(ModelState):
194
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
195
dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
196
positions=input_batch.positions,
197
+ for_cudagraph_capture=for_capture,
198
)
199
return attn_metadata
vllm/v1/worker/gpu/model_states/interface.py
+25
@@ -17,6 +17,24 @@ from vllm.v1.worker.gpu.states import RequestState
17
from vllm.v1.worker.utils import AttentionGroup
18
19
20
+class ModelSpecificAttnMetadata:
21
+ """Base class for model-specific attention metadata."""
22
+
23
+ def get_extra_common_attn_kwargs(
24
+ self,
25
+ kv_cache_group_id: int,
26
+ num_reqs: int,
27
+ ) -> dict[str, Any]:
28
+ return {}
29
+
30
+ def get_extra_attn_kwargs(
31
+ self,
32
+ attn_metadata_builder: Any,
33
+ num_reqs: int,
34
+ ) -> dict[str, Any]:
35
+ return {}
36
+
37
+
38
class ModelState(ABC):
39
@abstractmethod
40
def __init__(
@@ -38,6 +56,13 @@ class ModelState(ABC):
56
def apply_staged_writes(self) -> None:
57
return None
58
59
+ def postprocess_state(
60
+ self,
61
+ input_batch: InputBatch,
62
+ num_sampled: torch.Tensor,
63
+ ) -> None:
64
+ return None
65
+
66
@abstractmethod
67
def get_mm_embeddings(
68
self,
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
new
+150
@@ -0,0 +1,150 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+from typing import Any
5
+
6
+import numpy as np
7
+import torch
8
+import torch.nn as nn
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.config.compilation import CUDAGraphMode
12
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
13
+from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
14
+from vllm.v1.kv_cache_interface import KVCacheConfig
15
+from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
16
+from vllm.v1.worker.gpu.input_batch import InputBatch
17
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
18
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
19
+from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
20
+from vllm.v1.worker.utils import AttentionGroup
21
+
22
+
23
+@dataclass
24
+class MambaHybridAttnMetadata(ModelSpecificAttnMetadata):
25
+ is_prefilling: torch.Tensor
26
+ num_accepted_tokens: torch.Tensor | None = None
27
+ num_decode_draft_tokens_cpu: torch.Tensor | None = None
28
+
29
+ def get_extra_common_attn_kwargs(
30
+ self,
31
+ kv_cache_group_id: int,
32
+ num_reqs: int,
33
+ ) -> dict[str, Any]:
34
+ return {"is_prefilling": self.is_prefilling[:num_reqs]}
35
+
36
+ def get_extra_attn_kwargs(
37
+ self,
38
+ attn_metadata_builder: Any,
39
+ num_reqs: int,
40
+ ) -> dict[str, Any]:
41
+ if not isinstance(
42
+ attn_metadata_builder,
43
+ (Mamba2AttentionMetadataBuilder, GDNAttentionMetadataBuilder),
44
+ ):
45
+ return {}
46
+ return {
47
+ "num_accepted_tokens": None
48
+ if self.num_accepted_tokens is None
49
+ else self.num_accepted_tokens[:num_reqs],
50
+ "num_decode_draft_tokens_cpu": None
51
+ if self.num_decode_draft_tokens_cpu is None
52
+ else self.num_decode_draft_tokens_cpu[:num_reqs],
53
+ }
54
+
55
+
56
+class MambaHybridModelState(DefaultModelState):
57
+ """Model state for hybrid attention + Mamba / linear-attention models."""
58
+
59
+ def __init__(
60
+ self,
61
+ vllm_config: VllmConfig,
62
+ model: nn.Module,
63
+ encoder_cache: EncoderCache | None,
64
+ device: torch.device,
65
+ ) -> None:
66
+ super().__init__(vllm_config, model, encoder_cache, device)
67
+ self.num_accepted_tokens_gpu = torch.ones(
68
+ self.max_num_reqs, dtype=torch.int32, device=self.device
69
+ )
70
+
71
+ def prepare_attn(
72
+ self,
73
+ input_batch: InputBatch,
74
+ cudagraph_mode: CUDAGraphMode,
75
+ block_tables: tuple[torch.Tensor, ...],
76
+ slot_mappings: torch.Tensor,
77
+ attn_groups: list[list[AttentionGroup]],
78
+ kv_cache_config: KVCacheConfig,
79
+ for_capture: bool = False,
80
+ ) -> dict[str, Any]:
81
+ if cudagraph_mode == CUDAGraphMode.FULL:
82
+ num_reqs = input_batch.num_reqs_after_padding
83
+ num_tokens = input_batch.num_tokens_after_padding
84
+ else:
85
+ num_reqs = input_batch.num_reqs
86
+ num_tokens = input_batch.num_tokens
87
+ query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
88
+ max_query_len = input_batch.num_scheduled_tokens.max().item()
89
+
90
+ is_prefilling = torch.zeros(num_reqs, dtype=torch.bool, device="cpu")
91
+ is_prefilling[: input_batch.num_reqs] = torch.from_numpy(
92
+ input_batch.is_prefilling_np
93
+ )
94
+ # During CUDAGraph capture, num_decode_draft_tokens_cpu and num_accepted_tokens
95
+ # are created by attn_metadata_builder.build_for_cudagraph_capture, so we only
96
+ # compute them during actual (non-capture) forward execution.
97
+ num_accepted_tokens = None
98
+ num_decode_draft_tokens_cpu = None
99
+ if not for_capture:
100
+ num_accepted_tokens = self.num_accepted_tokens_gpu.new_ones(num_reqs)
101
+ num_accepted_tokens[: input_batch.num_reqs] = self.num_accepted_tokens_gpu[
102
+ input_batch.idx_mapping
103
+ ]
104
+
105
+ # GDN uses >= 0 to select spec-decode rows, so non-decode rows
106
+ # need the -1 sentinel rather than a raw zero draft count.
107
+ num_decode_draft_tokens_np = np.full(num_reqs, -1, dtype=np.int32)
108
+ if input_batch.num_draft_tokens_per_req is not None:
109
+ spec_decode_mask = (
110
+ input_batch.num_draft_tokens_per_req > 0
111
+ ) & ~input_batch.is_prefilling_np
112
+ num_decode_draft_tokens_np[: input_batch.num_reqs] = np.where(
113
+ spec_decode_mask,
114
+ input_batch.num_draft_tokens_per_req,
115
+ -1,
116
+ )
117
+ num_decode_draft_tokens_cpu = torch.from_numpy(num_decode_draft_tokens_np)
118
+
119
+ mamba_attn_metadata = MambaHybridAttnMetadata(
120
+ is_prefilling=is_prefilling,
121
+ num_accepted_tokens=num_accepted_tokens,
122
+ num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
123
+ )
124
+ return build_attn_metadata(
125
+ attn_groups=attn_groups,
126
+ num_reqs=num_reqs,
127
+ num_tokens=num_tokens,
128
+ query_start_loc_gpu=input_batch.query_start_loc,
129
+ query_start_loc_cpu=query_start_loc_cpu,
130
+ max_query_len=max_query_len,
131
+ seq_lens=input_batch.seq_lens,
132
+ max_seq_len=self.max_model_len,
133
+ block_tables=block_tables,
134
+ slot_mappings=slot_mappings,
135
+ kv_cache_config=kv_cache_config,
136
+ dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
137
+ model_specific_attn_metadata=mamba_attn_metadata,
138
+ for_cudagraph_capture=for_capture,
139
+ )
140
+
141
+ def postprocess_state(
142
+ self,
143
+ input_batch: InputBatch,
144
+ num_sampled: torch.Tensor,
145
+ ) -> None:
146
+ # Chunked prefill does not sample a token, so num_sampled can be 0.
147
+ # Mamba treats num_accepted_tokens=1 as the neutral non-spec value.
148
+ self.num_accepted_tokens_gpu[input_batch.idx_mapping] = torch.clamp(
149
+ num_sampled, min=1
150
+ )
vllm/v1/worker/gpu/model_states/whisper.py
+28
-4
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
from typing import Any
5
6
import numpy as np
@@ -13,11 +14,33 @@ from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
14
from vllm.v1.worker.gpu.input_batch import InputBatch
15
from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
16
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
16
-from vllm.v1.worker.gpu.model_states.interface import ModelState
17
+from vllm.v1.worker.gpu.model_states.interface import (
18
+ ModelSpecificAttnMetadata,
19
+ ModelState,
20
+)
21
from vllm.v1.worker.gpu.states import RequestState
22
from vllm.v1.worker.utils import AttentionGroup
23
24
25
+@dataclass
26
+class WhisperAttnMetadata(ModelSpecificAttnMetadata):
27
+ encoder_seq_lens: dict[int, tuple[torch.Tensor, np.ndarray]]
28
+
29
+ def get_extra_common_attn_kwargs(
30
+ self,
31
+ kv_cache_group_id: int,
32
+ num_reqs: int,
33
+ ) -> dict[str, Any]:
34
+ encoder_seq_lens = self.encoder_seq_lens.get(kv_cache_group_id)
35
+ if encoder_seq_lens is None:
36
+ return {}
37
+ encoder_seq_lens_gpu, encoder_seq_lens_cpu = encoder_seq_lens
38
+ return {
39
+ "encoder_seq_lens": encoder_seq_lens_gpu[:num_reqs],
40
+ "encoder_seq_lens_cpu": encoder_seq_lens_cpu[:num_reqs],
41
+ }
42
+
43
+
44
class WhisperModelState(ModelState):
45
def __init__(
46
self,
@@ -111,8 +134,8 @@ class WhisperModelState(ModelState):
134
else:
135
num_reqs = input_batch.num_reqs
136
num_tokens = input_batch.num_tokens
114
- encoder_seq_lens = self._get_encoder_seq_lens(
115
- input_batch.req_ids, attn_groups, for_capture
137
+ whisper_attn_metadata = WhisperAttnMetadata(
138
+ self._get_encoder_seq_lens(input_batch.req_ids, attn_groups, for_capture)
139
)
140
141
query_start_loc_cpu = torch.from_numpy(input_batch.query_start_loc_np)
@@ -136,7 +159,8 @@ class WhisperModelState(ModelState):
159
kv_cache_config=kv_cache_config,
160
seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
161
dcp_local_seq_lens=input_batch.dcp_local_seq_lens,
139
- encoder_seq_lens=encoder_seq_lens,
162
+ model_specific_attn_metadata=whisper_attn_metadata,
163
+ for_cudagraph_capture=for_capture,
164
)
165
return attn_metadata
166
vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py
+1
@@ -106,6 +106,7 @@ class DecodeEagleCudaGraphManager(EagleCudaGraphManagerBase):
106
block_tables,
107
attn_groups,
108
kv_cache_config,
109
+ skip_attn=(desc.cg_mode == CUDAGraphMode.PIECEWISE),
110
)
111
attn_metadata, slot_mappings = attn_state
112