[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:
Dima
2026-06-03 17:11:10 -07:00
committed by GitHub
co-authored by Dmytro Kuntso
parent bdbf08fc02
commit 128adabfe0
+6 -1
View File
@@ -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(