minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Skip init_mha_chunk_metadata in trtllm_mla when not needed (#26017)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-05-22 16:34:16 -07:00
committed by GitHub
parent cadfa2d025
commit c112f7623a
1 changed files with 8 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+8
View File
@@ -577,6 +577,14 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
"""Get the fill value for sequence lengths in CUDA graph.""" """Get the fill value for sequence lengths in CUDA graph."""
return 1 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): def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Initialize the metadata for a forward pass.""" """Initialize the metadata for a forward pass."""
# Delegate to parent for non-decode modes. # Delegate to parent for non-decode modes.