mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-22 21:50:15 +00:00
[Bugfix] Fix Gemma4 MTP block_table batch_size mismatch under concurrent load (#43982)
Signed-off-by: Dmytro Kuntso <[email protected]> Co-authored-by: Dmytro Kuntso <[email protected]>
This commit is contained in:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user