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

[Feature] Xiaomi MiMo-V2-Flash day0 support (#15207)

Browse Source 
Co-authored-by: 谢学扬 <xiexueyang@xiaomi.com>
Co-authored-by: tz <tangzhen3@xiaomi.com>
Co-authored-by: 李家乐 <lijiale10@xiaomi.com>
Co-authored-by: 张晨 <zhangchen50@xiaomi.com>
Co-authored-by: Shaohui Liu <liushaohui3@xiaomi.com>
Co-authored-by: 王晨 <wangchen77@xiaomi.com>
Co-authored-by: jiangzihan <jiangzihan@xiaomi.com>
Co-authored-by: xiexueyang <xyxie_wangyi@163.com>
Co-authored-by: Linghao Zhang <zhanglinghao@xiaomi.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
Co-authored-by: JoyFuture <35593546+JoyFuture@users.noreply.github.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
Co-authored-by: root <root@bj9-ml-g8h20e-k8s-slave106-20251106.alicn.idc.xiaomi.com>
This commit is contained in:
Yingchun Lai
2025-12-19 11:40:07 +08:00
committed by GitHub
co-authored by 谢学扬 tz 李家乐 张晨 Shaohui Liu 王晨 jiangzihan xiexueyang Linghao Zhang ispobock Liangsheng Yin JoyFuture Liangsheng Yin Qiaolin Yu root
parent a0985dd5e5
commit 160a06cab2
38 changed files with 5396 additions and 169 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/tp_worker.py
+29 -1
View File
@@ -216,6 +216,7 @@ class TpModelWorker(BaseTpWorker):
is_draft_worker: bool = False,
req_to_token_pool: Optional[ReqToTokenPool] = None,
token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None,
is_mtp_worker: bool = False,
):
# Parse args
self.tp_size = server_args.tp_size
@@ -223,6 +224,9 @@ class TpModelWorker(BaseTpWorker):
self.moe_ep_rank = moe_ep_rank
self.pp_rank = pp_rank
# MTP model runners
self.model_runner_list = []
# Init model and tokenizer
self.model_config = ModelConfig.from_server_args(
server_args,
@@ -261,7 +265,31 @@ class TpModelWorker(BaseTpWorker):
is_draft_worker=is_draft_worker,
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
draft_model_idx=0 if is_mtp_worker else None,
)
if is_mtp_worker:
self.model_runner_list.append(self.model_runner)
for i in range(1, server_args.speculative_num_steps):
self.model_runner_list.append(
ModelRunner(
model_config=self.model_config,
mem_fraction_static=server_args.mem_fraction_static,
gpu_id=gpu_id,
tp_rank=tp_rank,
tp_size=server_args.tp_size,
moe_ep_rank=moe_ep_rank,
moe_ep_size=server_args.ep_size,
pp_rank=pp_rank,
pp_size=server_args.pp_size,
nccl_port=nccl_port,
dp_rank=dp_rank,
server_args=server_args,
is_draft_worker=is_draft_worker,
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
draft_model_idx=i,
)
)
if server_args.skip_tokenizer_init:
self.tokenizer = self.processor = None
else:
@@ -305,7 +333,7 @@ class TpModelWorker(BaseTpWorker):
), "If configured, max_queued_requests must be at least 1 for any work to be scheduled."
self.max_req_len = min(
self.model_config.context_len - 1,
self.max_total_num_tokens - 1,
self.model_runner.max_token_pool_size - 1,
)
self.max_req_input_len = self.max_req_len - 5
assert (