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

model: support Kimi-K2.5 (#17789)

Browse Source 
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Yuhao Yang
2026-01-27 10:57:00 +08:00
committed by GitHub
parent 0519b0935f
commit 479ab7a4e7
11 changed files with 1053 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/srt/layers/attention/vision.py
View File

@@ -13,7 +13,11 @@ from einops import rearrange
from sglang.jit_kernel.norm import can_use_fused_inplace_qknorm as can_use_jit_qk_norm
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import get_attention_tp_rank, get_attention_tp_size
from sglang.srt.layers.dp_attention import (
get_attention_tp_group,
get_attention_tp_rank,
get_attention_tp_size,
)
from sglang.srt.models.utils import apply_qk_norm
from sglang.srt.utils import (
get_bool_env_var,
@@ -692,6 +696,7 @@ class VisionAttention(nn.Module):
quant_config=quant_config,
tp_rank=self.tp_rank,
tp_size=self.tp_size,
reduce_results=False,
prefix=add_prefix("proj", prefix),
)
self.aux_stream = aux_stream
@@ -914,6 +919,8 @@ class VisionAttention(nn.Module):
# [b, s, h * head_size] --> [b, s, h * head_size]
output, _ = self.proj(output)
if self.tp_size > 1:
output = get_attention_tp_group().all_reduce(output)
else:
# [b * s, h, head_size] --> [s, b, h * head_size]
context_layer = rearrange(
@@ -922,6 +929,8 @@ class VisionAttention(nn.Module):
# [s, b, h * head_size] --> [s, b, h * head_size]
output, _ = self.proj(context_layer)
if self.tp_size > 1:
output = get_attention_tp_group().all_reduce(output)
# [s, b, h * head_size] --> [b, s, h * head_size]
output = output.view(bsz, s, -1)