From c112f7623a6b9cdb485046f9e02671e7d02e9104 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Fri, 22 May 2026 16:34:16 -0700 Subject: [PATCH] Skip init_mha_chunk_metadata in trtllm_mla when not needed (#26017) --- python/sglang/srt/layers/attention/trtllm_mla_backend.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 72b8e6d5f..695e739f1 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -577,6 +577,14 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): """Get the fill value for sequence lengths in CUDA graph.""" return 1 + def init_mha_chunk_metadata(self, forward_batch: "ForwardBatch") -> None: + has_prefix = any(forward_batch.extend_prefix_lens_cpu) + fallback_to_flashinfer_impl = ( + self.disable_chunked_prefix_cache and has_prefix + ) or is_in_piecewise_cuda_graph() + if fallback_to_flashinfer_impl: + super().init_mha_chunk_metadata(forward_batch) + def init_forward_metadata(self, forward_batch: ForwardBatch): """Initialize the metadata for a forward pass.""" # Delegate to parent for non-decode modes.