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 12, 13:30 UTC
[Core] Generalize Encoder-Decoder `seq_lens` computation to avoid Whisper hardcoded logic (#29268)
Signed-off-by: NickLucche <nlucches@redhat.com> Co-authored-by: Ekagra Ranjan <3116519+ekagra-ranjan@users.noreply.github.com>
Nicolò Lucchesi committed
Nov 25, 2025 at 12:32 UTC
798e87db5c21e017f77edcdf7e50b1ac79d65c54
3 files changed
+51
-34
vllm/attention/layers/cross_attention.py
+22
-22
@@ -25,15 +25,6 @@ from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec
25
logger = init_logger(__name__)
26
27
28
-def _get_max_encoder_len(vllm_config: "VllmConfig") -> int:
29
- """Gets the max number of encoder input tokens from the config."""
30
- sc = vllm_config.scheduler_config
31
- assert sc and isinstance(sc.max_num_encoder_input_tokens, int), (
32
- "max_num_encoder_input_tokens must be int for enc-dec models"
33
- )
34
- return sc.max_num_encoder_input_tokens
35
-
36
-
28
def _get_cross_slot_mapping(
29
encoder_seq_lens: np.ndarray,
30
block_table_tensor: torch.Tensor,
@@ -93,23 +84,32 @@ def create_cross_attention_backend(
84
) -> AttentionMetadata:
85
new_metadata = copy(common_attn_metadata)
86
new_metadata.causal = False
96
- max_encoder_len = _get_max_encoder_len(self.vllm_config)
87
+ max_encoder_len = int(new_metadata.encoder_seq_lens_cpu.max())
88
new_metadata.max_seq_len = max_encoder_len
98
-
99
- new_metadata.seq_lens = torch.full(
100
- (new_metadata.num_reqs,),
101
- max_encoder_len,
102
- dtype=torch.int32,
103
- device=self.device,
89
+ # Any computed tokens indicated decode step>1 (no chunked prefill)
90
+ num_cache_decodes = (
91
+ (common_attn_metadata.num_computed_tokens_cpu > 0).sum().item()
92
)
105
- new_metadata.seq_lens_cpu = torch.full(
106
- (new_metadata.num_reqs,),
107
- max_encoder_len,
108
- dtype=torch.int32,
109
- device="cpu",
93
+ if num_cache_decodes > 0:
94
+ # CrossAttn KV cache has already been populated on first decoder step,
95
+ # skip slot_mapping calculation for requests that do not need
96
+ # reshape_and_cache.
97
+ num_tokens = common_attn_metadata.num_computed_tokens_cpu.numpy()
98
+ new_metadata.encoder_seq_lens_cpu = np.where(
99
+ num_tokens > 0, 0, new_metadata.encoder_seq_lens_cpu
100
+ )
101
+
102
+ # seq_lens is provided by model runner: initial encoder input length is
103
+ # needed here to know how many tokens to attend to from the cached
104
+ # cross-attention KV cache.
105
+ new_metadata.seq_lens = common_attn_metadata.encoder_seq_lens
106
+ new_metadata.seq_lens_cpu = torch.from_numpy(
107
+ common_attn_metadata.encoder_seq_lens_cpu
108
)
109
+
110
+ # NOTE (NickLucche) use `new_metadata` instead of `common_*` (initial) here
111
new_metadata.slot_mapping = _get_cross_slot_mapping(
112
- new_metadata.encoder_seq_lens,
112
+ new_metadata.encoder_seq_lens_cpu,
113
new_metadata.block_table_tensor,
114
self.kv_cache_spec,
115
self.device,
vllm/v1/attention/backends/utils.py
+2
-1
@@ -89,7 +89,8 @@ class CommonAttentionMetadata:
89
num_logits_indices: int | None = None
90
91
# Needed by CrossAttentionBuilder
92
- encoder_seq_lens: np.ndarray | None = None
92
+ encoder_seq_lens: torch.Tensor | None = None
93
+ encoder_seq_lens_cpu: np.ndarray | None = None
94
95
dcp_local_seq_lens: torch.Tensor | None = None
96
dcp_local_seq_lens_cpu: torch.Tensor | None = None
vllm/v1/worker/gpu_model_runner.py
+27
-11
@@ -475,6 +475,7 @@ class GPUModelRunner(
475
self.max_num_reqs + 1, dtype=torch.int32
476
)
477
self.seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
478
+ self.encoder_seq_lens = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
479
if self.dcp_world_size > 1:
480
self.dcp_local_seq_lens = self._make_buffer(
481
self.max_num_reqs, dtype=torch.int32
@@ -1202,21 +1203,35 @@ class GPUModelRunner(
1203
1204
def _get_encoder_seq_lens(
1205
self,
1205
- scheduled_encoder_inputs: dict[str, list[int]],
1206
+ num_scheduled_tokens: dict[str, int],
1207
kv_cache_spec: KVCacheSpec,
1208
num_reqs: int,
1208
- ) -> np.ndarray | None:
1209
+ ) -> tuple[torch.Tensor | None, np.ndarray | None]:
1210
if not isinstance(kv_cache_spec, CrossAttentionSpec):
1210
- return None
1211
+ return None, None
1212
1213
# Build encoder_seq_lens array mapping request indices to
1214
# encoder lengths for inputs scheduled in this batch
1214
- encoder_seq_lens = np.zeros(num_reqs, dtype=np.int32)
1215
- for req_id in scheduled_encoder_inputs:
1215
+ for req_id in num_scheduled_tokens:
1216
req_index = self.input_batch.req_id_to_index[req_id]
1217
- encoder_seq_lens[req_index] = self.max_encoder_len
1217
+ req_state = self.requests[req_id]
1218
+ if req_state.mm_features is None:
1219
+ self.encoder_seq_lens.np[req_index] = 0
1220
+ continue
1221
+
1222
+ # Get the total number of encoder input tokens for running encoder requests
1223
+ # whether encoding is finished or not so that cross-attention knows how
1224
+ # many encoder tokens to attend to.
1225
+ encoder_input_tokens = sum(
1226
+ feature.mm_position.length for feature in req_state.mm_features
1227
+ )
1228
+ self.encoder_seq_lens.np[req_index] = encoder_input_tokens
1229
+
1230
+ self.encoder_seq_lens.copy_to_gpu(num_reqs)
1231
+ encoder_seq_lens = self.encoder_seq_lens.gpu[:num_reqs]
1232
+ encoder_seq_lens_cpu = self.encoder_seq_lens.np[:num_reqs]
1233
1219
- return encoder_seq_lens
1234
+ return encoder_seq_lens, encoder_seq_lens_cpu
1235
1236
def _prepare_inputs(
1237
self,
@@ -1482,7 +1497,7 @@ class GPUModelRunner(
1497
logits_indices: torch.Tensor | None = None,
1498
use_spec_decode: bool = False,
1499
for_cudagraph_capture: bool = False,
1485
- scheduled_encoder_inputs: dict[str, list[int]] | None = None,
1500
+ num_scheduled_tokens: dict[str, int] | None = None,
1501
cascade_attn_prefix_lens: list[list[int]] | None = None,
1502
) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]:
1503
"""
@@ -1547,8 +1562,8 @@ class GPUModelRunner(
1562
for kv_cache_gid, kv_cache_group in enumerate(
1563
self.kv_cache_config.kv_cache_groups
1564
):
1550
- encoder_seq_lens = self._get_encoder_seq_lens(
1551
- scheduled_encoder_inputs or {},
1565
+ encoder_seq_lens, encoder_seq_lens_cpu = self._get_encoder_seq_lens(
1566
+ num_scheduled_tokens or {},
1567
kv_cache_group.kv_cache_spec,
1568
num_reqs,
1569
)
@@ -1591,6 +1606,7 @@ class GPUModelRunner(
1606
num_logits_indices=num_logits_indices,
1607
causal=True,
1608
encoder_seq_lens=encoder_seq_lens,
1609
+ encoder_seq_lens_cpu=encoder_seq_lens_cpu,
1610
dcp_local_seq_lens=dcp_local_seq_lens,
1611
dcp_local_seq_lens_cpu=dcp_local_seq_lens_cpu,
1612
)
@@ -2828,7 +2844,7 @@ class GPUModelRunner(
2844
ubatch_slices=ubatch_slices,
2845
logits_indices=logits_indices,
2846
use_spec_decode=use_spec_decode,
2831
- scheduled_encoder_inputs=scheduler_output.scheduled_encoder_inputs,
2847
+ num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
2848
cascade_attn_prefix_lens=cascade_attn_prefix_lens,
2849
)
2850
)