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

Allow Optional key/value in unified_attention_with_output split-op (MLA absorb fix) (#26515)

Browse Source 
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-28 01:04:07 -07:00
committed by GitHub
co-authored by Claude Opus 4.7
parent b429a30428
commit 8ca09a30f1
1 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/radix_attention.py
+4 -2
View File
@@ -151,8 +151,8 @@ class RadixAttention(nn.Module):
@register_split_op() @register_split_op()
def unified_attention_with_output( def unified_attention_with_output(
query: torch.Tensor, query: torch.Tensor,
key: torch.Tensor, key: Optional[torch.Tensor],
value: torch.Tensor, value: Optional[torch.Tensor],
output: torch.Tensor, output: torch.Tensor,
save_kv_cache: bool, save_kv_cache: bool,
layer_id: int, layer_id: int,
@@ -174,7 +174,9 @@ def unified_attention_with_output(
real_num_tokens = forward_batch.num_token_non_padded_cpu real_num_tokens = forward_batch.num_token_non_padded_cpu
query = query[:real_num_tokens] query = query[:real_num_tokens]
if key is not None:
key = key[:real_num_tokens] key = key[:real_num_tokens]
if value is not None:
value = value[:real_num_tokens] value = value[:real_num_tokens]
kwargs = {} kwargs = {}