diff --git a/vllm/v1/spec_decode/gemma4.py b/vllm/v1/spec_decode/gemma4.py index b0a02774faf..7f67ae9f499 100644 --- a/vllm/v1/spec_decode/gemma4.py +++ b/vllm/v1/spec_decode/gemma4.py @@ -81,11 +81,16 @@ class Gemma4Proposer(SpecDecodeBaseProposer): """ per_group_attn_metadata: list[object] = [] per_layer_attn_metadata: dict[str, object] = {} + batch_size = common_attn_metadata.batch_size() for attn_group in self.draft_attn_groups: gid = attn_group.kv_cache_group_id if gid in self._per_group_block_tables: cm = copy(common_attn_metadata) - cm.block_table_tensor = self._per_group_block_tables[gid] + # Slice to actual batch size to match cu_seqlens_q dimension. + # The stored block tables may be padded (num_reqs_padded) from + # the target forward pass, but the drafter operates on the + # unpadded batch. + cm.block_table_tensor = self._per_group_block_tables[gid][:batch_size] else: cm = common_attn_metadata attn_metadata = attn_group.get_metadata_builder().build_for_drafting(