Fix BatchMLAPagedAttentionWrapper query/qo_inptr mismatch for EAGLE (#15601)

This commit is contained in:
Shu Wang
2026-02-27 13:35:45 -06:00
committed by GitHub
parent 6a1480ce45
commit 1b75d0d1a9

View File

@@ -616,6 +616,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
):
bs = forward_batch.batch_size
self.forward_decode_metadata = TRTLLMMLADecodeMetadata()
# This is necessary because the backend instance persists across forward passes,
# and forward_prefill_metadata from a previous regular extend call could still be set.
if (
forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend(include_v2=True)
):
self.forward_prefill_metadata = None
# Get maximum sequence length.
if getattr(forward_batch, "seq_lens_cpu", None) is not None:
max_seq = forward_batch.seq_lens_cpu.max().item()