[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>

python/sglang/srt/managers/overlap_utils.py

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