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

Support Kimi Linear (#12469)

Browse Source 
Co-authored-by: yizhang2077 <1109276519@qq.com>
This commit is contained in:
Ke Bao
2025-11-01 05:03:35 +08:00
committed by GitHub
parent 30ad107028
commit a4bf5c6ad2
18 changed files with 2847 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -29,7 +29,12 @@ from typing import Callable, List, Optional, Tuple, Union
import torch
import torch.distributed as dist
from sglang.srt.configs import FalconH1Config, NemotronHConfig, Qwen3NextConfig
from sglang.srt.configs import (
FalconH1Config,
KimiLinearConfig,
NemotronHConfig,
Qwen3NextConfig,
)
from sglang.srt.configs.device_config import DeviceConfig
from sglang.srt.configs.load_config import LoadConfig, LoadFormat
from sglang.srt.configs.model_config import (
@@ -1358,9 +1363,16 @@ class ModelRunner:
return config
return None
@property
def kimi_linear_config(self):
config = self.model_config.hf_config
if isinstance(config, KimiLinearConfig):
return config
return None
@property
def mambaish_config(self):
return self.mamba2_config or self.hybrid_gdn_config
return self.mamba2_config or self.hybrid_gdn_config or self.kimi_linear_config
def set_num_token_hybrid(self):
if (
@@ -1691,7 +1703,7 @@ class ModelRunner:
end_layer=self.end_layer,
index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config),
)
elif self.use_mla_backend:
elif self.use_mla_backend and not self.mambaish_config:
assert not is_nsa_model
self.token_to_kv_pool = MLATokenToKVPool(
self.max_total_num_tokens,
@@ -1735,6 +1747,12 @@ class ModelRunner:
device=self.device,
)
elif config := self.mambaish_config:
extra_args = {}
if self.use_mla_backend:
extra_args = {
"kv_lora_rank": self.model_config.kv_lora_rank,
"qk_rope_head_dim": self.model_config.qk_rope_head_dim,
}
self.token_to_kv_pool = HybridLinearKVPool(
page_size=self.page_size,
size=self.max_total_num_tokens,
@@ -1750,6 +1768,8 @@ class ModelRunner:
enable_kvcache_transpose=False,
device=self.device,
mamba_pool=self.req_to_token_pool.mamba_pool,
use_mla=self.use_mla_backend,
**extra_args,
)
else:
self.token_to_kv_pool = MHATokenToKVPool(