Unverified Commit e94cfd51 authored by Vadim Gimpelson's avatar Vadim Gimpelson Committed by GitHub
Browse files

[BUG] Qwen3-next MTP. Fix attn metadata build bug (#26564)


Signed-off-by: default avatarVadim Gimpelson <vadim.gimpelson@gmail.com>
parent 7c12763b
......@@ -41,7 +41,6 @@ from vllm.v1.sample.metadata import SamplingMetadata
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
from vllm.v1.utils import CpuGpuBuffer
from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
from vllm.v1.worker.ubatching import dbo_current_ubatch_id
logger = init_logger(__name__)
......@@ -234,11 +233,11 @@ class EagleProposer:
assert self.runner is not None
# FIXME: need to consider multiple kv_cache_groups
ubatch_id = dbo_current_ubatch_id()
attn_metadata_builder = self.runner.attn_groups[0][0].metadata_builders[
ubatch_id
]
if self.attn_metadata_builder is None:
attn_metadata_builder = self._get_attention_metadata_builder()
else:
attn_metadata_builder = self.attn_metadata_builder
attn_metadata = attn_metadata_builder.build_for_drafting(
common_attn_metadata=common_attn_metadata, draft_index=0
)
......@@ -1076,7 +1075,7 @@ class EagleProposer:
inputs_embeds=inputs_embeds,
)
def _get_attention_metadata_builder(self) -> list[AttentionMetadataBuilder]:
def _get_attention_metadata_builder(self) -> AttentionMetadataBuilder:
"""Find and return the attention metadata builders for EAGLE layers.
Returns:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment