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 2, 01:00 UTC
[Bugfix] Fix Gemma4 MTP block_table batch_size mismatch under concurrent load (#43982)
Signed-off-by: Dmytro Kuntso <dkuntso@amazon.co.uk> Co-authored-by: Dmytro Kuntso <dkuntso@amazon.co.uk>
Dima committed
Jun 4, 2026 at 01:11 UTC
128adabfe0fe15dd40838d12166b3a1ca48b2b09
1 file changed
+6
-1
vllm/v1/spec_decode/gemma4.py
+6
-1
@@ -81,11 +81,16 @@ class Gemma4Proposer(SpecDecodeBaseProposer):
81
"""
82
per_group_attn_metadata: list[object] = []
83
per_layer_attn_metadata: dict[str, object] = {}
84
+ batch_size = common_attn_metadata.batch_size()
85
for attn_group in self.draft_attn_groups:
86
gid = attn_group.kv_cache_group_id
87
if gid in self._per_group_block_tables:
88
cm = copy(common_attn_metadata)
88
- cm.block_table_tensor = self._per_group_block_tables[gid]
89
+ # Slice to actual batch size to match cu_seqlens_q dimension.
90
+ # The stored block tables may be padded (num_reqs_padded) from
91
+ # the target forward pass, but the drafter operates on the
92
+ # unpadded batch.
93
+ cm.block_table_tensor = self._per_group_block_tables[gid][:batch_size]
94
else:
95
cm = common_attn_metadata
96
attn_metadata = attn_group.get_metadata_builder().build_for_drafting(