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 14, 05:30 UTC
[BUG] Qwen3-next MTP. Fix attn metadata build bug (#26564)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Oct 10, 2025 at 22:59 UTC
e94cfd51da5f07527bb6ce3305825ed2e61c5aca
1 file changed
+6
-7
vllm/v1/spec_decode/eagle.py
+6
-7
@@ -41,7 +41,6 @@ from vllm.v1.sample.metadata import SamplingMetadata
41
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
42
from vllm.v1.utils import CpuGpuBuffer
43
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
44
-from vllm.v1.worker.ubatching import dbo_current_ubatch_id
44
45
logger = init_logger(__name__)
46
@@ -234,11 +233,11 @@ class EagleProposer:
233
234
assert self.runner is not None
235
237
- # FIXME: need to consider multiple kv_cache_groups
238
- ubatch_id = dbo_current_ubatch_id()
239
- attn_metadata_builder = self.runner.attn_groups[0][0].metadata_builders[
240
- ubatch_id
241
- ]
236
+ if self.attn_metadata_builder is None:
237
+ attn_metadata_builder = self._get_attention_metadata_builder()
238
+ else:
239
+ attn_metadata_builder = self.attn_metadata_builder
240
+
241
attn_metadata = attn_metadata_builder.build_for_drafting(
242
common_attn_metadata=common_attn_metadata, draft_index=0
243
)
@@ -1076,7 +1075,7 @@ class EagleProposer:
1075
inputs_embeds=inputs_embeds,
1076
)
1077
1079
- def _get_attention_metadata_builder(self) -> list[AttentionMetadataBuilder]:
1078
+ def _get_attention_metadata_builder(self) -> AttentionMetadataBuilder:
1079
"""Find and return the attention metadata builders for EAGLE layers.
1080
1081
Returns: