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

fix(attn): delegate init_mha_chunk_metadata in HybridLinearAttnBackend (#27316)

Browse Source 
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-06-04 17:44:23 -07:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 7425bebb6c
commit 7dc7376697
4 changed files with 219 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+10
View File
@@ -807,6 +807,16 @@ class HybridLinearAttnBackend(AttentionBackend):
for attn_backend in self.attn_backend_list:
attn_backend.init_forward_metadata(forward_batch)
def init_mha_chunk_metadata(
self, forward_batch: ForwardBatch, disable_flashinfer_ragged: bool = False
):
# Hybrid MLA models (Ring/Ling, Kimi-Linear) resolve this via
# get_attn_backend(), which returns this wrapper; delegate to the
# full-attn backend so its chunked/one-shot prefill metadata is planned.
init = getattr(self.full_attn_backend, "init_mha_chunk_metadata", None)
if init is not None:
init(forward_batch, disable_flashinfer_ragged)
def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
for attn_backend in self.attn_backend_list:
attn_backend.init_cuda_graph_state(max_bs, max_num_tokens)