From 8ca09a30f18d8461169b84163ef9210dbf597972 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 28 May 2026 01:04:07 -0700 Subject: [PATCH] Allow Optional key/value in unified_attention_with_output split-op (MLA absorb fix) (#26515) Co-authored-by: Claude Opus 4.7 (1M context) --- python/sglang/srt/layers/radix_attention.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index c3409c8c2..f3697a090 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/python/sglang/srt/layers/radix_attention.py @@ -151,8 +151,8 @@ class RadixAttention(nn.Module): @register_split_op() def unified_attention_with_output( query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, + key: Optional[torch.Tensor], + value: Optional[torch.Tensor], output: torch.Tensor, save_kv_cache: bool, layer_id: int, @@ -174,8 +174,10 @@ def unified_attention_with_output( real_num_tokens = forward_batch.num_token_non_padded_cpu query = query[:real_num_tokens] - key = key[:real_num_tokens] - value = value[:real_num_tokens] + if key is not None: + key = key[:real_num_tokens] + if value is not None: + value = value[:real_num_tokens] kwargs = {} if q_rope is not None: