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

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>
2025-12-19 11:40:07 +08:00
parent a0985dd5e5
commit 160a06cab2
38 changed files with 5396 additions and 169 deletions
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Optional

 import torch

+from sglang.srt.speculative.spec_utils import spec_need_hidden_states
 from sglang.srt.utils import get_compiler_backend

 if TYPE_CHECKING:
@@ -73,7 +74,6 @@ class FutureMap:
        # Get a reference for each tensor
        topk_p0 = draft_input.topk_p[0]
        topk_index0 = draft_input.topk_index[0]
-        hidden_states0 = draft_input.hidden_states[0]
        verified_id0 = draft_input.verified_id[0]
        new_seq_lens0 = draft_input.new_seq_lens[0]

@@ -87,11 +87,6 @@ class FutureMap:
            dtype=topk_index0.dtype,
            device=self.device,
        )
-        self.hidden_states_buf = torch.empty(
-            (self.future_buffer_len, *hidden_states0.shape),
-            dtype=hidden_states0.dtype,
-            device=self.device,
-        )
        self.verified_id_buf = torch.empty(
            (self.future_buffer_len, *verified_id0.shape),
            dtype=verified_id0.dtype,
@@ -103,6 +98,14 @@ class FutureMap:
            device=self.device,
        )

+        if spec_need_hidden_states():
+            hidden_states0 = draft_input.hidden_states[0]
+            self.hidden_states_buf = torch.empty(
+                (self.future_buffer_len, *hidden_states0.shape),
+                dtype=hidden_states0.dtype,
+                device=self.device,
+            )
+
    def alloc_future_indices(self, bs: int) -> FutureIndices:
        """Update the circular buffer pointer and allocate future indices."""
        cur_future_ct = self.future_ct
@@ -122,9 +125,10 @@ class FutureMap:
            indices = draft_input.future_indices.indices
            draft_input.topk_p = self.topk_p_buf[indices]
            draft_input.topk_index = self.topk_index_buf[indices]
-            draft_input.hidden_states = self.hidden_states_buf[indices]
            draft_input.verified_id = self.verified_id_buf[indices]
            draft_input.new_seq_lens = self.new_seq_lens_buf[indices]
+            if spec_need_hidden_states():
+                draft_input.hidden_states = self.hidden_states_buf[indices]
        else:
            _resolve_future_token_ids(model_worker_batch.input_ids, self.token_ids_buf)

@@ -158,6 +162,7 @@ class FutureMap:

        self.topk_p_buf[intv] = draft_input.topk_p
        self.topk_index_buf[intv] = draft_input.topk_index
-        self.hidden_states_buf[intv] = draft_input.hidden_states
        self.verified_id_buf[intv] = draft_input.verified_id
        self.new_seq_lens_buf[intv] = draft_input.new_seq_lens
+        if spec_need_hidden_states():
+            self.hidden_states_buf[intv] = draft_input.hidden_states
@@ -1217,6 +1217,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
    # Diffusion LLM
    dllm_config: Optional[DllmConfig] = None

+    # For hidden states before normal
+    return_hidden_states_before_norm: bool = False
+
    @classmethod
    def init_new(
        cls,
@@ -2113,6 +2116,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
            dllm_config=self.dllm_config,
            reqs=self.reqs,
            has_grammar=self.has_grammar,
+            return_hidden_states_before_norm=self.return_hidden_states_before_norm,
            mamba_track_indices=self.mamba_track_indices,
            mamba_track_mask=self.mamba_track_mask,
            mamba_track_seqlens=self.mamba_track_seqlens,
@@ -2242,6 +2246,9 @@ class ModelWorkerBatch:
    reqs: Optional[List[Req]] = None
    has_grammar: bool = False

+    # For hidden states before normal
+    return_hidden_states_before_norm: bool = False
+
    # For mamba state tracking
    mamba_track_indices: Optional[torch.Tensor] = None  # shape: [b], int64
    mamba_track_mask: Optional[torch.Tensor] = None  # shape: [b], bool
@@ -275,6 +275,7 @@ class Scheduler(
        self.spec_algorithm = SpeculativeAlgorithm.from_string(
            server_args.speculative_algorithm
        )
+        self.enable_mtp = server_args.enable_mtp
        self.gpu_id = gpu_id
        self.page_size = server_args.page_size
        self.enable_hierarchical_cache = server_args.enable_hierarchical_cache
@@ -479,9 +480,37 @@ class Scheduler(
        # algorithms should register their factory instead of patching this code.
        if self.spec_algorithm.is_eagle():
            draft_worker_kwargs["enable_overlap"] = self.enable_overlap
-        self.draft_worker = self.spec_algorithm.create_draft_worker(
-            **draft_worker_kwargs
-        )
+
+        # FIXME: refactor the draft worker registration logic
+        if self.enable_mtp:
+            if self.enable_overlap:
+                from sglang.srt.speculative.mtp_worker_v2 import MTPWorkerV2
+
+                self.draft_worker = MTPWorkerV2(
+                    gpu_id=self.gpu_id,
+                    tp_rank=self.tp_rank,
+                    moe_ep_rank=self.moe_ep_rank,
+                    server_args=self.server_args,
+                    nccl_port=self.port_args.nccl_port,
+                    target_worker=self.tp_worker,
+                    dp_rank=self.dp_rank,
+                )
+            else:
+                from sglang.srt.speculative.mtp_worker import MTPWorker
+
+                self.draft_worker = MTPWorker(
+                    gpu_id=self.gpu_id,
+                    tp_rank=self.tp_rank,
+                    moe_ep_rank=self.moe_ep_rank,
+                    server_args=self.server_args,
+                    nccl_port=self.port_args.nccl_port,
+                    target_worker=self.tp_worker,
+                    dp_rank=self.dp_rank,
+                )
+        else:
+            self.draft_worker = self.spec_algorithm.create_draft_worker(
+                **draft_worker_kwargs
+            )

        # Dispatch the model worker
        if self.spec_algorithm.is_none():
@@ -548,7 +577,7 @@ class Scheduler(
    def init_cache_with_memory_pool(self):
        server_args = self.server_args

-        # Hybrid memory pool configs
+        # Hybrid memory pool
        self.is_hybrid_swa = self.tp_worker.is_hybrid_swa
        self.is_hybrid_ssm = (
            self.tp_worker.model_runner.hybrid_gdn_config is not None
@@ -592,9 +621,13 @@ class Scheduler(

                self.tree_cache = ChunkCache(params)
            else:
-
                from sglang.srt.mem_cache.chunk_cache import SWAChunkCache

+                params.is_local_attention = (
+                    "Llama4ForConditionalGeneration"
+                    in self.model_config.hf_config.architectures
+                )
+
                self.tree_cache = SWAChunkCache(params)
        else:

@@ -796,11 +829,14 @@ class Scheduler(
        if self.draft_worker is None or self.spec_algorithm.is_ngram():
            draft_token_to_kv_pool = None
        elif self.spec_algorithm.is_eagle() and self.enable_overlap:
-            draft_token_to_kv_pool = (
-                self.draft_worker.draft_worker.draft_runner.token_to_kv_pool
-            )
-            model_config = self.draft_worker.draft_worker.draft_runner.model_config
+            if self.enable_mtp:
+                draft_runner = self.draft_worker.draft_worker.draft_runner_list[0]
+            else:
+                draft_runner = self.draft_worker.draft_worker.draft_runner
+            draft_token_to_kv_pool = draft_runner.token_to_kv_pool
+            model_config = draft_runner.model_config
        else:
+            # todo: should we fix this when enabling mtp or it doesn't matter since we only enable mtp in decode node thus we don't transfer draft kvs between P and D?
            draft_token_to_kv_pool = self.draft_worker.model_runner.token_to_kv_pool
            model_config = self.draft_worker.model_config

@@ -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 (