diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 253bb92d9..1ce665583 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -848,10 +848,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): page_table = self._get_layer_page_table(layer, forward_batch) - if ( - forward_batch.forward_mode.is_target_verify() - or forward_batch.forward_mode.is_draft_extend_v2() - ): + if forward_batch.forward_mode.is_target_verify(): o = flashinfer.decode.trtllm_batch_decode_with_kv_cache( query=q, kv_cache=kv_cache,