From 1b75d0d1a979a324a24daf6a10fd8491b579f1c9 Mon Sep 17 00:00:00 2001 From: Shu Wang Date: Fri, 27 Feb 2026 13:35:45 -0600 Subject: [PATCH] Fix BatchMLAPagedAttentionWrapper query/qo_inptr mismatch for EAGLE (#15601) --- python/sglang/srt/layers/attention/trtllm_mla_backend.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index e5489ce7a..b74124da2 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -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()