EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Support MHA with chunked prefix cache for flashinfer/flashmla backend, support page size > 1 for MHA chunked prefix (#8616)

Browse Source 
Co-authored-by: xuyongfei.xyf <xuyongfei.xyf@antgroup.com>
This commit is contained in:
Yongfei Xu
2025-08-22 09:19:44 +08:00
committed by GitHub
parent 704ced1b2e
commit 9708d353b7
6 changed files with 184 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -241,6 +241,9 @@ class ForwardBatch:
prefix_chunk_num_tokens: Optional[List[int]] = None
# KV Indices for each chunk
prefix_chunk_kv_indices: Optional[List[torch.Tensor]] = None
# For MLA chunked prefix cache used in chunked prefill
# Tell attention backend whether lse needs to be returned
mha_return_lse: Optional[bool] = None
# For multimodal
mm_inputs: Optional[List[MultimodalInputs]] = None

3
python/sglang/srt/model_executor/model_runner.py
View File

@@ -518,9 +518,6 @@ class ModelRunner:
if not self.use_mla_backend:
server_args.disable_chunked_prefix_cache = True
elif self.page_size > 1:
logger.info("Disable chunked prefix cache when page size > 1.")
server_args.disable_chunked_prefix_cache = True
if not server_args.disable_chunked_prefix_cache:
logger.info("Chunked prefix cache is turned on.")