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, 11:00 UTC
[BugFix] Fix new nightly failures (#29578)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Nov 27, 2025 at 16:45 UTC
be493e0b3cfb5810d254e9845217878a39a4853b
2 files changed
+37
-1
vllm/v1/attention/backends/utils.py
+26
@@ -100,6 +100,32 @@ class CommonAttentionMetadata:
100
dcp_local_seq_lens_cpu: torch.Tensor | None = None
101
"""Sequence lengths of the local rank in decode context parallelism world"""
102
103
+ # TODO(lucas): remove once we have FULL-CG spec-decode support
104
+ def unpadded(
105
+ self, num_actual_tokens: int, num_actual_reqs: int
106
+ ) -> "CommonAttentionMetadata":
107
+ maybe_slice_reqs = lambda x: x[:num_actual_reqs] if x is not None else None
108
+ return CommonAttentionMetadata(
109
+ query_start_loc=self.query_start_loc[: num_actual_reqs + 1],
110
+ query_start_loc_cpu=self.query_start_loc_cpu[: num_actual_reqs + 1],
111
+ seq_lens=self.seq_lens[:num_actual_reqs],
112
+ seq_lens_cpu=self.seq_lens_cpu[:num_actual_reqs],
113
+ num_computed_tokens_cpu=self.num_computed_tokens_cpu[:num_actual_reqs],
114
+ num_reqs=num_actual_reqs,
115
+ num_actual_tokens=num_actual_tokens,
116
+ max_query_len=self.max_query_len,
117
+ max_seq_len=self.max_seq_len,
118
+ block_table_tensor=self.block_table_tensor[:num_actual_reqs],
119
+ slot_mapping=self.slot_mapping[:num_actual_tokens],
120
+ causal=self.causal,
121
+ logits_indices_padded=self.logits_indices_padded,
122
+ num_logits_indices=self.num_logits_indices,
123
+ encoder_seq_lens=maybe_slice_reqs(self.encoder_seq_lens),
124
+ encoder_seq_lens_cpu=maybe_slice_reqs(self.encoder_seq_lens_cpu),
125
+ dcp_local_seq_lens=maybe_slice_reqs(self.dcp_local_seq_lens),
126
+ dcp_local_seq_lens_cpu=maybe_slice_reqs(self.dcp_local_seq_lens_cpu),
127
+ )
128
+
129
130
def slice_query_start_locs(
131
query_start_loc: torch.Tensor,
vllm/v1/worker/gpu_model_runner.py
+11
-1
@@ -1551,7 +1551,7 @@ class GPUModelRunner(
1551
# Encoder-only layers do not have KV cache, so we need to
1552
# create a dummy block table and slot mapping for them.
1553
blk_table_tensor = torch.zeros(
1554
- (num_tokens_padded, 1),
1554
+ (num_reqs_padded, 1),
1555
dtype=torch.int32,
1556
device=self.device,
1557
)
@@ -1652,6 +1652,16 @@ class GPUModelRunner(
1652
for layer_name in attn_group.layer_names:
1653
attn_metadata[layer_name] = attn_metadata_i
1654
1655
+ if spec_decode_common_attn_metadata is not None and (
1656
+ num_reqs != num_reqs_padded or num_tokens != num_tokens_padded
1657
+ ):
1658
+ # Currently the drafter still only uses piecewise cudagraphs (and modifies
1659
+ # the attention metadata in directly), and therefore does not want to use
1660
+ # padded attention metadata.
1661
+ spec_decode_common_attn_metadata = (
1662
+ spec_decode_common_attn_metadata.unpadded(num_tokens, num_reqs)
1663
+ )
1664
+
1665
return attn_metadata, spec_decode_common_attn_metadata
1666
1667
def _compute_cascade_attn_prefix_lens(