Fix BatchMLAPagedAttentionWrapper query/qo_inptr mismatch for EAGLE (#15601)
This commit is contained in:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user