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

33
python/sglang/srt/layers/logits_processor.py
View File

@@ -144,6 +144,8 @@ class LogitsMetadata:
# Whether this batch is prefill-only (no token generation needed)
is_prefill_only: bool = False
return_hidden_states_before_norm: bool = False
@classmethod
def from_forward_batch(cls, forward_batch: ForwardBatch):
if (
@@ -194,6 +196,7 @@ class LogitsMetadata:
global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu,
global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
dp_padding_mode=DpPaddingMode.SUM_LEN,
return_hidden_states_before_norm=forward_batch.return_hidden_states_before_norm,
)
def compute_dp_attention_metadata(self):
@@ -381,6 +384,7 @@ class LogitsProcessor(nn.Module):
lm_head: VocabParallelEmbedding,
logits_metadata: Union[LogitsMetadata, ForwardBatch],
aux_hidden_states: Optional[torch.Tensor] = None,
hidden_states_before_norm: Optional[torch.Tensor] = None,
) -> LogitsProcessorOutput:
if isinstance(logits_metadata, ForwardBatch):
logits_metadata = LogitsMetadata.from_forward_batch(logits_metadata)
@@ -407,6 +411,7 @@ class LogitsProcessor(nn.Module):
or logits_metadata.forward_mode.is_draft_extend_v2()
):
pruned_states = hidden_states
pruned_states_before_norm = hidden_states_before_norm
if aux_hidden_states is not None:
aux_pruned_states = [hidden for hidden in aux_hidden_states]
sample_indices = None
@@ -432,6 +437,11 @@ class LogitsProcessor(nn.Module):
- 1
)
pruned_states = hidden_states[last_index]
pruned_states_before_norm = (
hidden_states_before_norm[last_index]
if hidden_states_before_norm is not None
else None
)
if aux_hidden_states is not None:
aux_pruned_states = [hidden[last_index] for hidden in aux_hidden_states]
sample_indices = None
@@ -464,7 +474,7 @@ class LogitsProcessor(nn.Module):
sample_indices = []
input_logprob_indices_pt = 0
input_logprob_indices = []
pt, pruned_states = 0, []
pt, pruned_states, pruned_states_before_norm = 0, [], []
token_to_seq_idx = []
for idx, (extend_logprob_start_len, extend_len) in enumerate(
@@ -484,6 +494,10 @@ class LogitsProcessor(nn.Module):
# by a caller.
assert extend_len > start_len
pruned_states.append(hidden_states[pt + start_len : pt + extend_len])
if hidden_states_before_norm is not None:
pruned_states_before_norm.append(
hidden_states_before_norm[pt + start_len : pt + extend_len]
)
# Map each token to its sequence index, for chunked computation
# of input logprobs
token_to_seq_idx.extend([idx] * (extend_len - start_len))
@@ -501,6 +515,10 @@ class LogitsProcessor(nn.Module):
# Set the last token of the last sequence
token_to_seq_idx.append(len(logits_metadata.extend_seq_lens_cpu) - 1)
pruned_states = torch.cat(pruned_states)
if hidden_states_before_norm is not None:
pruned_states_before_norm = torch.cat(pruned_states_before_norm)
else:
pruned_states_before_norm = None
sample_indices = torch.tensor(
sample_indices, device=pruned_states.device, dtype=torch.int64
)
@@ -515,6 +533,7 @@ class LogitsProcessor(nn.Module):
)
hidden_states_to_store: Optional[torch.Tensor] = None
hidden_states_to_store_before_norm: Optional[torch.Tensor] = None
if logits_metadata.capture_hidden_mode.need_capture():
if logits_metadata.capture_hidden_mode.is_full():
if aux_hidden_states is not None:
@@ -522,6 +541,7 @@ class LogitsProcessor(nn.Module):
hidden_states_to_store = aux_hidden_states
else:
hidden_states_to_store = hidden_states
hidden_states_to_store_before_norm = hidden_states_before_norm
elif logits_metadata.capture_hidden_mode.is_last():
# Get the last token hidden states. If sample_indices is None,
# pruned states only contain the last tokens already.
@@ -538,11 +558,22 @@ class LogitsProcessor(nn.Module):
if sample_indices is not None
else pruned_states
)
hidden_states_to_store_before_norm = (
pruned_states_before_norm[sample_indices]
if sample_indices is not None
else pruned_states_before_norm
)
else:
assert False, "Should never reach"
del hidden_states
if (
logits_metadata.return_hidden_states_before_norm
and hidden_states_to_store_before_norm is not None
):
hidden_states_to_store = hidden_states_to_store_before_norm
if not logits_metadata.extend_return_logprob:
# Compute logits for both input and sampled tokens.
logits = self._get_logits(pruned_states, lm_head, logits_metadata)