From 160a06cab23ffedf3069df2721cc976ec6df63f4 Mon Sep 17 00:00:00 2001 From: Yingchun Lai Date: Fri, 19 Dec 2025 11:40:07 +0800 Subject: [PATCH] [Feature] Xiaomi `MiMo-V2-Flash` day0 support (#15207) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 谢学扬 Co-authored-by: tz Co-authored-by: 李家乐 Co-authored-by: 张晨 Co-authored-by: Shaohui Liu Co-authored-by: 王晨 Co-authored-by: jiangzihan Co-authored-by: xiexueyang Co-authored-by: Linghao Zhang Co-authored-by: ispobock Co-authored-by: Liangsheng Yin Co-authored-by: JoyFuture <35593546+JoyFuture@users.noreply.github.com> Co-authored-by: Liangsheng Yin Co-authored-by: Qiaolin Yu Co-authored-by: root --- .../srt/batch_overlap/two_batch_overlap.py | 2 + python/sglang/srt/configs/load_config.py | 3 + python/sglang/srt/configs/model_config.py | 79 +- python/sglang/srt/disaggregation/decode.py | 7 + .../decode_schedule_batch_mixin.py | 10 +- .../srt/disaggregation/mooncake/conn.py | 10 +- python/sglang/srt/disaggregation/prefill.py | 7 + python/sglang/srt/entrypoints/http_server.py | 2 +- .../srt/function_call/function_call_parser.py | 2 + .../sglang/srt/function_call/mimo_detector.py | 281 +++++ .../attention/flashattention_backend.py | 75 +- python/sglang/srt/layers/linear.py | 38 +- python/sglang/srt/layers/logits_processor.py | 33 +- python/sglang/srt/layers/quantization/fp8.py | 41 +- python/sglang/srt/managers/overlap_utils.py | 21 +- python/sglang/srt/managers/schedule_batch.py | 7 + python/sglang/srt/managers/scheduler.py | 54 +- python/sglang/srt/managers/tp_worker.py | 30 +- python/sglang/srt/mem_cache/allocator.py | 6 + .../sglang/srt/mem_cache/cache_init_params.py | 1 + python/sglang/srt/mem_cache/chunk_cache.py | 13 +- python/sglang/srt/mem_cache/memory_pool.py | 49 +- .../sglang/srt/mem_cache/swa_radix_cache.py | 3 +- .../srt/model_executor/forward_batch_info.py | 6 + .../sglang/srt/model_executor/model_runner.py | 145 ++- python/sglang/srt/model_loader/loader.py | 17 + python/sglang/srt/models/mimo_v2_flash.py | 927 ++++++++++++++++ .../sglang/srt/models/mimo_v2_flash_nextn.py | 366 +++++++ python/sglang/srt/server_args.py | 21 + .../eagle_draft_extend_cuda_graph_runner.py | 20 +- .../sglang/srt/speculative/eagle_info_v2.py | 7 + python/sglang/srt/speculative/eagle_worker.py | 44 +- .../mtp_draft_extend_cuda_graph_runner.py | 655 ++++++++++++ python/sglang/srt/speculative/mtp_utils.py | 350 +++++++ python/sglang/srt/speculative/mtp_worker.py | 989 ++++++++++++++++++ .../sglang/srt/speculative/mtp_worker_v2.py | 750 +++++++++++++ python/sglang/srt/speculative/spec_utils.py | 53 +- .../test_function_call_parser.py | 441 ++++++++ 38 files changed, 5396 insertions(+), 169 deletions(-) create mode 100644 python/sglang/srt/function_call/mimo_detector.py create mode 100644 python/sglang/srt/models/mimo_v2_flash.py create mode 100644 python/sglang/srt/models/mimo_v2_flash_nextn.py create mode 100644 python/sglang/srt/speculative/mtp_draft_extend_cuda_graph_runner.py create mode 100644 python/sglang/srt/speculative/mtp_utils.py create mode 100644 python/sglang/srt/speculative/mtp_worker.py create mode 100644 python/sglang/srt/speculative/mtp_worker_v2.py diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index da266528e..12589a879 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -728,6 +728,7 @@ class TboForwardBatchPreparer: tbo_split_seq_index=None, tbo_parent_token_range=(start_token_index, end_token_index), tbo_children=None, + original_global_num_tokens_cpu=None, global_num_tokens_gpu=None, global_num_tokens_cpu=None, global_dp_buffer_len=global_dp_buffer_len, @@ -743,6 +744,7 @@ class TboForwardBatchPreparer: top_logprobs_nums=None, token_ids_logprobs=None, next_token_logits_buffer=None, + return_hidden_states_before_norm=False, ) ) diff --git a/python/sglang/srt/configs/load_config.py b/python/sglang/srt/configs/load_config.py index b9f16e03b..c2fd4332c 100644 --- a/python/sglang/srt/configs/load_config.py +++ b/python/sglang/srt/configs/load_config.py @@ -89,6 +89,9 @@ class LoadConfig: None # Path to rollout quantization profile (e.g., /root/profile.7b.pt) ) + # For multi-layer MTP + draft_model_idx: Optional[int] = None + def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index c749dfdaa..6e601a4ae 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -100,6 +100,7 @@ class ModelConfig: model_impl: Union[str, ModelImpl] = ModelImpl.AUTO, sampling_defaults: str = "openai", quantize_and_serve: bool = False, + is_mtp: bool = False, encoder_only: bool = False, language_only: bool = False, ) -> None: @@ -111,6 +112,7 @@ class ModelConfig: self.model_impl = model_impl self.sampling_defaults = sampling_defaults self.quantize_and_serve = quantize_and_serve + self.is_mtp = is_mtp # Validate quantize_and_serve configuration self._validate_quantize_and_serve_config() @@ -158,18 +160,6 @@ class ModelConfig: self.attention_chunk_size = getattr( self.hf_text_config, "attention_chunk_size", None ) - self.is_hybrid_swa = is_hybrid_model( - self.hf_config.architectures, - hybrid_kvcache_ratio=hybrid_kvcache_ratio, - context_length=context_length, - attention_chunk_size=self.attention_chunk_size, - ) - if self.is_hybrid_swa is not None: - self.swa_attention_layer_ids, self.full_attention_layer_ids = ( - get_hybrid_layer_ids( - self.hf_config.architectures, self.hf_text_config.num_hidden_layers - ) - ) self.is_generation = is_generation_model( self.hf_config.architectures, is_embedding ) @@ -204,6 +194,9 @@ class ModelConfig: self._derive_context_length(context_length) self._derive_model_shapes() + # Update hybrid model + self._derive_hybrid_model(hybrid_kvcache_ratio) + # Verify quantization self._verify_quantization() @@ -259,6 +252,7 @@ class ModelConfig: sampling_defaults=server_args.sampling_defaults, quantize_and_serve=server_args.quantize_and_serve, override_config_file=server_args.decrypted_config_file, + is_mtp=server_args.enable_mtp, language_only=server_args.language_only, encoder_only=server_args.encoder_only, is_draft_model=is_draft_model, @@ -286,6 +280,11 @@ class ModelConfig: if is_draft_model and self.hf_config.architectures[0] == "MiMoForCausalLM": self.hf_config.architectures[0] = "MiMoMTP" + if ( + is_draft_model + and self.hf_config.architectures[0] == "MiMoV2FlashForCausalLM" + ): + self.hf_config.architectures[0] = "MiMoV2MTP" if is_draft_model and self.hf_config.architectures[0] in [ "BailingMoeV2ForCausalLM", "BailingMoeForCausalLM", @@ -301,6 +300,28 @@ class ModelConfig: self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 + def _derive_hybrid_model(self, hybrid_kvcache_ratio: Optional[float] = None): + # Use self.context_len after it has been initialized to prevent using context_len which may be None. + self.is_hybrid_swa = is_hybrid_model( + self.hf_config.architectures, + hybrid_kvcache_ratio=hybrid_kvcache_ratio, + context_length=self.context_len, + attention_chunk_size=self.attention_chunk_size, + ) + if self.is_hybrid_swa is not None: + self.swa_attention_layer_ids, self.full_attention_layer_ids = ( + get_hybrid_layer_ids( + self.hf_config.architectures, + self.hf_text_config.num_hidden_layers, + getattr(self.hf_text_config, "hybrid_layer_pattern", None), + ) + ) + + self.is_hybrid_swa_compress = self.hf_config.architectures[0] in [ + "MiMoV2FlashForCausalLM", + "MiMoV2MTP", + ] + def _derive_context_length(self, context_length: int): is_draft_model = self.is_draft_model derived_context_len = get_context_length(self.hf_text_config) @@ -342,6 +363,11 @@ class ModelConfig: "head_dim", self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads, ) + self.v_head_dim = getattr( + self.hf_text_config, + "v_head_dim", + self.head_dim, + ) # FIXME: temporary special judge for MLA architecture if ( @@ -526,6 +552,15 @@ class ModelConfig: # parallel size so each GPU has at least one KV head. return max(1, total_num_kv_heads // tensor_parallel_size) + def get_swa_num_kv_heads(self, tensor_parallel_size) -> int: + """Similar to get_num_kv_heads(), but for SWA.""" + if not self.is_hybrid_swa_compress: + return 0 + + # For MiMoV2FlashForCausalLM models + total_num_kv_heads = self.hf_text_config.swa_num_key_value_heads + return max(1, total_num_kv_heads // tensor_parallel_size) + # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py def _parse_quant_hf_config(self): quant_cfg = getattr(self.hf_config, "quantization_config", None) @@ -1106,6 +1141,11 @@ def is_hybrid_model( context_length: Optional[int], attention_chunk_size: Optional[int], ): + if model_architectures[0] in [ + "MiMoV2FlashForCausalLM", + "MiMoV2MTP", + ]: + return 1 if hybrid_kvcache_ratio is None: return None elif ( @@ -1118,7 +1158,11 @@ def is_hybrid_model( return None -def get_hybrid_layer_ids(model_architectures: List[str], num_hidden_layers: int): +def get_hybrid_layer_ids( + model_architectures: List[str], + num_hidden_layers: int, + hybrid_layer_pattern: Optional[List[int]] = None, +): if "Llama4ForConditionalGeneration" in model_architectures: swa_attention_layer_ids = [ i for i in range(num_hidden_layers) if (i + 1) % 4 != 0 @@ -1126,6 +1170,15 @@ def get_hybrid_layer_ids(model_architectures: List[str], num_hidden_layers: int) full_attention_layer_ids = [ i for i in range(num_hidden_layers) if (i + 1) % 4 == 0 ] + elif "MiMoV2FlashForCausalLM" in model_architectures: + swa_attention_layer_ids = [ + i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 1 + ] + full_attention_layer_ids = [ + i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0 + ] + elif "MiMoV2MTP" in model_architectures: + return [0], [] else: swa_attention_layer_ids = None full_attention_layer_ids = None diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index b0f090530..07059f594 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -240,6 +240,13 @@ class DecodePreallocQueue: self.prefill_pp_size = prefill_pp_size self.kv_manager = self._init_kv_manager() + if self.scheduler.tp_worker.is_hybrid_swa: + # FIXME: current SWA allocation allocate full kv cache size in prefill + self.max_total_num_tokens = min( + self.max_total_num_tokens, + self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens, + ) + def _init_kv_manager(self) -> BaseKVManager: kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) kv_args = kv_args_class() diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index efa979460..f4dea976d 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -132,13 +132,13 @@ class ScheduleBatchDisaggregationDecodeMixin: # Simulate the eagle run. if self.spec_algorithm.is_eagle(): - - b = len(self.reqs) - topk = server_args.speculative_eagle_topk + num_states = server_args.speculative_eagle_topk + if server_args.enable_mtp: + num_states *= server_args.speculative_num_steps topk_p = torch.stack( [ torch.as_tensor( - req.output_topk_p[:topk], + req.output_topk_p[:num_states], device=self.device, dtype=torch.float32, ) @@ -149,7 +149,7 @@ class ScheduleBatchDisaggregationDecodeMixin: topk_index = torch.stack( [ torch.as_tensor( - req.output_topk_index[:topk], + req.output_topk_index[:num_states], device=self.device, dtype=torch.int64, ) diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index 56a748e52..32e8c0b69 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -300,18 +300,24 @@ class MooncakeKVManager(CommonKVManager): src_k_ptrs, src_v_ptrs, dst_k_ptrs, dst_v_ptrs, layers_current_pp_stage = ( self.get_mha_kv_ptrs_with_pp(src_data_ptrs, dst_data_ptrs) ) + # item_lens structure: [k_layer0, k_layer1, ..., k_layerN, v_layer0, v_layer1, ..., v_layerN] + # Use correct item lengths for K and V separately + if layers_current_pp_stage > len(dst_k_ptrs): + logger.error( + f"layers_current_pp_stage is out of range: {layers_current_pp_stage=}, {len(dst_k_ptrs)}" + ) layers_params = [ ( src_k_ptrs[layer_id], dst_k_ptrs[layer_id], - item_lens[layer_id], + item_lens[layer_id], # K item length ) for layer_id in range(layers_current_pp_stage) ] + [ ( src_v_ptrs[layer_id], dst_v_ptrs[layer_id], - item_lens[layer_id], + item_lens[layers_current_pp_stage + layer_id], # V item length ) for layer_id in range(layers_current_pp_stage) ] diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index bd0896b09..cf63ff48a 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -109,6 +109,13 @@ class PrefillBootstrapQueue: self.transfer_backend = transfer_backend self.kv_manager = self._init_kv_manager() + if self.scheduler.tp_worker.is_hybrid_swa: + # FIXME: current SWA allocation allocate full kv cache size in prefill + self.max_total_num_tokens = min( + self.max_total_num_tokens, + self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens, + ) + def _init_kv_manager(self) -> BaseKVManager: kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS) kv_args = kv_args_class() diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index 6eaf578a6..99e07a9a0 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -1582,7 +1582,7 @@ def _execute_server_warmup( i * (2**63 // server_args.dp_size) + (i % server_args.tp_size) for i in range(server_args.dp_size) ], - "input_ids": [[0, 1, 2, 3]] * server_args.dp_size, + "input_ids": [[10, 11, 12, 13]] * server_args.dp_size, } res = requests.post( url + request_name, diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index ef39571c7..24287842b 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -19,6 +19,7 @@ from sglang.srt.function_call.gpt_oss_detector import GptOssDetector from sglang.srt.function_call.internlm_detector import InternlmDetector from sglang.srt.function_call.kimik2_detector import KimiK2Detector from sglang.srt.function_call.llama32_detector import Llama32Detector +from sglang.srt.function_call.mimo_detector import MiMoDetector from sglang.srt.function_call.minimax_m2 import MinimaxM2Detector from sglang.srt.function_call.mistral_detector import MistralDetector from sglang.srt.function_call.pythonic_detector import PythonicDetector @@ -48,6 +49,7 @@ class FunctionCallParser: "gpt-oss": GptOssDetector, "kimi_k2": KimiK2Detector, "llama3": Llama32Detector, + "mimo": MiMoDetector, "mistral": MistralDetector, "pythonic": PythonicDetector, "qwen": Qwen25Detector, diff --git a/python/sglang/srt/function_call/mimo_detector.py b/python/sglang/srt/function_call/mimo_detector.py new file mode 100644 index 000000000..c9cef1c89 --- /dev/null +++ b/python/sglang/srt/function_call/mimo_detector.py @@ -0,0 +1,281 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import ast +import html +import json +import logging +import re +from typing import Any, Dict, List + +from sglang.srt.entrypoints.openai.protocol import Tool +from sglang.srt.environ import envs +from sglang.srt.function_call.base_format_detector import BaseFormatDetector +from sglang.srt.function_call.core_types import StreamingParseResult, _GetInfoFunc + +logger = logging.getLogger(__name__) + + +def _get_param_type(func_name: str, param_name: str, tools: List[Tool]) -> str: + """Get parameter type from tool schema.""" + for tool in tools: + if tool.function.name == func_name: + props = tool.function.parameters.get("properties", {}) + if param_name in props: + return props[param_name].get("type", "string") + return "string" + + +def _convert_param_value( + param_value: str, param_name: str, func_name: str, tools: List[Tool] +) -> Any: + """ + Convert parameter value based on its type in the schema. + Adapted from vllm-project/vllm (vllm/entrypoints/openai/tool_parsers/qwen3coder_tool_parser.py) + """ + param_value = html.unescape(param_value) + + # Handle null value for any type + if param_value.lower() == "null": + return None + + param_type = _get_param_type(func_name, param_name, tools) + + if param_type in ["string", "str", "text", "varchar", "char", "enum"]: + return param_value + elif ( + param_type.startswith("int") + or param_type.startswith("integer") + or param_type.startswith("uint") + or param_type.startswith("long") + or param_type.startswith("short") + or param_type.startswith("unsigned") + ): + try: + return int(param_value) + except (ValueError, TypeError): + logger.warning( + "Parsed value '%s' of parameter '%s' is not an " + "integer in tool '%s', degenerating to string.", + param_value, + param_name, + func_name, + ) + return param_value + elif param_type.startswith("num") or param_type.startswith("float"): + try: + float_param_value = float(param_value) + return ( + float_param_value + if float_param_value - int(float_param_value) != 0 + else int(float_param_value) + ) + except (ValueError, TypeError): + logger.warning( + "Parsed value '%s' of parameter '%s' is not a float " + "in tool '%s', degenerating to string.", + param_value, + param_name, + func_name, + ) + return param_value + elif param_type in ["boolean", "bool", "binary"]: + param_value = param_value.lower() + if param_value not in ["true", "false"]: + logger.warning( + "Parsed value '%s' of parameter '%s' is not a boolean " + "(`true` or `false`) in tool '%s', degenerating to " + "false.", + param_value, + param_name, + func_name, + ) + return param_value == "true" + else: + if ( + param_type in ["object", "array", "arr"] + or param_type.startswith("dict") + or param_type.startswith("list") + ): + try: + param_value = json.loads(param_value) + return param_value + except (json.JSONDecodeError, TypeError, ValueError): + logger.warning( + "Parsed value '%s' of parameter '%s' cannot be " + "parsed with json.loads in tool '%s', will try " + "other methods to parse it.", + param_value, + param_name, + func_name, + ) + try: + param_value = ast.literal_eval(param_value) # safer + except (ValueError, SyntaxError, TypeError): + logger.warning( + "Parsed value '%s' of parameter '%s' cannot be " + "converted via Python `ast.literal_eval()` in tool " + "'%s', degenerating to string.", + param_value, + param_name, + func_name, + ) + return param_value + + +class MiMoDetector(BaseFormatDetector): + """ + Detector for MiMo function call format. + + Format: + + + pwd && ls + + + """ + + def __init__(self): + super().__init__() + self.bot_token = "" + self.eot_token = "" + self.tool_call_regex = re.compile(r"(.*?)", re.DOTALL) + self.func_regex = re.compile(r"]+)>(.*?)", re.DOTALL) + self.param_regex = re.compile( + r"]+)>(.*?)", re.DOTALL + ) + + def has_tool_call(self, text: str) -> bool: + return self.bot_token in text + + def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult: + """Parse complete text for tool calls.""" + idx = text.find(self.bot_token) + if idx == -1: + return StreamingParseResult(normal_text=text, calls=[]) + + normal_text = text[:idx] + tool_indices = self._get_tool_indices(tools) + + calls = [] + last_end = idx + + for match in self.tool_call_regex.finditer(text): + tool_call_body = match.group(1) + + parsed = self._parse_tool_call(tool_call_body, tools) + + if parsed: + func_name = parsed.get("name") + if func_name not in tool_indices: + # Unknown function + logger.warning(f"Unknown function: {func_name}") + if not envs.SGLANG_FORWARD_UNKNOWN_TOOLS.get(): + # Return tool call block as normal text + normal_text += text[last_end : match.end()] + last_end = match.end() + continue + calls.extend(self.parse_base_json(parsed, tools)) + + last_end = match.end() + + return StreamingParseResult(normal_text=normal_text, calls=calls) + + def parse_streaming_increment( + self, new_text: str, tools: List[Tool] + ) -> StreamingParseResult: + """ + Streaming parsing: buffer until complete tool call block. + """ + self._buffer += new_text + current_text = self._buffer + + start = current_text.find(self.bot_token) + if start == -1: + if self.current_tool_id > 0: + # Already processing tool calls, keep buffering + # (more tool calls might come, don't discard text yet) + return StreamingParseResult(normal_text="") + else: + # No tool calls seen yet, return as normal text + self._buffer = "" + return StreamingParseResult(normal_text=current_text) + + # Find end token AFTER the start token + end = current_text.find(self.eot_token, start) + if end == -1: + # Incomplete tool call, return text before start and keep buffering + normal_text = current_text[:start] + self._buffer = current_text[start:] + return StreamingParseResult(normal_text=normal_text) + + # Parse the complete tool call block + result = self.detect_and_parse(current_text[: end + len(self.eot_token)], tools) + + if result.calls: + # Valid tool call - initialize tracking if first one + if self.current_tool_id == -1: + self.current_tool_id = 0 + self.prev_tool_call_arr = [] + self.streamed_args_for_tool = [""] + + while len(self.prev_tool_call_arr) <= self.current_tool_id: + self.prev_tool_call_arr.append({}) + while len(self.streamed_args_for_tool) <= self.current_tool_id: + self.streamed_args_for_tool.append("") + + call = result.calls[0] + self.prev_tool_call_arr[self.current_tool_id] = { + "name": call.name, + "arguments": json.loads(call.parameters) if call.parameters else {}, + } + self.streamed_args_for_tool[self.current_tool_id] = call.parameters + call.tool_index = self.current_tool_id + self.current_tool_id += 1 + + self._buffer = current_text[end + len(self.eot_token) :] + return result + + def _parse_tool_call( + self, tool_call_body: str, tools: List[Tool] + ) -> Dict[str, Any]: + """ + Parse content inside .... + + Structure: + tool_call_body contains: ...params... + """ + # Match complete body block + func_match = self.func_regex.search(tool_call_body) + if not func_match: + return None + + func_name = func_match.group(1).strip() + func_body = func_match.group(2) + + params = {} + for param_match in self.param_regex.finditer(func_body): + param_name = param_match.group(1).strip() + param_value = param_match.group(2) + params[param_name] = _convert_param_value( + param_value, param_name, func_name, tools + ) + + return {"name": func_name, "parameters": params} + + def supports_structural_tag(self) -> bool: + return False + + def structure_info(self) -> _GetInfoFunc: + raise NotImplementedError diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 672605560..ed4ec1273 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -340,6 +340,7 @@ class FlashAttentionBackend(AttentionBackend): self.full_to_swa_index_mapping = ( model_runner.token_to_kv_pool.full_to_swa_index_mapping ) + self.token_to_kv_pool = model_runner.token_to_kv_pool self.topk = model_runner.server_args.speculative_eagle_topk or 0 self.speculative_num_steps = speculative_num_steps self.speculative_num_draft_tokens = ( @@ -792,6 +793,15 @@ class FlashAttentionBackend(AttentionBackend): cu_seqlens_k = swa_spec_metadata.cu_seqlens_k else: page_table = metadata.page_table + if self.is_hybrid_swa: + _, is_swa = forward_batch.token_to_kv_pool.layers_mapping[ + layer.layer_id + ] + if is_swa: + page_table = self.token_to_kv_pool.translate_loc_from_full_to_swa( + page_table + ) + window_size = (self.attention_chunk_size, 0) cu_seqlens_q = metadata.cu_seqlens_q cache_seqlens = metadata.cache_seqlens_int32 max_seqlen_q = metadata.max_seq_len_q @@ -807,7 +817,7 @@ class FlashAttentionBackend(AttentionBackend): -1, self.page_size, layer.tp_k_head_num, layer.head_dim ) value_cache = value_cache.view( - -1, self.page_size, layer.tp_v_head_num, layer.head_dim + -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim ) if layer.is_cross_attention: page_table = metadata.encoder_page_table @@ -1098,7 +1108,7 @@ class FlashAttentionBackend(AttentionBackend): -1, self.page_size, layer.tp_k_head_num, layer.head_dim ) value_cache = value_cache.view( - -1, self.page_size, layer.tp_v_head_num, layer.head_dim + -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim ) if layer.is_cross_attention: @@ -1143,6 +1153,17 @@ class FlashAttentionBackend(AttentionBackend): ) else: page_table = metadata.page_table + if self.is_hybrid_swa: + _, is_swa = forward_batch.token_to_kv_pool.layers_mapping[ + layer.layer_id + ] + if is_swa: + page_table = ( + self.token_to_kv_pool.translate_loc_from_full_to_swa( + page_table + ) + ) + window_size = (self.attention_chunk_size, 0) cache_seqlens = metadata.cache_seqlens_int32 cu_seqlens_k = metadata.cu_seqlens_k max_seqlen_q = metadata.max_seq_len_q @@ -1743,7 +1764,7 @@ class FlashAttentionBackend(AttentionBackend): self.target_verify_metadata_topk_swa[bs] = metadata_swa metadata.swa_spec_metadata = metadata_swa - elif forward_mode.is_draft_extend(): + elif forward_mode.is_draft_extend(include_v2=True): metadata.cache_seqlens_int32 = self.draft_extend_metadata["cache_seqlens"][ :bs ] @@ -2048,6 +2069,54 @@ class FlashAttentionBackend(AttentionBackend): ] metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size) + elif forward_mode.is_draft_extend_v2(): + metadata = self.draft_extend_metadata[bs] + metadata.cache_seqlens_int32.copy_(seq_lens) + + metadata.max_seq_len_k = seq_lens_cpu.max().item() + metadata.cu_seqlens_k[1:].copy_( + torch.cumsum(metadata.cache_seqlens_int32, dim=0, dtype=torch.int32) + ) + + extend_seq_lens_tensor = getattr(spec_info, "extend_seq_lens_tensor", None) + extend_seq_lens_cpu = getattr(spec_info, "extend_seq_lens_cpu", None) + if extend_seq_lens_tensor is not None: + extend_seq_lens = extend_seq_lens_tensor.to(torch.int32) + elif extend_seq_lens_cpu is not None: + extend_seq_lens = torch.as_tensor( + extend_seq_lens_cpu, + dtype=torch.int32, + device=device, + ) + else: + default_extend = getattr( + spec_info, "num_tokens_per_batch", self.speculative_num_steps + 1 + ) + extend_seq_lens = torch.full( + (bs,), default_extend, dtype=torch.int32, device=device + ) + extend_seq_lens_cpu = [default_extend] * bs + + if extend_seq_lens_cpu: + metadata.max_seq_len_q = int(max(extend_seq_lens_cpu)) + else: + metadata.max_seq_len_q = getattr( + spec_info, "num_tokens_per_batch", self.speculative_num_steps + 1 + ) + + metadata.cu_seqlens_q[1:].copy_( + torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32) + ) + + max_seq_pages = ( + metadata.max_seq_len_k + self.page_size - 1 + ) // self.page_size + page_indices = self.req_to_token[ + req_pool_indices[:, None], + self.draft_extend_metadata["strided_indices"][:max_seq_pages], + ] + metadata.page_table[:, :max_seq_pages].copy_(page_indices // self.page_size) + if encoder_lens is not None: # Only support encoder size 1 for now metadata.encoder_max_seq_len_k = encoder_lens[0] diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index e64fd5070..428f3a261 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -305,6 +305,7 @@ class ColumnParallelLinear(LinearBase): tp_rank: Optional[int] = None, tp_size: Optional[int] = None, use_presharded_weights: bool = False, + skip_block_quant_check: bool = False, ): super().__init__( input_size, output_size, skip_bias_add, params_dtype, quant_config, prefix @@ -338,6 +339,7 @@ class ColumnParallelLinear(LinearBase): input_size=self.input_size, output_size=self.output_size, params_dtype=self.params_dtype, + skip_block_quant_check=skip_block_quant_check, weight_loader=( self.weight_loader_v2 if self.quant_method.__class__.__name__ in WEIGHT_LOADER_V2_SUPPORTED @@ -815,9 +817,12 @@ class QKVParallelLinear(ColumnParallelLinear): tp_rank: Optional[int] = None, tp_size: Optional[int] = None, load_presharded_attn: bool = False, + v_head_size: Optional[int] = None, + skip_block_quant_check: bool = False, ): self.hidden_size = hidden_size self.head_size = head_size + self.v_head_size = v_head_size if v_head_size is not None else head_size self.total_num_heads = total_num_heads if total_num_kv_heads is None: total_num_kv_heads = total_num_heads @@ -837,14 +842,17 @@ class QKVParallelLinear(ColumnParallelLinear): self.num_kv_head_replicas = 1 self.q_proj_shard_size = self.num_heads * self.head_size self.kv_proj_shard_size = self.num_kv_heads * self.head_size + self.v_proj_shard_size = self.num_kv_heads * self.v_head_size input_size = self.hidden_size output_size = ( - (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size - ) + self.num_heads * self.head_size + + self.num_kv_heads * self.head_size + + self.num_kv_heads * self.v_head_size + ) * tp_size self.output_sizes = [ self.num_heads * self.head_size * tp_size, # q_proj self.num_kv_heads * self.head_size * tp_size, # k_proj - self.num_kv_heads * self.head_size * tp_size, # v_proj + self.num_kv_heads * self.v_head_size * tp_size, # v_proj ] self.use_presharded_weights = load_presharded_attn quant_config = None if _disable_hip_linear_quant else quant_config @@ -861,6 +869,7 @@ class QKVParallelLinear(ColumnParallelLinear): tp_rank=tp_rank, tp_size=tp_size, use_presharded_weights=self.use_presharded_weights, + skip_block_quant_check=skip_block_quant_check, ) def _get_shard_offset_mapping(self, loaded_shard_id: str): @@ -868,7 +877,8 @@ class QKVParallelLinear(ColumnParallelLinear): "q": 0, "k": self.num_heads * self.head_size, "v": (self.num_heads + self.num_kv_heads) * self.head_size, - "total": (self.num_heads + 2 * self.num_kv_heads) * self.head_size, + "total": (self.num_heads + self.num_kv_heads) * self.head_size + + self.num_kv_heads * self.v_head_size, } return shard_offset_mapping.get(loaded_shard_id) @@ -876,7 +886,7 @@ class QKVParallelLinear(ColumnParallelLinear): shard_size_mapping = { "q": self.num_heads * self.head_size, "k": self.num_kv_heads * self.head_size, - "v": self.num_kv_heads * self.head_size, + "v": self.num_kv_heads * self.v_head_size, } return shard_size_mapping.get(loaded_shard_id) @@ -903,7 +913,7 @@ class QKVParallelLinear(ColumnParallelLinear): ( "v", (self.total_num_heads + self.total_num_kv_heads) * self.head_size, - self.total_num_kv_heads * self.head_size, + self.total_num_kv_heads * self.v_head_size, ), ] @@ -1055,7 +1065,7 @@ class QKVParallelLinear(ColumnParallelLinear): ( "v", (self.total_num_heads + self.total_num_kv_heads) * self.head_size, - self.total_num_kv_heads * self.head_size, + self.total_num_kv_heads * self.v_head_size, ), ] use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit", False) @@ -1089,11 +1099,12 @@ class QKVParallelLinear(ColumnParallelLinear): "v": ( (self.total_num_heads + self.total_num_kv_heads) * self.head_size, - self.total_num_kv_heads * self.head_size, + self.total_num_kv_heads * self.v_head_size, ), "total": ( - (self.total_num_heads + 2 * self.total_num_kv_heads) - * self.head_size, + (self.total_num_heads + self.total_num_kv_heads) + * self.head_size + + self.total_num_kv_heads * self.v_head_size, 0, ), } @@ -1121,7 +1132,7 @@ class QKVParallelLinear(ColumnParallelLinear): shard_size = self.num_kv_heads * self.head_size elif loaded_shard_id == "v": shard_offset = (self.num_heads + self.num_kv_heads) * self.head_size - shard_size = self.num_kv_heads * self.head_size + shard_size = self.num_kv_heads * self.v_head_size # Special case for Quantized Weights. # If quantized, we need to adjust the offset and size to account # for the packing. @@ -1145,10 +1156,11 @@ class QKVParallelLinear(ColumnParallelLinear): ), "v": ( (self.num_heads + self.num_kv_heads) * self.head_size, - self.num_kv_heads * self.head_size, + self.num_kv_heads * self.v_head_size, ), "total": ( - (self.num_heads + 2 * self.num_kv_heads) * self.head_size, + (self.num_heads + self.num_kv_heads) * self.head_size + + self.num_kv_heads * self.v_head_size, 0, ), } diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index f258e4890..48e735e6d 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -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) diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 830a6752c..e7f2b2a28 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -234,6 +234,7 @@ class Fp8LinearMethod(LinearMethodBase): input_size: int, output_size: int, params_dtype: torch.dtype, + skip_block_quant_check: bool = False, **extra_weight_attrs, ): output_size_per_partition = sum(output_partition_sizes) @@ -245,25 +246,31 @@ class Fp8LinearMethod(LinearMethodBase): self.quant_config.weight_block_size[0], self.quant_config.weight_block_size[1], ) - # Required by row parallel - if tp_size > 1 and input_size // input_size_per_partition == tp_size: - if input_size_per_partition % block_k != 0: - raise ValueError( - f"Weight input_size_per_partition = " - f"{input_size_per_partition} is not divisible by " - f"weight quantization block_k = {block_k}." - ) - # Required by column parallel or enabling merged weights - if ( - tp_size > 1 and output_size // output_size_per_partition == tp_size - ) or len(output_partition_sizes) > 1: - for output_partition_size in output_partition_sizes: - if output_partition_size % block_n != 0: + + if skip_block_quant_check: + logger.warning_once( + f"Skipping block quantization checks for weight partition." + ) + else: + # Required by row parallel + if tp_size > 1 and input_size // input_size_per_partition == tp_size: + if input_size_per_partition % block_k != 0: raise ValueError( - f"Weight output_partition_size = " - f"{output_partition_size} is not divisible by " - f"weight quantization block_n = {block_n}." + f"Weight input_size_per_partition = " + f"{input_size_per_partition} is not divisible by " + f"weight quantization block_k = {block_k}." ) + # Required by column parallel or enabling merged weights + if ( + tp_size > 1 and output_size // output_size_per_partition == tp_size + ) or len(output_partition_sizes) > 1: + for output_partition_size in output_partition_sizes: + if output_partition_size % block_n != 0: + raise ValueError( + f"Weight output_partition_size = " + f"{output_partition_size} is not divisible by " + f"weight quantization block_n = {block_n}." + ) layer.logical_widths = output_partition_sizes layer.input_size_per_partition = input_size_per_partition diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index c47b69afb..3056cdc54 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index d49bae768..7445c5fcc 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 97304cdb8..e491e41e4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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 diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index dff873edd..8d8c85e8f 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -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 ( diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index 4fefac941..e8edb50da 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/python/sglang/srt/mem_cache/allocator.py @@ -291,6 +291,12 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.is_not_in_free_group = True self.free_group = [] + def get_cpu_copy(self, indices): + return self._kvcache.get_cpu_copy(indices) + + def load_cpu_copy(self, kv_cache_cpu, indices): + return self._kvcache.load_cpu_copy(kv_cache_cpu, indices) + @triton.jit def alloc_extend_kernel( diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index c1258deb3..a833bec98 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/python/sglang/srt/mem_cache/cache_init_params.py @@ -26,3 +26,4 @@ class CacheInitParams: enable_kv_cache_events: bool = False enable_mamba_extra_buffer: bool = False + is_local_attention: bool = False diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index a543a62d1..756dd3024 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -82,6 +82,7 @@ class SWAChunkCache(ChunkCache): def __init__(self, params: CacheInitParams): assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator) super().__init__(params) + self.is_local_attention = params.is_local_attention def evict_swa( self, @@ -89,10 +90,14 @@ class SWAChunkCache(ChunkCache): prelen: int, attention_chunk_size: int, ): - if prelen >= req.evicted_seqlen_local + attention_chunk_size: - new_evicted_seqlen_local = attention_chunk_size * ( - prelen // attention_chunk_size - ) + thresh = req.evicted_seqlen_local + attention_chunk_size * 2 + if self.is_local_attention: + thresh -= attention_chunk_size + + if prelen >= thresh: + new_evicted_seqlen_local = ( + prelen // attention_chunk_size * attention_chunk_size + ) - (attention_chunk_size if not self.is_local_attention else 0) free_slots = self.req_to_token_pool.req_to_token[ req.req_pool_idx, req.evicted_seqlen_local : new_evicted_seqlen_local ] diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index ce9c8e35f..733663df9 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -614,6 +614,10 @@ class MHATokenToKVPool(KVCache): layer_num: int, device: str, enable_memory_saver: bool, + v_head_dim: Optional[int] = None, + swa_head_num: Optional[int] = None, + swa_head_dim: Optional[int] = None, + swa_v_head_dim: Optional[int] = None, start_layer: Optional[int] = None, end_layer: Optional[int] = None, enable_alt_stream: bool = True, @@ -629,8 +633,13 @@ class MHATokenToKVPool(KVCache): start_layer, end_layer, ) - self.head_num = head_num - self.head_dim = head_dim + self.head_num = swa_head_num if swa_head_num is not None else head_num + self.head_dim = swa_head_dim if swa_head_dim is not None else head_dim + self.v_head_dim = ( + swa_v_head_dim + if swa_v_head_dim is not None + else v_head_dim if v_head_dim is not None else head_dim + ) self._create_buffers() @@ -708,7 +717,7 @@ class MHATokenToKVPool(KVCache): ] self.v_buffer = [ torch.zeros( - (self.size + self.page_size, self.head_num, self.head_dim), + (self.size + self.page_size, self.head_num, self.v_head_dim), dtype=self.store_dtype, device=self.device, ) @@ -1289,6 +1298,9 @@ class SWAKVPool(KVCache): layer_num=self.swa_layer_nums, **kwargs, ) + kwargs.pop("swa_head_num", None) + kwargs.pop("swa_head_dim", None) + kwargs.pop("swa_v_head_dim", None) self.full_kv_pool = token_to_kv_pool_class( size=size, dtype=dtype, @@ -1306,7 +1318,7 @@ class SWAKVPool(KVCache): k_size, v_size = self.get_kv_size_bytes() self.mem_usage = (k_size + v_size) / GB logger.info( - f"SWAKVPool mem usage: {self.mem_usage} GB, swa size: {self.size_swa}, full size: {self.size}" + f"SWAKVPool mem usage: {self.mem_usage:.2f} GB, swa size: {self.size_swa}, full size: {self.size}" ) def get_kv_size_bytes(self): @@ -1392,6 +1404,35 @@ class SWAKVPool(KVCache): layer_id_override=layer_id_pool, ) + def get_cpu_copy(self, indices): + # For SWA, we need to copy KV cache from both full and SWA pools + # The indices are for the full pool, and we use mapping to get SWA indices + full_kv_cpu = self.full_kv_pool.get_cpu_copy(indices) + + # Get SWA indices through the mapping + # Note: SWA allocation always creates 1:1 mapping, so no need to filter + if self.full_to_swa_index_mapping is not None: + swa_indices = self.full_to_swa_index_mapping[indices] + swa_kv_cpu = self.swa_kv_pool.get_cpu_copy(swa_indices) + else: + swa_kv_cpu = None + + return {"full": full_kv_cpu, "swa": swa_kv_cpu} + + def load_cpu_copy(self, kv_cache_cpu, indices): + # Load KV cache back from CPU to both full and SWA pools + # Note: indices here are NEW indices (newly allocated), different from get_cpu_copy indices + full_kv_cpu = kv_cache_cpu["full"] + swa_kv_cpu = kv_cache_cpu["swa"] + + # Load full KV cache to the new indices + self.full_kv_pool.load_cpu_copy(full_kv_cpu, indices) + + # Load SWA KV cache if it exists + if swa_kv_cpu is not None and self.full_to_swa_index_mapping is not None: + swa_indices = self.full_to_swa_index_mapping[indices] + self.swa_kv_pool.load_cpu_copy(swa_kv_cpu, swa_indices) + class MLATokenToKVPool(KVCache): def __init__( diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 3cbc85c7f..901b67a99 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -485,7 +485,8 @@ class SWARadixCache(BasePrefixCache): ) # free the unaligned tail - self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:]) + if not self.is_eagle: + self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:]) # Remove req slot release the cache lock self.req_to_token_pool.free(req.req_pool_idx) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index dbdeea8c7..df4081bb5 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -346,6 +346,7 @@ class ForwardBatch: attn_backend: AttentionBackend = None # For DP attention + original_global_num_tokens_cpu: Optional[List[int]] = None global_num_tokens_cpu: Optional[List[int]] = None global_num_tokens_gpu: Optional[torch.Tensor] = None # Has to be None when cuda graph is captured. @@ -391,6 +392,9 @@ class ForwardBatch: # Record the split metadata of the sequence number of NSA context parallels. nsa_cp_metadata: Optional[NSAContextParallelMetadata] = None + # For hidden states before normal + return_hidden_states_before_norm: bool = False + @classmethod def init_new( cls, @@ -434,6 +438,7 @@ class ForwardBatch: token_type_ids=batch.token_type_ids, tbo_split_seq_index=batch.tbo_split_seq_index, dimensions=batch.dimensions, + return_hidden_states_before_norm=batch.return_hidden_states_before_norm, ) device = model_runner.device @@ -462,6 +467,7 @@ class ForwardBatch: global_num_tokens = batch.global_num_tokens global_num_tokens_for_logprob = batch.global_num_tokens_for_logprob + ret.original_global_num_tokens_cpu = batch.global_num_tokens ret.global_num_tokens_cpu = global_num_tokens ret.global_num_tokens_gpu = torch.tensor( global_num_tokens, dtype=torch.int64 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ebc418731..0fb96a700 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -296,6 +296,7 @@ class ModelRunner: is_draft_worker: bool = False, req_to_token_pool: Optional[ReqToTokenPool] = None, token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None, + draft_model_idx: Optional[int] = None, ): # Parse args self.mem_fraction_static = mem_fraction_static @@ -324,10 +325,13 @@ class ModelRunner: self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator self.is_hybrid_swa = model_config.is_hybrid_swa + self.is_hybrid_swa_compress = model_config.is_hybrid_swa_compress self.use_mla_backend = self.model_config.attention_arch == AttentionArch.MLA self.attention_chunk_size = model_config.attention_chunk_size self.forward_pass_id = 0 self.init_new_workspace = False + self.kv_cache_memory = 0 + self.draft_model_idx = draft_model_idx self.remote_instance_transfer_engine = None self.remote_instance_transfer_engine_session_id = "" @@ -481,6 +485,8 @@ class ModelRunner: self.model_config.num_attention_layers, ) ) + if self.model_config.hf_config.architectures[0] == "MiMoV2MTP": + model_num_layers = 1 self.start_layer = getattr(self.model, "start_layer", 0) self.end_layer = getattr(self.model, "end_layer", model_num_layers) self.num_effective_layers = self.end_layer - self.start_layer @@ -493,6 +499,23 @@ class ModelRunner: ) ), "PP is not compatible with MTP models." + # Consider PP, so use start_layer and end_layer. + full_attention_layer_ids = [ + layer_idx + for layer_idx in range(self.start_layer, self.end_layer + 1) + if hasattr(self.model_config, "full_attention_layer_ids") + and layer_idx in self.model_config.full_attention_layer_ids + ] + swa_attention_layer_ids = [ + layer_idx + for layer_idx in range(self.start_layer, self.end_layer + 1) + if hasattr(self.model_config, "swa_attention_layer_ids") + and layer_idx in self.model_config.swa_attention_layer_ids + ] + # Update back to model_config. + self.model_config.swa_attention_layer_ids = swa_attention_layer_ids + self.model_config.full_attention_layer_ids = full_attention_layer_ids + # Apply torchao quantization torchao_applied = getattr(self.model, "torchao_applied", False) # In layered loading, torchao may have been applied @@ -811,6 +834,7 @@ class ModelRunner: remote_instance_weight_loader_transfer_engine=self.remote_instance_transfer_engine, modelopt_config=modelopt_config, rl_quant_profile=self.server_args.rl_quant_profile, + draft_model_idx=self.draft_model_idx, ) if self.device == "cpu": self.model_config = adjust_config_with_unaligned_cpu_tp( @@ -1431,6 +1455,8 @@ class ModelRunner: ) elif config := self.mambaish_config: num_layers = len(config.full_attention_layer_ids) + elif self.model_config.full_attention_layer_ids: + num_layers = len(self.model_config.full_attention_layer_ids) else: num_layers = self.num_effective_layers if self.use_mla_backend: @@ -1468,9 +1494,8 @@ class ModelRunner: else: cell_size = ( self.model_config.get_num_kv_heads(get_attention_tp_size()) - * self.model_config.head_dim + * (self.model_config.head_dim + self.model_config.v_head_dim) * num_layers - * 2 * torch._utils._element_size(self.kv_cache_dtype) ) @@ -1491,12 +1516,24 @@ class ModelRunner: // scale_block_size ) + if self.model_config.hf_config.architectures[0] == "MiMoV2FlashForCausalLM": + cell_size += ( + self.model_config.get_swa_num_kv_heads(get_attention_tp_size()) + * ( + self.model_config.hf_text_config.swa_head_dim + + self.model_config.hf_text_config.swa_v_head_dim + ) + * len(self.model_config.swa_attention_layer_ids) + * torch._utils._element_size(self.kv_cache_dtype) + ) rest_memory = available_gpu_memory - total_gpu_memory * ( 1 - self.mem_fraction_static ) if self.mambaish_config is not None: rest_memory = self.handle_max_mamba_cache(rest_memory) - max_num_token = int(rest_memory * (1 << 30) // cell_size) + self.kv_cache_memory = int(rest_memory * (1 << 30)) + max_num_token = int(self.kv_cache_memory // cell_size) + logger.info(f"The available memory for KV cache is {rest_memory:.2f} GB.") return max_num_token def handle_max_mamba_cache(self, total_rest_memory): @@ -1578,6 +1615,14 @@ class ModelRunner: return config.llm_config return None + @property + def max_token_pool_size(self): + """Return the max token pool size considering hybrid swa settings.""" + if self.is_hybrid_swa: + return min(self.swa_max_total_num_tokens, self.max_total_num_tokens) + else: + return self.max_total_num_tokens + @property def kimi_linear_config(self): config = self.model_config.hf_config @@ -1590,6 +1635,7 @@ class ModelRunner: return self.mamba2_config or self.hybrid_gdn_config or self.kimi_linear_config def set_num_token_hybrid(self): + page_size = self.server_args.page_size if ( "Llama4ForConditionalGeneration" in self.model_config.hf_config.architectures @@ -1607,44 +1653,33 @@ class ModelRunner: 4 * self.max_total_num_tokens - 12 * self.max_total_num_tokens * temp_ratio // (3 * temp_ratio + 1) ) - self.swa_max_total_num_tokens = int( - self.swa_max_total_num_tokens - // self.server_args.page_size - * self.server_args.page_size + self.swa_max_total_num_tokens = ( + self.swa_max_total_num_tokens // page_size * page_size ) - self.full_max_total_num_tokens = int( - self.full_max_total_num_tokens - // self.server_args.page_size - * self.server_args.page_size + self.full_max_total_num_tokens = ( + self.full_max_total_num_tokens // page_size * page_size + ) + self.max_total_num_tokens = self.full_max_total_num_tokens + elif "MiMoV2MTP" in self.model_config.hf_config.architectures: + assert self.is_draft_worker + # MiMoV2MTP uses SWA, so set full KV cache to 0 + self.full_max_total_num_tokens = 0 + self.swa_max_total_num_tokens = ( + self.max_total_num_tokens // page_size * page_size + ) + self.max_total_num_tokens = self.swa_max_total_num_tokens + elif self.model_config.hf_config.architectures[0] == "MiMoV2FlashForCausalLM": + self.full_max_total_num_tokens = ( + self.max_total_num_tokens // page_size * page_size + ) + self.swa_max_total_num_tokens = ( + self.max_total_num_tokens // page_size * page_size ) self.max_total_num_tokens = self.full_max_total_num_tokens else: assert self.sliding_window_size is not None and self.sliding_window_size > 0 - full_attention_layer_ids = [] - swa_attention_layer_ids = [] - - try: - layers = self.model.model.layers - except: - try: - layers = self.model.language_model.model.layers - except: - try: - layers = self.model.language_model.layers - except: - self.is_hybrid_swa = False - return - - for layer in layers: - if ( - layer.self_attn.attn.sliding_window_size is None - or layer.self_attn.attn.sliding_window_size == -1 - ): - full_attention_layer_ids.append(layer.layer_id) - else: - swa_attention_layer_ids.append(layer.layer_id) - self.model_config.swa_attention_layer_ids = swa_attention_layer_ids - self.model_config.full_attention_layer_ids = full_attention_layer_ids + full_layers_num = len(self.model_config.full_attention_layer_ids) + swa_layers_num = len(self.model_config.swa_attention_layer_ids) # Algorithm: # Existing max_total_num_tokens is per layer and assume all layers have the same number of tokens. @@ -1653,8 +1688,6 @@ class ModelRunner: total_tokens = ( self.max_total_num_tokens * self.model_config.num_hidden_layers ) - full_layers_num = len(full_attention_layer_ids) - swa_layers_num = len(swa_attention_layer_ids) swa_full_tokens_ratio = self.server_args.swa_full_tokens_ratio # Solve the equations: @@ -1667,9 +1700,9 @@ class ModelRunner: ) self.max_total_num_tokens = self.full_max_total_num_tokens - logger.info( - f"Use Sliding window memory pool. full_layer_tokens={self.full_max_total_num_tokens}, swa_layer_tokens={self.swa_max_total_num_tokens}" - ) + logger.info( + f"Use sliding window memory pool. full_layer_tokens={self.full_max_total_num_tokens}, swa_layer_tokens={self.swa_max_total_num_tokens}" + ) def can_run_piecewise_cuda_graph(self): if self.server_args.enable_torch_compile: @@ -1778,10 +1811,9 @@ class ModelRunner: else: # We are sharing the `token_to_kv_pool`, and both verify and draft tokens # can be concurrently allocated, so we should give a headroom for it. - self.server_args.draft_runner_cache_size = ( - self.max_total_num_tokens + extra_tokens = ( # draft - + max_num_reqs + max_num_reqs * self.server_args.speculative_num_steps * self.server_args.speculative_eagle_topk # verify @@ -1791,7 +1823,9 @@ class ModelRunner: ) # Target worker and draft worker shares the same indices for the # token_to_kv_pool, so we should make sure to match max_total_num_tokens. - self.max_total_num_tokens = self.server_args.draft_runner_cache_size + self.max_total_num_tokens += extra_tokens + self.server_args.draft_runner_cache_size = self.max_total_num_tokens + self.server_args.max_num_reqs = max_num_reqs if max_total_tokens is not None: @@ -1988,6 +2022,18 @@ class ModelRunner: ) else: if self.is_hybrid_swa: + kwargs = {} + if self.is_hybrid_swa_compress: + kwargs = { + "swa_head_num": max( + 1, + self.model_config.hf_text_config.swa_num_key_value_heads + // get_attention_tp_size(), + ), + "swa_head_dim": self.model_config.hf_text_config.swa_head_dim, + "swa_v_head_dim": self.model_config.hf_text_config.swa_v_head_dim, + "v_head_dim": self.model_config.hf_text_config.v_head_dim, + } self.token_to_kv_pool = SWAKVPool( size=self.full_max_total_num_tokens, size_swa=self.swa_max_total_num_tokens, @@ -2000,6 +2046,7 @@ class ModelRunner: full_attention_layer_ids=self.model_config.full_attention_layer_ids, enable_kvcache_transpose=False, device=self.device, + **kwargs, ) elif config := self.mambaish_config: extra_args = {} @@ -2117,6 +2164,14 @@ class ModelRunner: ) else: assert self.is_draft_worker + if self.is_hybrid_swa: + assert ( + self.token_to_kv_pool_allocator.__class__ + == SWATokenToKVPoolAllocator + ) + self.token_to_kv_pool.full_to_swa_index_mapping = ( + self.token_to_kv_pool_allocator.full_to_swa_index_mapping + ) logger.info( f"Memory pool end. " diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 697af0b55..7590b3c10 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -498,6 +498,23 @@ class DefaultModelLoader(BaseModelLoader): else: weights_iterator = pt_weights_iterator(hf_weights_files) + if self.load_config.draft_model_idx is not None: + import re + + pattern = r"model.mtp.layers.(\d+)." + filtered_weights = [] + for name, tensor in weights_iterator: + group = re.match(pattern, name) + if group is not None: + idx = int(group.group(1)) + if idx != self.load_config.draft_model_idx: + continue + new_name = name.replace(group.group(), "model.mtp.layers.0.") + else: + new_name = name + filtered_weights.append((source.prefix + new_name, tensor)) + return tuple(filtered_weights) + # Apply the prefix. return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) diff --git a/python/sglang/srt/models/mimo_v2_flash.py b/python/sglang/srt/models/mimo_v2_flash.py new file mode 100644 index 000000000..d2563ddc1 --- /dev/null +++ b/python/sglang/srt/models/mimo_v2_flash.py @@ -0,0 +1,927 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import logging +from typing import Any, Dict, Iterable, Optional, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn + +from sglang.srt.distributed import ( + get_moe_expert_parallel_world_size, + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) +from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + enable_moe_dense_fully_dp, +) +from sglang.srt.layers.dp_attention import ( + get_attention_tp_rank, + get_attention_tp_size, + is_dp_attention_enabled, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import get_moe_a2a_backend, get_moe_runner_backend +from sglang.srt.layers.moe.ep_moe.layer import DeepEPMoE, get_moe_impl_class +from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer, get_layer_id +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, + kv_cache_scales_loader, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import LazyValue, add_prefix, make_layers + +MiMoV2FlashConfig = None + +logger = logging.getLogger(__name__) + + +class MiMoV2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + prefix: str = "", + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + self.tp_size = tp_size + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=add_prefix("down_proj", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x, forward_batch: ForwardBatch = None): + if (self.tp_size == 1) and x.shape[0] == 0: + return x + + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MoEGate(nn.Module): + def __init__( + self, + config, + quant_config, + prefix: str = "", + is_nextn: bool = False, + ): + super().__init__() + self.is_nextn = is_nextn + self.dtype = torch.float32 + self.weight = nn.Parameter( + torch.empty((config.n_routed_experts, config.hidden_size), dtype=self.dtype) + ) + if config.topk_method == "noaux_tc": + correction_bias_dtype = ( + torch.bfloat16 + if quant_config is not None + and quant_config.get_name() == "modelopt_fp4" + and get_moe_runner_backend().is_flashinfer_trtllm() + else self.dtype + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty((config.n_routed_experts), dtype=correction_bias_dtype) + ) + else: + self.e_score_correction_bias = None + + def forward(self, hidden_states): + logits = F.linear(hidden_states.to(self.dtype), self.weight, None) + + return logits + + +class MiMoV2MoE(nn.Module): + + def __init__( + self, + config: MiMoV2FlashConfig, + layer_id: int, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + is_nextn: bool = False, + ): + super().__init__() + self.tp_size = get_tensor_model_parallel_world_size() + + self.config = config + self.layer_id = layer_id + + if self.tp_size > config.n_routed_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.n_routed_experts}." + ) + + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for now." + ) + + self.gate = MoEGate( + config=config, + quant_config=quant_config, + prefix=add_prefix("gate", prefix), + is_nextn=is_nextn, + ) + + experts_type = get_moe_impl_class(quant_config) + self.experts = experts_type( + num_experts=config.n_routed_experts + + get_global_server_args().ep_num_redundant_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + layer_id=self.layer_id, + quant_config=quant_config, + routed_scaling_factor=1.0, + prefix=add_prefix("experts", prefix), + ) + + self.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=config.norm_topk_prob, + use_grouped_topk=True, + num_expert_group=config.n_group, + topk_group=config.topk_group, + correction_bias=self.gate.e_score_correction_bias, + quant_config=quant_config, + routed_scaling_factor=1.0, + apply_routed_scaling_factor_on_output=self.experts.should_fuse_routed_scaling_factor_in_topk, + # Some Fp4 MoE backends require the output format to be bypassed but the MTP layers are unquantized + # and requires the output format to be standard. We use quant_config to determine the output format. + output_format=TopKOutputFormat.STANDARD if quant_config is None else None, + ) + + # todo : implement tbo forward needed + if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake(): + # TODO: we will support tp < ep in the future + self.ep_size = get_moe_expert_parallel_world_size() + self.num_experts = ( + config.n_routed_experts + + get_global_server_args().ep_num_redundant_experts + ) + self.renormalize = config.norm_topk_prob + self.topk_group = config.topk_group + self.num_expert_group = config.n_group + self.correction_bias = ( + self.gate.e_score_correction_bias.data + if self.gate.e_score_correction_bias is not None + else None + ) + + self._enable_a2a_moe = ( + get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake() + ) + + def get_moe_weights(self): + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] + ] + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + should_allreduce_fusion: bool = False, + ) -> torch.Tensor: + if not self._enable_a2a_moe: + return self.forward_normal( + hidden_states, + should_allreduce_fusion, + ) + else: + return self.forward_deepep(hidden_states, forward_batch) + + def forward_normal( + self, + hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, + ) -> torch.Tensor: + + if hidden_states.shape[0] > 0: + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + else: + topk_output = self.topk.empty_topk_output(hidden_states.device) + + final_hidden_states = self.experts(hidden_states, topk_output) + + if self.tp_size > 1 and not should_allreduce_fusion: + final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) + + return final_hidden_states + + def forward_deepep( + self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + ) -> torch.Tensor: + if hidden_states.shape[0] > 0: + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=forward_batch.num_token_non_padded, + expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ), + ) + else: + topk_output = self.topk.empty_topk_output(hidden_states.device) + + final_hidden_states = self.experts( + hidden_states=hidden_states, topk_output=topk_output + ) + + return final_hidden_states + + +class MiMoV2Attention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: Optional[int] = None, + v_head_dim: Optional[int] = None, + v_scale: Optional[float] = None, + sliding_window_size: int = -1, # if is -1 ,normal attention,else ,window attention + attention_bias: bool = False, + attention_sink_bias: bool = False, + layer_id: int = 0, + rope_theta: float = 1000000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 32768, + quant_config: Optional[QuantizationConfig] = None, + partial_rotary_factor: float = 1.0, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + + self.total_num_heads = num_heads + assert self.total_num_heads % attn_tp_size == 0 + self.num_heads = self.total_num_heads // attn_tp_size + self.total_num_kv_heads = num_kv_heads + if self.total_num_kv_heads >= attn_tp_size: + # Number of KV heads is greater than TP size, so we partition + # the KV heads across multiple tensor parallel GPUs. + assert self.total_num_kv_heads % attn_tp_size == 0 + else: + # Number of KV heads is less than TP size, so we replicate + # the KV heads across multiple tensor parallel GPUs. + assert attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) + self.head_dim = head_dim + self.v_head_dim = v_head_dim if v_head_dim is not None else head_dim + + self.q_size = self.num_heads * self.head_dim + self.k_size = self.num_kv_heads * self.head_dim + self.v_size = self.num_kv_heads * self.v_head_dim + + self.v_scale = v_scale + + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + v_head_size=self.v_head_dim, + bias=attention_bias, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + prefix=add_prefix("qkv_proj", prefix), + skip_block_quant_check=True, + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.v_head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + reduce_results=False, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + partial_rotary_factor=partial_rotary_factor, + ) + + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + v_head_dim=self.v_head_dim, + sliding_window_size=sliding_window_size, # if is -1 ,normal attention,else ,window attention + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + ) + + self.attention_sink_bias = ( + torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False) + if attention_sink_bias + else None + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + + # [t, h, dr] + q, k = self.rotary_emb(positions, q, k) + # [t, h, d] + + if self.v_scale is not None: + v = v * self.v_scale + attn_output = self.attn(q, k, v, forward_batch, sinks=self.attention_sink_bias) + output, _ = self.o_proj(attn_output) + return output + + +class MiMoV2DecoderLayer(nn.Module): + def __init__( + self, + config: MiMoV2FlashConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.layer_id = layer_id + + rope_theta = getattr(config, "rope_theta", 1000000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 32768) + + if self.is_swa_layer(): + self.self_attn = MiMoV2Attention( + hidden_size=self.hidden_size, + num_heads=config.swa_num_attention_heads, + num_kv_heads=config.swa_num_key_value_heads, + head_dim=config.swa_head_dim, + v_head_dim=getattr(config, "swa_v_head_dim", None), + v_scale=getattr(config, "attention_value_scale", None), + sliding_window_size=config.sliding_window_size, + attention_bias=config.attention_bias, + attention_sink_bias=getattr( + config, "add_swa_attention_sink_bias", False + ), + layer_id=layer_id, + rope_theta=getattr(config, "swa_rope_theta", rope_theta), + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), + prefix=add_prefix("self_attn", prefix), + ) + else: + self.self_attn = MiMoV2Attention( + hidden_size=self.hidden_size, + num_heads=self.config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + head_dim=config.head_dim, + v_head_dim=getattr(config, "v_head_dim", None), + v_scale=getattr(config, "attention_value_scale", None), + sliding_window_size=-1, # normal attention + attention_bias=config.attention_bias, + attention_sink_bias=getattr( + config, "add_full_attention_sink_bias", False + ), + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), + prefix=add_prefix("self_attn", prefix), + ) + + self.is_layer_sparse = self.is_moe_layer(layer_id) + is_previous_layer_sparse = self.is_moe_layer(layer_id - 1) + is_next_layer_sparse = self.is_moe_layer(layer_id + 1) + + if self.is_layer_sparse: + self.mlp = MiMoV2MoE( + config=config, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + layer_id=layer_id, + ) + else: + if enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + self.mlp = MiMoV2MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.layernorm_epsilon + ) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=config.num_hidden_layers, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, + ) + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + # Self Attention + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, residual, forward_batch + ) + + if hidden_states.shape[0] != 0: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + + hidden_states = self.mlp(hidden_states, forward_batch) + + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + + return hidden_states, residual + + def is_moe_layer(self, layer_idx: int) -> bool: + return ( + hasattr(self.config, "moe_layer_freq") + and 0 <= layer_idx < len(self.config.moe_layer_freq) + and not isinstance(self.config.moe_layer_freq, int) + and self.config.moe_layer_freq[layer_idx] + ) + + def is_swa_layer(self) -> bool: + return self.config.hybrid_layer_pattern[self.layer_id] == 1 + + +class MiMoV2Model(nn.Module): + def __init__( + self, + config: MiMoV2FlashConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + decoder_layer_type: type[nn.Module] = MiMoV2DecoderLayer, + ) -> None: + super().__init__() + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self.pp_group = get_pp_group() + + if self.pp_group.is_first_rank: + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + enable_tp=not is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + else: + self.embed_tokens = PPMissingLayer() + + # Use the provided decoder layer type or default to MiMoV2DecoderLayer + decoder_layer_type = decoder_layer_type or MiMoV2DecoderLayer + self.layers, self.start_layer, self.end_layer = make_layers( + config.num_hidden_layers, + layer_fn=lambda idx, prefix: decoder_layer_type( + layer_id=idx, + config=config, + quant_config=quant_config, + prefix=prefix, + ), + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=add_prefix("layers", prefix), + ) + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + else: + self.norm = PPMissingLayer(return_tuple=True) + + def get_input_embedding(self, input_ids: torch.Tensor) -> torch.Tensor: + if hasattr(self.config, "scale_emb"): + return self.get_input_embeddings()(input_ids) * self.config.scale_emb + else: + return self.get_input_embeddings()(input_ids) + + def get_input_embeddings(self) -> nn.Embedding: + return self.embed_tokens + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + if self.pp_group.is_first_rank: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + for i in range(self.start_layer, self.end_layer): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + forward_batch, + residual, + ) + + hidden_states_before_norm = None + if not self.pp_group.is_last_rank: + return PPProxyTensors( + { + "hidden_states": hidden_states, + "residual": residual, + } + ) + else: + if hidden_states.shape[0] > 0: + if residual is None: + hidden_states_before_norm = hidden_states + hidden_states = self.norm(hidden_states) + else: + hidden_states_before_norm = hidden_states + residual + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states, hidden_states_before_norm + + # If this function is called, it should always initialize KV cache scale + # factors (or else raise an exception). Thus, handled exceptions should + # make sure to leave KV cache scale factors in a known good (dummy) state + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_size() + for layer_idx, scaling_factor in kv_cache_scales_loader( + quantization_param_path, + attn_tp_rank, + attn_tp_size, + self.config.num_hidden_layers, + self.config.__class__.model_type, + ): + if not isinstance(self.layers[layer_idx], nn.Identity): + layer_self_attn = self.layers[layer_idx].self_attn + if hasattr(layer_self_attn.attn, "k_scale"): + layer_self_attn.attn.k_scale = scaling_factor + layer_self_attn.attn.v_scale = scaling_factor + else: + raise RuntimeError( + "Self attention has no KV cache scaling " "factor attribute!" + ) + + +class MiMoV2FlashForCausalLM(nn.Module): + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: MiMoV2FlashConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + self.model = MiMoV2Model( + config, quant_config=quant_config, prefix=add_prefix("model", prefix) + ) + + if self.pp_group.is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + else: + # ranks other than the last rank will have a placeholder layer + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config) + + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, MiMoV2MoE) + } + ) + + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + + def get_input_embedding(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.get_input_embedding(input_ids) + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> torch.Tensor: + hidden_states, hidden_states_before_norm = self.model( + input_ids, + positions, + forward_batch, + input_embeds, + pp_proxy_tensors=pp_proxy_tensors, + ) + + if self.pp_group.is_last_rank: + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + forward_batch, + hidden_states_before_norm=hidden_states_before_norm, + ) + else: + return hidden_states + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + # (param_name, weight_name, expert_id, shard_id) + expert_params_mapping = DeepEPMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts, + ) + + params_dict = dict(self.named_parameters()) + + for name, loaded_weight in weights: + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + + if self.config.tie_word_embeddings and "lm_head.weight" in name: + if self.pp_group.world_size > 1 and self.pp_group.is_last_rank: + # Handle pp weight tying here + # find the embed_tokens.weight in the weights + embed_token_weights = next( + filter(lambda x: x[0] == "model.embed_tokens.weight", weights) + )[1] + loaded_weight = embed_token_weights + else: + continue + + # TODO: skip mtp weights for now, need to implement mtp + if "mtp" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if ("mlp.experts." in name) and name not in params_dict: + continue + + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if name in params_dict.keys(): + param = params_dict[name] + if "attention_sink_bias" in name: + start = get_attention_tp_rank() * param.numel() + param.data.copy_( + loaded_weight[start : start + param.numel()] + ) + else: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning(f"Parameter {name} not found in params_dict") + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + self.model.load_kv_cache_scales(quantization_param_path) + + @classmethod + def get_model_config_for_expert_location(cls, config): + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=getattr(config, "n_routed_experts", 1), + num_groups=getattr(config, "n_group", None), + ) + + +EntryClass = MiMoV2FlashForCausalLM diff --git a/python/sglang/srt/models/mimo_v2_flash_nextn.py b/python/sglang/srt/models/mimo_v2_flash_nextn.py new file mode 100644 index 000000000..1ce9ce47e --- /dev/null +++ b/python/sglang/srt/models/mimo_v2_flash_nextn.py @@ -0,0 +1,366 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import logging +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_tensor_model_parallel_world_size +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + enable_moe_dense_fully_dp, +) +from sglang.srt.layers.dp_attention import ( + get_attention_tp_rank, + is_dp_attention_enabled, +) +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_loader.weight_utils import default_weight_loader +from sglang.srt.models.mimo_v2_flash import ( + MiMoV2Attention, + MiMoV2FlashForCausalLM, + MiMoV2MLP, +) +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import add_prefix + +MiMoV2FlashConfig = None + +logger = logging.getLogger(__name__) + + +class MiMoV2MTPLayer(nn.Module): + def __init__( + self, + config: MiMoV2FlashConfig, + layer_id: int = 0, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + + rope_theta = getattr(config, "rope_theta", 1000000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 32768) + + self.self_attn = MiMoV2Attention( + hidden_size=self.hidden_size, + num_heads=config.swa_num_attention_heads, + num_kv_heads=config.swa_num_key_value_heads, + head_dim=config.swa_head_dim, + v_head_dim=getattr(config, "swa_v_head_dim", None), + v_scale=getattr(config, "attention_value_scale", None), + sliding_window_size=config.sliding_window_size, + attention_bias=config.attention_bias, + attention_sink_bias=getattr(config, "add_swa_attention_sink_bias", False), + layer_id=layer_id, + rope_theta=getattr(config, "swa_rope_theta", rope_theta), + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), + prefix=add_prefix("self_attn", prefix), + ) + self.is_layer_sparse = False + is_previous_layer_sparse = True + is_next_layer_sparse = False + + if enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + self.mlp = MiMoV2MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.layernorm_epsilon + ) + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + num_layers=1, + is_layer_sparse=self.is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, + ) + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: + + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, residual, forward_batch + ) + + if hidden_states.shape[0] != 0: + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + + return hidden_states, residual + + +class MiMoV2ModelNextN(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + enable_tp=not is_dp_attention_enabled(), + prefix=add_prefix("embed_tokens", prefix), + ) + + self.enorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + self.hnorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + + self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + + self.mtp_block = MiMoV2MTPLayer( + config, + 0, + quant_config=quant_config, + prefix=add_prefix("decoder", prefix), + ) + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: torch.Tensor = None, + ) -> torch.Tensor: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + if hidden_states.shape[0] > 0: + hidden_states = self.eh_proj( + torch.cat( + ( + self.enorm(hidden_states), + self.hnorm(forward_batch.spec_info.hidden_states), + ), + dim=-1, + ) + ) + hidden_states, residual = self.mtp_block( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + residual=None, + ) + hidden_states_before_norm = None + if not forward_batch.forward_mode.is_idle(): + if residual is not None: + hidden_states_before_norm = hidden_states + residual + hidden_states, _ = self.final_layernorm(hidden_states, residual) + else: + hidden_states_before_norm = hidden_states + hidden_states = self.final_layernorm(hidden_states) + + return hidden_states, hidden_states_before_norm + + +class MiMoV2MTP(MiMoV2FlashForCausalLM): + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + self.tp_size = get_tensor_model_parallel_world_size() + self.quant_config = quant_config + + self.model = MiMoV2ModelNextN( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + self.logits_processor = LogitsProcessor(config) + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + hidden_states, hidden_states_before_norm = self.model( + input_ids, positions, forward_batch + ) + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + forward_batch, + hidden_states_before_norm=hidden_states_before_norm, + ) + + def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + name = self.map_model_name_to_mtp_param_name(name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + + if weight_name not in name: + continue + if "mtp_block" not in name: + break + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + + if "mtp_block" not in name and ( + "embed_tokens" not in name + and "lm_head" not in name + and "enorm" not in name + and "hnorm" not in name + and "eh_proj" not in name + and "final_layernorm" not in name + ): + continue + if name in params_dict.keys(): + param = params_dict[name] + if "attention_sink_bias" in name: + start = get_attention_tp_rank() * param.numel() + param.data.copy_(loaded_weight[start : start + param.numel()]) + else: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning(f"Parameter {name} not found in params_dict") + + def map_model_name_to_mtp_param_name(self, name: str) -> str: + import re + + if "pre_mlp_layernorm" in name: + name = name.replace("pre_mlp_layernorm", "post_attention_layernorm") + + name_without_prefix = [ + "enorm", + "hnorm", + "eh_proj", + "final_layernorm", + ] + pattern = r"model.mtp.layers.(\d+)." + group = re.match(pattern, name) + if group is not None: + for sub_name in name_without_prefix: + if sub_name in name: + name = name.replace(group.group(), "model.") + return name + name = name.replace(group.group(), "model.mtp_block.") + return name + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +EntryClass = MiMoV2MTP diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a7c48fad4..7e8648031 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -438,6 +438,10 @@ class ServerArgs: speculative_ngram_branch_length: int = 18 speculative_ngram_capacity: int = 10 * 1000 * 1000 + # For Multi-Layer MTP + # FIXME: rename -> enable_multi_layer_mtp + enable_mtp: bool = False + # Expert parallelism ep_size: int = 1 moe_a2a_backend: Literal["none", "deepep", "mooncake", "ascend_fuseep"] = "none" @@ -1175,6 +1179,16 @@ class ServerArgs: ), "Triton kernel MoE is only supported when ep_size == 1" self.disable_hybrid_swa_memory = True + elif "MiMoV2FlashForCausalLM" in model_arch: + self.swa_full_tokens_ratio = 1.0 + logger.warning( + "Reset swa_full_tokens_ratio to 1.0 for MiMoV2FlashForCausalLM model" + ) + if self.enable_hierarchical_cache: + self.disable_hybrid_swa_memory = True + logger.warning( + "Disable hybrid SWA memory for MiMoV2FlashForCausalLM model with hierarchical cache" + ) elif "Llama4" in model_arch and self.device != "cpu": # Auto-select attention backend for Llama4 if not specified if self.attention_backend is None: @@ -3405,6 +3419,13 @@ class ServerArgs: help="The cache capacity for ngram speculative decoding.", ) + # Speculative decoding (MTP) + parser.add_argument( + "--enable-mtp", + action="store_true", + help="Enable multi-layer MTP speculative decoding.", + ) + # Expert parallelism parser.add_argument( "--expert-parallel-size", diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index c884bef8e..916136cec 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -127,7 +127,9 @@ class EAGLEDraftExtendCudaGraphRunner: self.seq_lens = torch.full( (self.max_bs,), self.seq_len_fill_value, dtype=torch.int32 ) - self.extend_seq_lens = torch.ones((self.max_bs,), dtype=torch.int32) + self.extend_seq_lens = torch.full( + (self.max_bs,), self.num_tokens_per_bs, dtype=torch.int32 + ) self.accept_length = torch.full( (self.max_bs,), self.num_tokens_per_bs, dtype=torch.int32 ) @@ -389,14 +391,16 @@ class EAGLEDraftExtendCudaGraphRunner: self.seq_lens.fill_(self.seq_len_fill_value) self.out_cache_loc.zero_() self.positions.zero_() - self.accept_length.fill_(1) - self.extend_seq_lens.fill_(1) + self.accept_length.fill_(self.num_tokens_per_bs) + self.extend_seq_lens.fill_(self.num_tokens_per_bs) # Common inputs self.input_ids[:num_tokens].copy_(forward_batch.input_ids) self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens) if forward_batch.extend_seq_lens is not None: self.extend_seq_lens[:raw_bs].copy_(forward_batch.extend_seq_lens) + else: + self.extend_seq_lens[:raw_bs].fill_(self.num_tokens_per_bs) self.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc) self.positions[:num_tokens].copy_(forward_batch.positions) if ( @@ -420,6 +424,16 @@ class EAGLEDraftExtendCudaGraphRunner: if forward_batch.extend_seq_lens_cpu is not None: self.extend_seq_lens_cpu[:raw_bs] = forward_batch.extend_seq_lens_cpu + else: + self.extend_seq_lens_cpu[:raw_bs] = [self.num_tokens_per_bs] * raw_bs + if bs > raw_bs: + self.extend_seq_lens_cpu[raw_bs:bs] = [self.num_tokens_per_bs] * ( + bs - raw_bs + ) + forward_batch.spec_info.extend_seq_lens_cpu = list( + self.extend_seq_lens_cpu[:bs] + ) + forward_batch.spec_info.extend_seq_lens_tensor = self.extend_seq_lens[:bs] if bs != raw_bs: forward_batch.spec_info.positions = self.positions[:num_tokens] diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index 3894c2176..caacbfff5 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -10,6 +10,7 @@ import triton.language as tl from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import ModelWorkerBatch, ScheduleBatch +from sglang.srt.mem_cache.chunk_cache import SWAChunkCache from sglang.srt.mem_cache.common import ( alloc_paged_token_slots_extend, alloc_token_slots, @@ -79,6 +80,12 @@ def assign_draft_cache_locs_page_size_1( @dataclass class EagleDraftInputV2Mixin: def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch): + if isinstance(batch.tree_cache, SWAChunkCache): + for req in batch.reqs: + batch.tree_cache.evict_swa( + req, req.seqlen - 1, batch.model_config.attention_chunk_size + ) + from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func bs = batch.batch_size() diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index db25c5066..f681ca158 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -19,6 +19,7 @@ from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.mem_cache.chunk_cache import SWAChunkCache from sglang.srt.mem_cache.common import ( alloc_paged_token_slots_extend, alloc_token_slots, @@ -53,6 +54,7 @@ from sglang.srt.speculative.spec_utils import ( draft_tp_context, fast_topk, generate_token_bitmask, + get_last_loc_large_page_size_large_top_k, load_token_map, select_top_k_tokens, ) @@ -366,6 +368,12 @@ class EAGLEWorker(TpModelWorker): ) def _draft_preprocess_decode(self, batch: ScheduleBatch): + if isinstance(batch.tree_cache, SWAChunkCache): + for req in batch.reqs: + batch.tree_cache.evict_swa( + req, req.seqlen - 1, batch.model_config.attention_chunk_size + ) + # Parse args num_seqs = batch.batch_size() spec_info = batch.spec_info @@ -1083,39 +1091,3 @@ def get_last_loc_large_page_size_top_k_1( prefix_lens, ) return prefix_lens, seq_lens, last_loc - - -# Disable torch.compile for this function because it will be -# even slower. -# @torch.compile(dynamic=True) -def get_last_loc_large_page_size_large_top_k( - req_to_token: torch.Tensor, - req_pool_indices: torch.Tensor, - seq_lens: torch.Tensor, - speculative_num_steps: int, - topk: int, - page_size: int, -): - prefix_lens = seq_lens - last_page_lens = prefix_lens % page_size - num_new_pages_per_topk = ( - last_page_lens + speculative_num_steps + page_size - 1 - ) // page_size - seq_lens = prefix_lens // page_size * page_size + num_new_pages_per_topk * ( - page_size * topk - ) - extend_lens = seq_lens - prefix_lens - last_loc = get_last_loc( - req_to_token, - req_pool_indices, - prefix_lens, - ) - - return ( - prefix_lens, - seq_lens, - last_loc, - num_new_pages_per_topk, - extend_lens, - last_page_lens, - ) diff --git a/python/sglang/srt/speculative/mtp_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/mtp_draft_extend_cuda_graph_runner.py new file mode 100644 index 000000000..f94e51ae4 --- /dev/null +++ b/python/sglang/srt/speculative/mtp_draft_extend_cuda_graph_runner.py @@ -0,0 +1,655 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +from __future__ import annotations + +import bisect +import logging +import time +from typing import TYPE_CHECKING, Callable + +import torch + +from sglang.srt.layers.dp_attention import DpPaddingMode, set_dp_buffer_len +from sglang.srt.model_executor.cuda_graph_runner import ( + CUDA_GRAPH_CAPTURE_FAILED_MSG, + CudaGraphRunner, + DeepEPCudaGraphRunnerAdapter, + LogitsProcessorOutput, + get_batch_sizes_to_capture, + get_global_graph_memory_pool, + model_capture_mode, + set_global_graph_memory_pool, + set_is_extend_in_batch, + set_torch_compile_config, +) +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ForwardMode, +) +from sglang.srt.speculative.eagle_info import EagleDraftInput +from sglang.srt.speculative.mtp_utils import assign_new_state_triton +from sglang.srt.speculative.spec_utils import fast_topk +from sglang.srt.utils import ( + get_available_gpu_memory, + require_attn_tp_gather, + require_gathered_buffer, + require_mlp_sync, + require_mlp_tp_gather, +) + +if TYPE_CHECKING: + from sglang.srt.speculative.mtp_worker_v2 import MTPDraftWorker + + +logger = logging.getLogger(__name__) + + +class MTPDraftExtendCudaGraphRunner: + def __init__(self, mtp_worker: MTPDraftWorker, step: int): + # Parse args + self.step = step + self.mtp_worker = mtp_worker + self.model_runner = model_runner = mtp_worker.mtp_model_runner(self.step) + self.forward_mode = ForwardMode.DRAFT_EXTEND_V2 + + self.graphs = {} + self.output_buffers = {} + self.enable_torch_compile = model_runner.server_args.enable_torch_compile + self.disable_padding = model_runner.server_args.disable_cuda_graph_padding + self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args) + self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args) + self.require_mlp_sync = require_mlp_sync(model_runner.server_args) + self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args) + self.tp_size = self.model_runner.tp_size + self.dp_size = model_runner.server_args.dp_size + self.enable_pdmux = model_runner.server_args.enable_pdmux + self.speculative_num_steps = model_runner.server_args.speculative_num_steps + self.speculative_num_draft_tokens = ( + model_runner.server_args.speculative_num_draft_tokens + ) + self.topk = model_runner.server_args.speculative_eagle_topk + self.enable_profile_cuda_graph = ( + model_runner.server_args.enable_profile_cuda_graph + ) + self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner) + self.padded_static_len = -1 + self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() + + # For Attention Backend + self.num_tokens_per_bs = self.speculative_num_steps + 1 + step + self.max_bs = max(self.capture_bs) + self.max_num_token = self.max_bs * self.num_tokens_per_bs + + self.mtp_worker.draft_extend_attn_backend_list[self.step].init_cuda_graph_state( + self.max_bs, self.max_num_token + ) + self.seq_len_fill_value = self.mtp_worker.draft_extend_attn_backend_list[ + self.step + ].get_cuda_graph_seq_len_fill_value() + + def init_buffers_and_capture( + self, + cuda_graph_buffers, + offset, + next_cuda_graph_runner, + ): + self.next_cuda_graph_runner = next_cuda_graph_runner + self.seq_lens_cpu = cuda_graph_buffers["seq_lens_cpu"] + self.extend_seq_lens_cpu = [self.num_tokens_per_bs] * self.max_bs + + if self.enable_torch_compile: + set_torch_compile_config() + + # Graph inputs + with torch.device(self.model_runner.device): + # sliced buffers + # slice according to max_num_token + self.input_ids = cuda_graph_buffers["input_ids"][ + offset : offset + self.max_num_token + ] + self.out_cache_loc = cuda_graph_buffers["out_cache_loc"][ + offset : offset + self.max_num_token + ] + self.swa_out_cache_loc = cuda_graph_buffers["swa_out_cache_loc"][ + offset : offset + self.max_num_token + ] + self.positions = cuda_graph_buffers["positions"][ + offset : offset + self.max_num_token + ] + + # shared states + self.seq_lens = cuda_graph_buffers["seq_lens"] + self.req_pool_indices = cuda_graph_buffers["req_pool_indices"] + self.accept_length = cuda_graph_buffers["accept_length"] + + self.extend_seq_lens = torch.full( + (self.max_bs,), + self.num_tokens_per_bs, + dtype=torch.int32, + ) + self.extend_start_loc = torch.arange( + 0, + self.max_bs * self.num_tokens_per_bs, + step=self.num_tokens_per_bs, + dtype=torch.int32, + ) + + self.mrope_positions = torch.zeros( + (3, self.max_num_token), dtype=torch.int64 + ) + + self.hidden_states = torch.zeros( + (self.max_num_token, self.model_runner.model_config.hidden_size), + dtype=self.model_runner.dtype, + ) + + if self.require_gathered_buffer: + if self.require_mlp_tp_gather: + self.global_num_tokens_gpu = torch.zeros( + (self.dp_size,), dtype=torch.int32 + ) + self.global_num_tokens_for_logprob_gpu = torch.zeros( + (self.dp_size,), dtype=torch.int32 + ) + else: + assert self.require_attn_tp_gather + self.global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32) + self.global_num_tokens_for_logprob_gpu = torch.zeros( + (1,), dtype=torch.int32 + ) + else: + self.global_num_tokens_gpu = None + self.global_num_tokens_for_logprob_gpu = None + + if hasattr( + self.model_runner.model_config.hf_config, "draft_vocab_size" + ): # llama_eagle + vocab_size = self.model_runner.model_config.hf_config.draft_vocab_size + elif hasattr( + self.model_runner.model_config.hf_config, "hot_vocab_size" + ): # llama_eagle3 + vocab_size = self.model_runner.model_config.hf_config.hot_vocab_size + else: + vocab_size = self.model_runner.model_config.vocab_size + + self.next_token_logits_buffer = torch.zeros( + ( + ( + self.max_bs * self.num_tokens_per_bs + if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2 + else self.max_bs + ), + vocab_size, + ), + dtype=torch.float, + ) + + # Capture + try: + with model_capture_mode(): + self.capture() + except RuntimeError as e: + raise Exception( + f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}" + ) + + def can_run(self, forward_batch: ForwardBatch): + if self.require_mlp_tp_gather: + cuda_graph_bs = ( + max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs + if self.model_runner.spec_algorithm.is_eagle() + else max(forward_batch.global_num_tokens_cpu) + ) + else: + cuda_graph_bs = forward_batch.seq_lens.numel() + + is_bs_supported = ( + cuda_graph_bs in self.graphs + if self.disable_padding + else cuda_graph_bs <= self.max_bs + ) + + if self.require_mlp_sync: + is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph + + return is_bs_supported + + def _create_graph(self): + return torch.cuda.CUDAGraph() + + def _capture_init(self, run_once_fn): + for _ in range(2): + torch.cuda.synchronize() + self.model_runner.tp_group.barrier() + run_once_fn() + + def _capture_graph(self, graph, pool, stream, run_once_fn): + with torch.cuda.graph(graph, pool=pool, stream=stream): + out = run_once_fn() + return out + + def _replay(self, forward_batch: ForwardBatch): + self.graphs[self.bs].replay() + + def capture(self): + CudaGraphRunner.capture(self) + + def get_forward_batch(self, bs: int) -> ForwardBatch: + num_tokens = bs * self.num_tokens_per_bs + + # Graph inputs + input_ids = self.input_ids[:num_tokens] + req_pool_indices = self.req_pool_indices[:bs] + seq_lens = self.seq_lens[:bs] + seq_lens_cpu = self.seq_lens_cpu[:bs] + extend_seq_lens = self.extend_seq_lens[:bs] + extend_seq_lens_cpu = self.extend_seq_lens_cpu[:bs] + extend_start_loc = self.extend_start_loc[:bs] + accept_length = self.accept_length[:bs] + out_cache_loc = self.out_cache_loc[:num_tokens] + positions = self.positions[:num_tokens] + mrope_positions = self.mrope_positions[:, :num_tokens] + hidden_states = self.hidden_states[:num_tokens] + next_token_logits_buffer = self.next_token_logits_buffer[ + : bs if self.forward_mode == ForwardMode.DRAFT_EXTEND else num_tokens + ] + + if self.require_mlp_tp_gather: + self.global_num_tokens_gpu.copy_( + torch.tensor( + [num_tokens] * self.dp_size, + dtype=torch.int32, + device=self.input_ids.device, + ) + ) + self.global_num_tokens_for_logprob_gpu.copy_( + torch.tensor( + [num_tokens] * self.dp_size, + dtype=torch.int32, + device=self.input_ids.device, + ) + ) + global_dp_buffer_len = num_tokens * self.dp_size + elif self.require_attn_tp_gather: + self.global_num_tokens_gpu.copy_( + torch.tensor( + [num_tokens], + dtype=torch.int32, + device=self.input_ids.device, + ) + ) + self.global_num_tokens_for_logprob_gpu.copy_( + torch.tensor( + [bs], + dtype=torch.int32, + device=self.input_ids.device, + ) + ) + global_dp_buffer_len = num_tokens + else: + global_dp_buffer_len = None + + spec_info = EagleDraftInput( + hidden_states=hidden_states, + accept_length=accept_length, + ) + spec_info.positions = None + + # Forward batch + forward_batch = ForwardBatch( + forward_mode=self.forward_mode, + batch_size=bs, + input_ids=input_ids, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + next_token_logits_buffer=next_token_logits_buffer, + req_to_token_pool=self.model_runner.req_to_token_pool, + token_to_kv_pool=self.model_runner.token_to_kv_pool, + out_cache_loc=out_cache_loc, + seq_lens_sum=seq_lens.sum().item(), + return_logprob=False, + positions=positions, + mrope_positions=mrope_positions, + global_num_tokens_gpu=self.global_num_tokens_gpu, + global_num_tokens_for_logprob_gpu=self.global_num_tokens_for_logprob_gpu, + dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(), + global_dp_buffer_len=global_dp_buffer_len, + spec_algorithm=self.model_runner.spec_algorithm, + spec_info=spec_info, + capture_hidden_mode=CaptureHiddenMode.FULL, + attn_backend=self.mtp_worker.draft_extend_attn_backend_list[self.step], + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + padded_static_len=self.padded_static_len, + # added args + extend_start_loc=extend_start_loc, + extend_num_tokens=self.num_tokens_per_bs * bs, + num_token_non_padded_cpu=self.num_tokens_per_bs * bs, + return_hidden_states_before_norm=True, + ) + return forward_batch + + def capture_one_batch_size(self, bs: int, forward: Callable, stream_idx: int = 0): + graph = self._create_graph() + stream = self.stream + + self.deepep_adapter.capture(is_extend_in_batch=True) + + num_tokens = bs * self.num_tokens_per_bs + forward_batch = self.get_forward_batch(bs) + + self.mtp_worker.draft_extend_attn_backend_list[ + self.step + ].init_forward_metadata_capture_cuda_graph( + bs=bs, + num_tokens=num_tokens, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens=forward_batch.seq_lens, + encoder_lens=None, + forward_mode=self.forward_mode, + spec_info=forward_batch.spec_info, + ) + + # Run and capture + def run_once(): + # Clean intermediate result cache for DP attention + forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None + set_dp_buffer_len( + forward_batch.global_dp_buffer_len, + num_tokens, + forward_batch.dp_padding_mode.is_max_len(), + ) + set_is_extend_in_batch(False) + + # Backup two fields, which will be modified in-place in `draft_forward`. + output_cache_loc_backup = forward_batch.out_cache_loc + hidden_states_backup = forward_batch.spec_info.hidden_states + + ret = self.model_runner.model.forward( + forward_batch.input_ids, + forward_batch.positions, + forward_batch, + ) + + select_index = ( + torch.arange(bs, device=self.model_runner.device) + * (self.speculative_num_draft_tokens + self.step) + + self.accept_length[:bs] + - 1 + + self.step + ) + + probs = torch.softmax(ret.next_token_logits[select_index], dim=-1) + ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1) + + if self.next_cuda_graph_runner is not None: + padding_lens = ( + self.speculative_num_draft_tokens - self.accept_length[:bs] + ) + assign_new_state_triton( + ret.topk_index, + self.input_ids, + self.positions, + self.hidden_states, + self.out_cache_loc, + self.extend_seq_lens, + self.extend_start_loc, + self.next_cuda_graph_runner.input_ids, + self.next_cuda_graph_runner.positions, + self.next_cuda_graph_runner.hidden_states, + self.next_cuda_graph_runner.out_cache_loc, + self.next_cuda_graph_runner.extend_seq_lens, + self.next_cuda_graph_runner.extend_start_loc, + self.next_cuda_graph_runner.seq_lens, + padding_lens, + forward_batch.batch_size, + self.step, + forward_batch.req_pool_indices, + forward_batch.req_to_token_pool.req_to_token, + self.mtp_worker.req_to_hidden_states_pool, + ) + self.next_cuda_graph_runner.swa_out_cache_loc.copy_( + self.model_runner.token_to_kv_pool.translate_loc_from_full_to_swa( + self.next_cuda_graph_runner.out_cache_loc + ) + ) + + forward_batch.out_cache_loc = output_cache_loc_backup + forward_batch.spec_info.hidden_states = hidden_states_backup + return ret + + self._capture_init(run_once) + + out = self._capture_graph( + graph, get_global_graph_memory_pool(), stream, run_once + ) + + set_global_graph_memory_pool(graph.pool()) + return graph, out + + def init_replay_state( + self, forward_batch: ForwardBatch, bs: int, raw_bs: int, num_tokens: int + ): + # Common inputs + self.input_ids[:num_tokens].copy_(forward_batch.input_ids) + self.seq_lens[:raw_bs].copy_(forward_batch.seq_lens) + if forward_batch.extend_seq_lens is not None: + self.extend_seq_lens[:raw_bs].copy_(forward_batch.extend_seq_lens) + self.extend_start_loc[:raw_bs].copy_(forward_batch.extend_start_loc) + self.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc) + self.positions[:num_tokens].copy_(forward_batch.positions) + if ( + forward_batch.spec_info.hidden_states.shape[1] + == self.hidden_states.shape[1] + ): + self.hidden_states[:num_tokens].copy_(forward_batch.spec_info.hidden_states) + if forward_batch.spec_info.accept_length is not None: + self.accept_length[:raw_bs].copy_(forward_batch.spec_info.accept_length) + self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) + + if forward_batch.seq_lens_cpu is not None: + if bs != raw_bs: + self.seq_lens_cpu.fill_(self.seq_len_fill_value) + self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu) + + if forward_batch.extend_seq_lens_cpu is not None: + self.extend_seq_lens_cpu[:raw_bs] = forward_batch.extend_seq_lens_cpu + + def replay(self, forward_batch: ForwardBatch, init_state: bool = True): + assert forward_batch.out_cache_loc is not None + self.deepep_adapter.replay() + + # batch_size and num_seqs can be different in case there are finished examples + # in the batch, which will not be counted as num_seqs + raw_bs = forward_batch.batch_size + num_tokens = raw_bs * self.num_tokens_per_bs + # num_tokens = forward_batch.input_ids.shape[0] + if self.require_mlp_tp_gather: + max_batch_size = max(forward_batch.original_global_num_tokens_cpu) + index = bisect.bisect_left(self.capture_bs, max_batch_size) + else: + index = bisect.bisect_left(self.capture_bs, raw_bs) + + bs = self.capture_bs[index] + + if init_state: + self.init_replay_state(forward_batch, bs, raw_bs, num_tokens) + + if self.require_gathered_buffer: + self.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs) + self.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs) + + forward_batch.spec_info.hidden_states = self.hidden_states[:num_tokens] + forward_batch.spec_info.accept_length = self.accept_length[:bs] + forward_batch.spec_info.num_tokens_per_batch = self.num_tokens_per_bs + forward_batch.spec_info.num_tokens_for_logprob_per_batch = 1 + forward_batch.spec_info.positions = self.positions[:num_tokens] + forward_batch.spec_info.extend_seq_lens_tensor = self.extend_seq_lens[:bs] + + self.mtp_worker.draft_extend_attn_backend_list[ + self.step + ].init_forward_metadata_replay_cuda_graph( + bs=bs, + req_pool_indices=self.req_pool_indices, + seq_lens=self.seq_lens, + seq_lens_sum=forward_batch.seq_lens_sum + + (bs - raw_bs) * self.seq_len_fill_value, + encoder_lens=None, + forward_mode=self.forward_mode, + spec_info=forward_batch.spec_info, + seq_lens_cpu=self.seq_lens_cpu, + ) + + # Replay + self.raw_bs = raw_bs + self.bs = bs + self._replay(forward_batch) + out = self.output_buffers[bs] + + if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2: + # DRAFT_EXTEND_V2: all tokens calculations whether accepted or not. + unpadding_bs = num_tokens + elif bs != raw_bs: + forward_batch.spec_info.accept_length = self.accept_length[:raw_bs] + unpadding_bs = raw_bs + else: + unpadding_bs = None + + if unpadding_bs is not None: + out_copy = out + out = LogitsProcessorOutput( + next_token_logits=out.next_token_logits[:unpadding_bs], + hidden_states=out.hidden_states[:unpadding_bs], + ) + out.topk_p = out_copy.topk_p[:raw_bs] + out.topk_index = out_copy.topk_index[:raw_bs] + return out + + +class MTPMultiStepDraftExtendCudaGraphRunner: + def __init__(self, mtp_worker: MTPDraftWorker): + self.mtp_worker = mtp_worker + self.device = mtp_worker.device + self.gpu_id = mtp_worker.gpu_id + self.speculative_num_steps = mtp_worker.speculative_num_steps + self.draft_extend_attn_backend_list = mtp_worker.draft_extend_attn_backend_list + + self.runners = [] + self.cuda_graph_buffers = {} + self.seq_len_fill_value = 1 + self.max_bs = 1 + self.offsets = [0] + + self._init_and_capture() + + def _init_and_capture(self): + if self.mtp_worker.server_args.disable_cuda_graph: + self.runners = [None] * self.speculative_num_steps + return + + self.runners = [] + buffer_len_list = [] + + # 1. Capture loop + for step in range(self.speculative_num_steps): + if self.draft_extend_attn_backend_list[step]: + runner = MTPDraftExtendCudaGraphRunner(self.mtp_worker, step) + self.runners.append(runner) + + self.seq_len_fill_value = runner.seq_len_fill_value + self.max_bs = runner.max_bs + buffer_len_list.append(runner.max_num_token) + self.offsets.append(self.offsets[-1] + runner.max_num_token) + else: + self.runners.append(None) + + # 2. Allocate buffers + self.cuda_graph_buffers["seq_lens_cpu"] = torch.full( + (self.max_bs,), + self.seq_len_fill_value, + dtype=torch.int32, + ) + + with torch.device(self.device): + # Sliced buffers + self.cuda_graph_buffers["input_ids"] = torch.zeros( + (self.offsets[-1],), dtype=torch.int64 + ) + self.cuda_graph_buffers["out_cache_loc"] = torch.ones( + (self.offsets[-1],), dtype=torch.int64 + ) + self.cuda_graph_buffers["swa_out_cache_loc"] = torch.ones( + (self.offsets[-1],), dtype=torch.int64 + ) + self.cuda_graph_buffers["positions"] = torch.zeros( + (self.offsets[-1],), dtype=torch.int64 + ) + + # Shared states + self.cuda_graph_buffers["seq_lens"] = torch.full( + (self.max_bs,), + self.seq_len_fill_value, + dtype=torch.int32, + ) + self.cuda_graph_buffers["req_pool_indices"] = torch.zeros( + (self.max_bs,), dtype=torch.int32 + ) + self.cuda_graph_buffers["accept_length"] = torch.full( + (self.max_bs,), 1, dtype=torch.int32 + ) + + for step in range(self.speculative_num_steps - 1, -1, -1): + if self.runners[step] is not None: + tic = time.perf_counter() + before_mem = get_available_gpu_memory(self.device, self.gpu_id) + logger.info( + f"Capture draft extend cuda graph begin (step {step}). This can take up to several minutes. avail mem={before_mem:.2f} GB" + ) + + self.runners[step].init_buffers_and_capture( + self.cuda_graph_buffers, + self.offsets[step], + ( + self.runners[step + 1] + if step + 1 < self.speculative_num_steps + else None + ), + ) + + after_mem = get_available_gpu_memory(self.device, self.gpu_id) + logger.info( + f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB." + ) + + def reset_buffers(self, forward_batch, batch_result): + self.cuda_graph_buffers["input_ids"].zero_() + self.cuda_graph_buffers["seq_lens"].fill_(self.seq_len_fill_value) + self.cuda_graph_buffers["out_cache_loc"].zero_() + self.cuda_graph_buffers["swa_out_cache_loc"].zero_() + self.cuda_graph_buffers["positions"].zero_() + self.cuda_graph_buffers["accept_length"][: forward_batch.batch_size].copy_( + batch_result.accept_lens + ) + + def get_runner(self, step): + return self.runners[step] + + def get_last_runner(self): + return self.runners[-1] if self.runners else None + + def can_run(self, forward_batch): + return self.runners[0].can_run(forward_batch) diff --git a/python/sglang/srt/speculative/mtp_utils.py b/python/sglang/srt/speculative/mtp_utils.py new file mode 100644 index 000000000..f1ce9d4b0 --- /dev/null +++ b/python/sglang/srt/speculative/mtp_utils.py @@ -0,0 +1,350 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import torch +import triton +import triton.language as tl + + +@triton.jit +def rotate_input_ids_kernel( + input_ids_ptr, + extend_start_loc_ptr, + extend_seq_lens_ptr, + topk_index_ptr, + select_index_ptr, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + + start_loc = tl.load(extend_start_loc_ptr + pid) + seq_len = tl.load(extend_seq_lens_ptr + pid) + new_token = tl.load(topk_index_ptr + pid) + + num_elements_to_shift = seq_len - 1 + + for off in range(0, num_elements_to_shift, BLOCK_SIZE): + offsets = off + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_elements_to_shift + + read_ptr = input_ids_ptr + start_loc + offsets + 1 + val = tl.load(read_ptr, mask=mask) + tl.debug_barrier() + + write_ptr = input_ids_ptr + start_loc + offsets + tl.store(write_ptr, val, mask=mask) + tl.debug_barrier() + + if seq_len > 0: + if select_index_ptr is not None: + last_pos_ptr = input_ids_ptr + tl.load(select_index_ptr + pid) + else: + last_pos_ptr = input_ids_ptr + start_loc + seq_len - 1 + tl.store(last_pos_ptr, new_token) + + +def rotate_input_ids_triton( + input_ids, extend_start_loc, extend_seq_lens, topk_index, select_index=None +): + batch_size = extend_seq_lens.shape[0] + BLOCK_SIZE = 4096 if select_index is not None else 8 + grid = (batch_size,) + + rotate_input_ids_kernel[grid]( + input_ids, + extend_start_loc, + extend_seq_lens, + topk_index, + select_index, + BLOCK_SIZE=BLOCK_SIZE, + ) + return input_ids + + +@triton.jit +def assign_new_state_kernel( + # Source pointers + old_input_ids_ptr, + old_positions_ptr, + old_hidden_states_ptr, + old_out_cache_loc_ptr, + old_extend_seq_lens_ptr, + old_extend_start_loc_ptr, + # Destination pointers + input_ids_ptr, + positions_ptr, + hidden_states_ptr, + out_cache_loc_ptr, + extend_seq_lens_ptr, + extend_start_loc_ptr, + # Auxiliary data pointers + next_token_ids_ptr, + seq_lens_ptr, + padding_lens_ptr, + req_pool_indices_ptr, + req_to_token_ptr, + req_to_hidden_states_pool_ptr, + # Scalars and Strides + step, + stride_hidden_seq, + stride_hidden_dim, # hidden_states strides + stride_pool_req, + stride_pool_step, + stride_pool_dim, # pool strides + stride_req_token_0, + stride_req_token_1, # req_to_token strides + # Meta-parameters + HIDDEN_DIM: tl.constexpr, + BLOCK_SEQ: tl.constexpr, + BLOCK_HID: tl.constexpr, +): + pid = tl.program_id(0) + + seq_len: tl.tensor = tl.load(seq_lens_ptr + pid) + old_extend_len = tl.load(old_extend_seq_lens_ptr + pid) + old_start = tl.load(old_extend_start_loc_ptr + pid) + new_extend_len = old_extend_len + 1 + new_start = old_start + pid + + tl.store(extend_seq_lens_ptr + pid, new_extend_len) + tl.store(extend_start_loc_ptr + pid, new_start) + + offs_seq = tl.arange(0, BLOCK_SEQ) + mask_seq = offs_seq < old_extend_len + + old_ids = tl.load(old_input_ids_ptr + old_start + offs_seq, mask=mask_seq) + tl.store(input_ids_ptr + new_start + offs_seq, old_ids, mask=mask_seq) + padding_len = tl.load(padding_lens_ptr + pid) + tl.store( + input_ids_ptr + new_start + old_extend_len - padding_len, + tl.load(next_token_ids_ptr + pid), + ) + + old_pos = tl.load(old_positions_ptr + old_start + offs_seq, mask=mask_seq) + tl.store(positions_ptr + new_start + 1 + offs_seq, old_pos, mask=mask_seq) + tl.store( + positions_ptr + new_start, max(tl.load(old_positions_ptr + old_start) - 1, 0) + ) + + old_cache = tl.load(old_out_cache_loc_ptr + old_start + offs_seq, mask=mask_seq) + tl.store(out_cache_loc_ptr + new_start + 1 + offs_seq, old_cache, mask=mask_seq) + + req_idx = tl.load(req_pool_indices_ptr + pid) + token_idx_col = seq_len - old_extend_len - 1 + if token_idx_col >= 0: + req_token_ptr_loc = ( + req_to_token_ptr + + (req_idx * stride_req_token_0) + + (token_idx_col * stride_req_token_1) + ) + last_cache_loc = tl.load(req_token_ptr_loc) + tl.store(out_cache_loc_ptr + new_start, last_cache_loc) + + pool_vec_offset_base = ((req_idx + 1) * stride_pool_req) + ( + -(step + 1) * stride_pool_step + ) + + for off_h in range(0, HIDDEN_DIM, BLOCK_HID): + offs_h = off_h + tl.arange(0, BLOCK_HID) + mask_h = offs_h < HIDDEN_DIM + + for i in range(BLOCK_SEQ): + if i < old_extend_len: + old_h_ptr = ( + old_hidden_states_ptr + + (old_start + i) * stride_hidden_seq + + (offs_h * stride_hidden_dim) + ) + new_h_ptr = ( + hidden_states_ptr + + (new_start + 1 + i) * stride_hidden_seq + + (offs_h * stride_hidden_dim) + ) + + chunk_old = tl.load(old_h_ptr, mask=mask_h) + tl.store(new_h_ptr, chunk_old, mask=mask_h) + + pool_ptrs = ( + req_to_hidden_states_pool_ptr + + pool_vec_offset_base + + (offs_h * stride_pool_dim) + ) + pool_val = tl.load(pool_ptrs, mask=mask_h) + + new_h_start_ptrs = ( + hidden_states_ptr + + (new_start * stride_hidden_seq) + + (offs_h * stride_hidden_dim) + ) + tl.store(new_h_start_ptrs, pool_val, mask=mask_h) + + +def assign_new_state_triton( + next_token_ids: torch.Tensor, + old_input_ids: torch.Tensor, + old_positions: torch.Tensor, + old_hidden_states: torch.Tensor, + old_out_cache_loc: torch.Tensor, + old_extend_seq_lens: torch.Tensor, + old_extend_start_loc: torch.Tensor, + input_ids: torch.Tensor, + positions: torch.Tensor, + hidden_states: torch.Tensor, + out_cache_loc: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + padding_lens: torch.Tensor, + num_seqs: int, + step: int, + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + req_to_hidden_states_pool: torch.Tensor, +): + """ + Wrapper function to calculate offsets and launch the Triton kernel. + """ + hidden_dim = hidden_states.shape[1] + + BLOCK_SEQ = 8 + BLOCK_HID = 64 + + grid = (num_seqs,) + + assign_new_state_kernel[grid]( + # Pointers + old_input_ids, + old_positions, + old_hidden_states, + old_out_cache_loc, + old_extend_seq_lens, + old_extend_start_loc, + input_ids, + positions, + hidden_states, + out_cache_loc, + extend_seq_lens, + extend_start_loc, + next_token_ids, + seq_lens, + padding_lens, + req_pool_indices, + req_to_token, + req_to_hidden_states_pool, + # Constants/Strides + step, + old_hidden_states.stride(0), + old_hidden_states.stride(1), + req_to_hidden_states_pool.stride(0), + req_to_hidden_states_pool.stride(1), + req_to_hidden_states_pool.stride(2), + req_to_token.stride(0), + req_to_token.stride(1), + # Meta + HIDDEN_DIM=hidden_dim, + BLOCK_SEQ=BLOCK_SEQ, + BLOCK_HID=BLOCK_HID, + ) + + +@triton.jit +def assign_hidden_states_pool_kernel( + hidden_states_ptr, + req_pool_indices_ptr, + req_to_hidden_states_pool_ptr, + extend_seq_lens_ptr, + extend_start_loc_ptr, + stride_hidden_seq, + stride_hidden_dim, + stride_pool_req, + stride_pool_step, + stride_pool_dim, + HIDDEN_DIM: tl.constexpr, + pool_size: tl.constexpr, + BLOCK_HID: tl.constexpr, +): + pid = tl.program_id(0) + + extend_len = tl.load(extend_seq_lens_ptr + pid) + start_loc = tl.load(extend_start_loc_ptr + pid) + end_loc = start_loc + extend_len + + req_idx = tl.load(req_pool_indices_ptr + pid) + pool_vec_offset_base = req_idx * stride_pool_req + + for i in range(pool_size): + for off_h in range(0, HIDDEN_DIM, BLOCK_HID): + offs_h = off_h + tl.arange(0, BLOCK_HID) + mask_h = offs_h < HIDDEN_DIM + + hid_ptr = ( + hidden_states_ptr + + (end_loc - pool_size + i) * stride_hidden_seq + + offs_h * stride_hidden_dim + ) + hid_val = tl.load(hid_ptr, mask=mask_h) + + pool_ptr = ( + req_to_hidden_states_pool_ptr + + pool_vec_offset_base + + i * stride_pool_step + + offs_h * stride_pool_dim + ) + tl.store(pool_ptr, hid_val, mask=mask_h) + + +def assign_hidden_states_pool_triton( + hidden_states: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_hidden_states_pool: torch.Tensor, + pool_size: int, + num_seqs: int, + extend_seq_lens: torch.Tensor, + extend_start_loc: torch.Tensor, +): + grid = (num_seqs,) + assign_hidden_states_pool_kernel[grid]( + hidden_states, + req_pool_indices, + req_to_hidden_states_pool, + extend_seq_lens, + extend_start_loc, + hidden_states.stride(0), + hidden_states.stride(1), + req_to_hidden_states_pool.stride(0), + req_to_hidden_states_pool.stride(1), + req_to_hidden_states_pool.stride(2), + HIDDEN_DIM=hidden_states.shape[1], + pool_size=pool_size, + BLOCK_HID=64, + ) + + +def assign_hidden_states_pool_torch( + hidden_states: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_hidden_states_pool: torch.Tensor, + pool_size: int, + num_seqs: int, + extend_seq_lens: torch.Tensor, + extend_start_loc: torch.Tensor, +): + for req in range(num_seqs): + pool_idx = req_pool_indices[req] + extend_len = extend_seq_lens[req] + start_loc = extend_start_loc[req] + end_loc = start_loc + extend_len + req_to_hidden_states_pool[pool_idx, :pool_size, :].copy_( + hidden_states[end_loc - pool_size : end_loc, :] + ) diff --git a/python/sglang/srt/speculative/mtp_worker.py b/python/sglang/srt/speculative/mtp_worker.py new file mode 100644 index 000000000..233fd1992 --- /dev/null +++ b/python/sglang/srt/speculative/mtp_worker.py @@ -0,0 +1,989 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import logging +import time +from typing import List, Optional, Tuple + +import torch + +from sglang.srt.distributed import get_tp_group +from sglang.srt.layers.dp_attention import get_attention_tp_group +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.moe.utils import speculative_moe_backend_context +from sglang.srt.layers.sampler import get_token_ids_logprobs, get_top_logprobs +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.managers.scheduler import GenerationBatchResult +from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.mem_cache.chunk_cache import SWAChunkCache +from sglang.srt.mem_cache.common import ( + alloc_paged_token_slots_extend, + alloc_token_slots, +) +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ForwardMode, +) +from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative.draft_utils import DraftBackendFactory +from sglang.srt.speculative.eagle_info import ( + EagleDraftInput, + EagleVerifyInput, + EagleVerifyOutput, +) +from sglang.srt.speculative.eagle_utils import ( + build_tree_kernel_efficient, + organize_draft_results, +) +from sglang.srt.speculative.eagle_worker import get_last_loc_large_page_size_top_k_1 +from sglang.srt.speculative.mtp_draft_extend_cuda_graph_runner import ( + MTPDraftExtendCudaGraphRunner, +) +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.speculative.spec_utils import ( + assign_draft_cache_locs, + detect_nan, + draft_tp_context, + fast_topk, + generate_token_bitmask, + get_last_loc_large_page_size_large_top_k, + load_token_map, + select_top_k_tokens, +) +from sglang.srt.utils import ( + empty_context, + get_available_gpu_memory, + get_bool_env_var, + is_cuda, + is_npu, + next_power_of_2, +) + +_is_npu = is_npu() + +if is_cuda(): + from sgl_kernel import segment_packbits # noqa: F401 + +logger = logging.getLogger(__name__) +SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB") + + +class MTPWorker(TpModelWorker): + + def __init__( + self, + server_args: ServerArgs, + gpu_id: int, + tp_rank: int, + dp_rank: Optional[int], + moe_ep_rank: int, + nccl_port: int, + target_worker: TpModelWorker, + ): + # Parse arguments + self.server_args = server_args + self.topk = server_args.speculative_eagle_topk + self.speculative_num_steps = server_args.speculative_num_steps + self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + self.enable_nan_detection = server_args.enable_nan_detection + self.gpu_id = gpu_id + self.device = server_args.device + self.target_worker = target_worker + self.page_size = server_args.page_size + self.speculative_algorithm = SpeculativeAlgorithm.from_string( + server_args.speculative_algorithm + ) + self.draft_extend_attn_backend_list = [] + + # Override the context length of the draft model to be the same as the target model. + server_args.context_length = target_worker.model_runner.model_config.context_len + + # Do not capture cuda graph in `super().__init__()` + # It will be captured later. + backup_disable_cuda_graph = server_args.disable_cuda_graph + server_args.disable_cuda_graph = True + # Share the allocator with a target worker. + # Draft and target worker own their own KV cache pools. + self.req_to_token_pool, self.token_to_kv_pool_allocator = ( + target_worker.get_memory_pool() + ) + + # Load hot token ids + if self.speculative_algorithm.is_eagle3(): + if server_args.speculative_token_map is not None: + logger.warning( + "Speculative token map specified, but EAGLE3 models already have this. Ignoring the specified token map." + ) + self.hot_token_id = None + elif server_args.speculative_token_map is not None: + self.hot_token_id = load_token_map(server_args.speculative_token_map) + server_args.json_model_override_args = ( + f'{{"hot_vocab_size": {len(self.hot_token_id)}}}' + ) + else: + self.hot_token_id = None + + # Init draft worker + if server_args.enable_dp_attention and self.speculative_algorithm.is_eagle3(): + ctx = draft_tp_context(get_attention_tp_group()) + else: + ctx = empty_context() + with ctx, speculative_moe_backend_context(): + super().__init__( + server_args=server_args, + gpu_id=gpu_id, + tp_rank=tp_rank, + pp_rank=0, # FIXME + dp_rank=dp_rank, + moe_ep_rank=moe_ep_rank, + nccl_port=nccl_port, + is_draft_worker=True, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + is_mtp_worker=True, + ) + + embed, head = self.target_worker.model_runner.model.get_embed_and_head() + + if self.speculative_algorithm.is_eagle3(): + # most cases EAGLE3 models don't share lm_head + # but some models (e.g. nvidia/gpt-oss-120b-Eagle3) shares + if ( + hasattr(self.draft_model_runner.model, "load_lm_head_from_target") + and self.draft_model_runner.model.load_lm_head_from_target + ): + self.draft_model_runner.model.set_embed_and_head(embed, head) + else: + self.draft_model_runner.model.set_embed(embed) + + # grab hot token ids + if self.draft_model_runner.model.hot_token_id is not None: + self.hot_token_id = self.draft_model_runner.model.hot_token_id.to( + embed.device + ) + + else: + if self.hot_token_id is not None: + head = head.clone() + self.hot_token_id = self.hot_token_id.to(head.device) + head.data = head.data[self.hot_token_id] + + # Share the embedding and lm_head + for i in range(self.speculative_num_steps): + self.mtp_model_runner(i).model.set_embed_and_head(embed, head) + + # Init attention backend and cuda graphs + for i in range(self.speculative_num_steps): + self.mtp_model_runner(i).server_args.disable_cuda_graph = ( + backup_disable_cuda_graph + ) + self.draft_tp_context = ( + draft_tp_context if server_args.enable_dp_attention else empty_context + ) + with self.draft_tp_context( + self.mtp_model_runner(0).tp_group + ), speculative_moe_backend_context(): + self.init_attention_backend() + self.init_cuda_graphs() + + # Some dummy tensors + self.num_new_pages_per_topk = torch.empty( + (), dtype=torch.int64, device=self.device + ) + self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device) + + def init_attention_backend(self): + # Create multi-step attn backends and cuda graph runners + for step in range(self.speculative_num_steps): + draft_backend_factory = DraftBackendFactory( + self.server_args, + self.mtp_model_runner(step), + self.topk, + self.speculative_num_steps, + ) + + # Initialize draft extend attention backend (respects speculative_attention_mode setting) + self.draft_extend_attn_backend_list.append( + draft_backend_factory.create_draft_extend_backend() + ) + + def init_cuda_graphs(self): + """Capture cuda graphs.""" + self.cuda_graph_runner_for_draft_extend_list = [] + + if self.server_args.disable_cuda_graph: + return + + # Capture extend + for step in range(self.speculative_num_steps): + if self.draft_extend_attn_backend_list[step] and not _is_npu: + tic = time.perf_counter() + before_mem = get_available_gpu_memory(self.device, self.gpu_id) + logger.info( + f"Capture draft extend cuda graph begin. This can take up to several minutes. avail mem={before_mem:.2f} GB" + ) + self.cuda_graph_runner_for_draft_extend_list.append( + MTPDraftExtendCudaGraphRunner(self, step) + ) + after_mem = get_available_gpu_memory(self.device, self.gpu_id) + logger.info( + f"Capture draft extend cuda graph end. Time elapsed: {time.perf_counter() - tic:.2f} s. mem usage={(before_mem - after_mem):.2f} GB. avail mem={after_mem:.2f} GB." + ) + + def mtp_model_runner(self, layer_id: int): + return self.model_runner_list[layer_id] + + def forward_batch_generation(self, batch: ScheduleBatch) -> GenerationBatchResult: + """Run speculative decoding forward. + + NOTE: Many states of batch is modified as you go through. It is not guaranteed that + the final output batch have the same state as the input. + + Args: + batch: The batch to run forward. The state of the batch is modified as it runs. + Returns: + A tuple of the final logit output of the target model, next tokens accepted, + the batch id (used for overlap schedule), and number of accepted tokens. + """ + if batch.forward_mode.is_extend() or batch.is_extend_in_batch: + logits_output, next_token_ids, seq_lens_cpu = self.forward_target_extend( + batch + ) + with self.draft_tp_context( + self.mtp_model_runner(0).tp_group + ), speculative_moe_backend_context(): + self.forward_draft_extend( + batch, logits_output.hidden_states, next_token_ids, seq_lens_cpu + ) + return GenerationBatchResult( + logits_output=logits_output, + next_token_ids=next_token_ids, + num_accepted_tokens=0, + can_run_cuda_graph=False, + ) + else: + with self.draft_tp_context( + self.mtp_model_runner(0).tp_group + ), speculative_moe_backend_context(): + spec_info = self.draft(batch) + logits_output, verify_output, model_worker_batch, can_run_cuda_graph = ( + self.verify(batch, spec_info) + ) + + with self.draft_tp_context( + self.mtp_model_runner(0).tp_group + ), speculative_moe_backend_context(): + # NOTE: We should use `check_forward_draft_extend_after_decode` + # when DP attention is enabled, but it is slow. Skip it for now. + if ( + self.server_args.enable_dp_attention + or batch.spec_info.verified_id.shape[0] > 0 + ): + # decode is not finished + self.forward_draft_extend_after_decode(batch) + + return GenerationBatchResult( + logits_output=logits_output, + next_token_ids=verify_output.verified_id, + num_accepted_tokens=sum(verify_output.accept_length_per_req_cpu), + can_run_cuda_graph=can_run_cuda_graph, + ) + + def check_forward_draft_extend_after_decode(self, batch: ScheduleBatch): + local_need_forward = batch.spec_info.verified_id.shape[0] > 0 + if not self.server_args.enable_dp_attention: + return local_need_forward + + global_need_forward = torch.tensor( + [ + (local_need_forward), + ], + dtype=torch.int64, + ) + torch.distributed.all_reduce( + global_need_forward, group=get_tp_group().cpu_group + ) + global_need_forward_cnt = global_need_forward[0].item() + need_forward = global_need_forward_cnt > 0 + return need_forward + + def forward_target_extend( + self, batch: ScheduleBatch + ) -> Tuple[LogitsProcessorOutput, torch.Tensor, int, Optional[torch.Tensor]]: + """Run the target extend. + + Args: + batch: The batch to run. States could be modified. + + Returns: + logits_output: The output of logits. It will contain the full hidden states. + next_token_ids: Next token ids generated. + """ + # Forward with the target model and get hidden states. + # We need the full hidden states to prefill the KV cache of the draft model. + model_worker_batch = batch.get_model_worker_batch() + model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL + model_worker_batch.return_hidden_states_before_norm = True + batch_result = self.target_worker.forward_batch_generation(model_worker_batch) + logits_output, next_token_ids = ( + batch_result.logits_output, + batch_result.next_token_ids, + ) + return ( + logits_output, + next_token_ids, + model_worker_batch.seq_lens_cpu, + ) + + def _draft_preprocess_decode(self, batch: ScheduleBatch): + if isinstance(batch.tree_cache, SWAChunkCache): + for req in batch.reqs: + batch.tree_cache.evict_swa( + req, req.seqlen - 1, batch.model_config.attention_chunk_size + ) + + # Parse args + num_seqs = batch.batch_size() + spec_info = batch.spec_info + + # Accumulate penalty + if batch.sampling_info.penalizer_orchestrator.is_required: + # This is a relaxed version of penalties for speculative decoding. + batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens( + spec_info.verified_id.to(torch.int64) + ) + + # Allocate cache locations + # Layout of the out_cache_loc + # [ topk 0 ] [ topk 1 ] + # [iter=0, iter=1, iter=2] [iter=0, iter=1, iter=2] + if self.page_size == 1: + out_cache_loc, token_to_kv_pool_state_backup = alloc_token_slots( + batch.tree_cache, + num_seqs * self.speculative_num_steps * self.topk, + backup_state=True, + ) + duplicate_cache_len = 0 + source_cache_loc, target_cache_loc, last_page_lens_cumsum = None, None, None + else: + if self.topk == 1: + prefix_lens, seq_lens, last_loc = get_last_loc_large_page_size_top_k_1( + batch.req_to_token_pool.req_to_token, + batch.req_pool_indices, + batch.seq_lens, + self.speculative_num_steps, + ) + prefix_lens_cpu = batch.seq_lens_cpu + seq_lens_cpu = batch.seq_lens_cpu + self.speculative_num_steps + extend_num_tokens = num_seqs * self.speculative_num_steps + duplicate_cache_len = 0 + source_cache_loc, target_cache_loc, last_page_lens_cumsum = ( + None, + None, + None, + ) + else: + # In this case, the last partial page needs to be duplicated. + # KV cache layout in batch.req_to_token_pool.req_to_token: + # + # | -------- | -- xxxx .. | -- xxxx .. | -- xxxx .. | + # prefix top-k = 0 tok-k = 1 top-k = 2 + # + # "-" means prefix tokens + # "x" means speculative draft tokens + # "." means padded tokens + + # TODO(lmzheng): The current implementation is still a fake support + # for page size > 1. In the `assign_draft_cache_locs` below, + # we directly move the indices instead of the real kv cache. + # This only works when the kernel backend runs with page size = 1. + # If the kernel backend runs with page size > 1, we need to + # duplicate the real KV cache. The overhead of duplicating KV + # cache seems okay because the draft KV cache only has one layer. + # see a related copy operation in MHATokenToKVPool::move_kv_cache. + + ( + prefix_lens, + seq_lens, + last_loc, + self.num_new_pages_per_topk, + self.extend_lens, + _, + ) = get_last_loc_large_page_size_large_top_k( + batch.req_to_token_pool.req_to_token, + batch.req_pool_indices, + batch.seq_lens, + self.speculative_num_steps, + self.topk, + self.page_size, + ) + prefix_lens_cpu = batch.seq_lens_cpu + last_page_lens = prefix_lens_cpu % self.page_size + num_new_pages_per_topk = ( + last_page_lens + self.speculative_num_steps + self.page_size - 1 + ) // self.page_size + seq_lens_cpu = ( + prefix_lens_cpu // self.page_size * self.page_size + + num_new_pages_per_topk * (self.page_size * self.topk) + ) + extend_num_tokens = torch.sum((seq_lens_cpu - prefix_lens_cpu)).item() + + out_cache_loc, token_to_kv_pool_state_backup = ( + alloc_paged_token_slots_extend( + batch.tree_cache, + prefix_lens, + prefix_lens_cpu, + seq_lens, + seq_lens_cpu, + last_loc, + extend_num_tokens, + backup_state=True, + ) + ) + last_page_lens_cumsum = torch.cumsum(last_page_lens, dim=0) + duplicate_cache_len = torch.sum(last_page_lens).item() * (self.topk - 1) + target_cache_loc = torch.zeros( + duplicate_cache_len, dtype=torch.int32, device=self.device + ) + source_cache_loc = torch.zeros( + duplicate_cache_len, dtype=torch.int32, device=self.device + ) + + assign_draft_cache_locs[(num_seqs,)]( + batch.req_pool_indices, + batch.req_to_token_pool.req_to_token, + batch.seq_lens, + self.extend_lens, + self.num_new_pages_per_topk, + out_cache_loc, + source_cache_loc, + target_cache_loc, + last_page_lens_cumsum, + duplicate_cache_len, + batch.req_to_token_pool.req_to_token.shape[1], + self.topk, + self.speculative_num_steps, + self.page_size, + next_power_of_2(num_seqs), + next_power_of_2(self.speculative_num_steps), + ) + + if self.page_size > 1 and self.topk > 1: + # Remove padded slots + out_cache_loc = out_cache_loc[ + : num_seqs * self.topk * self.speculative_num_steps + ] + + batch.out_cache_loc = out_cache_loc + batch.seq_lens_sum = torch.sum(batch.seq_lens).item() + batch.return_hidden_states = False + spec_info.positions = batch.seq_lens.repeat_interleave(self.topk, dim=0) + self.token_to_kv_pool_allocator.restore_state(token_to_kv_pool_state_backup) + + def _draft_preprocess_idle(self, batch: ScheduleBatch): + batch.spec_info = EagleDraftInput.create_idle_input( + device=self.device, + hidden_size=self.model_config.hidden_size, + dtype=self.model_config.dtype, + topk=self.topk * self.speculative_num_steps, + capture_hidden_mode=CaptureHiddenMode.LAST, + ) + + def draft(self, batch: ScheduleBatch): + # Parse args + if batch.forward_mode.is_idle(): + self._draft_preprocess_idle(batch) + else: + self._draft_preprocess_decode(batch) + + spec_info = batch.spec_info + assert isinstance(spec_info, EagleDraftInput) + + spec_info.capture_hidden_mode = CaptureHiddenMode.LAST + spec_info.num_tokens_per_batch = self.topk + spec_info.num_tokens_for_logprob_per_batch = self.topk + batch.return_hidden_states = False + + # Get forward batch + model_worker_batch = batch.get_model_worker_batch() + assert model_worker_batch.capture_hidden_mode == CaptureHiddenMode.LAST + forward_batch = ForwardBatch.init_new( + model_worker_batch, self.mtp_model_runner(0) + ) + forward_batch.can_run_dp_cuda_graph = False + forward_batch.return_hidden_states_before_norm = True + + # Parse args + assert isinstance(spec_info, EagleDraftInput) + topk_p, topk_index, hidden_states = ( + spec_info.topk_p, + spec_info.topk_index, + spec_info.hidden_states, + ) + + # Return values + score_list: List[torch.Tensor] = [] + token_list: List[torch.Tensor] = [] + parents_list: List[torch.Tensor] = [] + + # Forward multiple steps + scores = None + input_ids, hidden_states, scores, tree_info = select_top_k_tokens( + 0, topk_p, topk_index, hidden_states, scores, self.topk + ) + if self.speculative_num_steps == 1: + score_list.append(tree_info[0]) + token_list.append(tree_info[1]) + parents_list.append(tree_info[2]) + else: + for i in range(self.speculative_num_steps): + score_list.append(tree_info[0][:, :, i].unsqueeze(-1)) + token_index = tree_info[1][:, i].unsqueeze(-1) + token_list.append(token_index) + if i == 0: + parents_list.append(tree_info[2]) + else: + parents_list.append( + torch.full( + (tree_info[2].size(0), 1), + i, + dtype=torch.long, + device=self.device, + ) + ) + + parent_list, top_scores_index, draft_tokens = organize_draft_results( + score_list, token_list, parents_list, self.speculative_num_draft_tokens + ) + + if batch.forward_mode.is_idle(): + return EagleVerifyInput.create_idle_input( + self.topk, + self.speculative_num_steps, + self.speculative_num_draft_tokens, + ) + + ( + tree_mask, + position, + retrive_index, + retrive_next_token, + retrive_next_sibling, + draft_tokens, + ) = build_tree_kernel_efficient( + spec_info.verified_id, + parent_list, + top_scores_index, + draft_tokens, + batch.seq_lens, + batch.seq_lens_sum, + self.topk, + self.speculative_num_steps, + self.speculative_num_draft_tokens, + ) + + return EagleVerifyInput( + draft_token=draft_tokens, + custom_mask=tree_mask, + positions=position, + retrive_index=retrive_index, + retrive_next_token=retrive_next_token, + retrive_next_sibling=retrive_next_sibling, + retrive_cum_len=None, + spec_steps=self.speculative_num_steps, + topk=self.topk, + draft_token_num=self.server_args.speculative_num_draft_tokens, + capture_hidden_mode=CaptureHiddenMode.FULL, + seq_lens_sum=forward_batch.seq_lens_sum, + seq_lens_cpu=forward_batch.seq_lens_cpu, + ) + + def clear_cache_pool(self): + # allocator and kv cache pool are shared with target worker + pass + + def verify(self, batch: ScheduleBatch, spec_info: EagleVerifyInput): + spec_info.prepare_for_verify(batch, self.page_size) + batch.return_hidden_states = False + batch.forward_mode = ( + ForwardMode.TARGET_VERIFY + if not batch.forward_mode.is_idle() + else ForwardMode.IDLE + ) + batch.spec_info = spec_info + + model_worker_batch = batch.get_model_worker_batch( + seq_lens_cpu_cache=spec_info.seq_lens_cpu + ) + assert model_worker_batch.capture_hidden_mode == spec_info.capture_hidden_mode + model_worker_batch.return_hidden_states_before_norm = True + + if batch.has_grammar: + retrieve_next_token_cpu = spec_info.retrive_next_token.cpu() + retrieve_next_sibling_cpu = spec_info.retrive_next_sibling.cpu() + draft_tokens_cpu = spec_info.draft_token.view( + spec_info.retrive_next_token.shape + ).cpu() + + # Forward + batch_result = self.target_worker.forward_batch_generation( + model_worker_batch, is_verify=True + ) + logits_output, can_run_cuda_graph = ( + batch_result.logits_output, + batch_result.can_run_cuda_graph, + ) + + vocab_mask = None + if batch.has_grammar: + # Generate the logit mask for structured output. + # Overlap the CPU operations for bitmask generation with the forward pass. + vocab_mask = generate_token_bitmask( + batch.reqs, + spec_info, + retrieve_next_token_cpu, + retrieve_next_sibling_cpu, + draft_tokens_cpu, + batch.sampling_info.vocab_size, + ) + + if vocab_mask is not None: + assert spec_info.grammar is not None + vocab_mask = vocab_mask.to(spec_info.retrive_next_token.device) + # NOTE (sk): otherwise, this vocab mask will be the one from the previous extend stage + # and will be applied to produce wrong results + batch.sampling_info.vocab_mask = None + + if self.enable_nan_detection: + detect_nan(logits_output) + + spec_info.hidden_states = logits_output.hidden_states + res: EagleVerifyOutput = spec_info.verify( + batch, + logits_output, + self.token_to_kv_pool_allocator, + self.page_size, + vocab_mask, + ) + + # Post process based on verified outputs. + # Pick indices that we care (accepted) + logits_output.next_token_logits = logits_output.next_token_logits[ + res.accepted_indices + ] + logits_output.hidden_states = logits_output.hidden_states[res.accepted_indices] + + if self.target_worker.model_runner.hybrid_gdn_config is not None: + accepted_length = ( + torch.tensor( + res.accept_length_per_req_cpu, + device=logits_output.hidden_states.device, + dtype=torch.int64, + ) + + 1 + ) + + # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask + # res.accepted_indices.shape[0] > 0 skips DP attn idle batch + if spec_info.topk > 1 and res.accepted_indices.shape[0] > 0: + # accepted_indices=[0,2,3,4,5,7,9,10,11], accepted_length=[4, 3, 2], cumulative_accepted_lengths=[4, 7, 9] + # first_token_indices_per_req=prepend(0, accepted_indices[cumulative_accepted_lengths[:-1]]) = [0, 5, 10] + # last_token_indices_per_req=accepted_indices[cumulative_accepted_lengths - 1] = [4, 9, 11] (last token ID of each req) + # max_relative_indices_per_req = [4,4,1]; those are the per-req spec-decoding step offsets that contain the correct mamba caches + cumulative_accepted_lengths = torch.cumsum(accepted_length, dim=0) + req_start_positions = torch.cat( + [ + torch.zeros( + 1, + dtype=cumulative_accepted_lengths.dtype, + device=cumulative_accepted_lengths.device, + ), + cumulative_accepted_lengths[:-1], + ] + ) + first_token_indices_per_req = res.accepted_indices[req_start_positions] + last_token_indices_per_req = res.accepted_indices[ + cumulative_accepted_lengths - 1 + ] + max_relative_indices_per_req = ( + last_token_indices_per_req - first_token_indices_per_req + ) + else: + max_relative_indices_per_req = accepted_length - 1 + self.target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify( + max_relative_indices_per_req, self.target_worker.model_runner.model + ) + + if batch.return_logprob: + self.add_logprob_values(batch, res, logits_output) + + # Prepare the batch for the next draft forwards. + batch.forward_mode = ( + ForwardMode.DECODE if not batch.forward_mode.is_idle() else ForwardMode.IDLE + ) + batch.spec_info = res.draft_input + + return logits_output, res, model_worker_batch, can_run_cuda_graph + + def add_logprob_values( + self, + batch: ScheduleBatch, + res: EagleVerifyOutput, + logits_output: LogitsProcessorOutput, + ): + # Extract args + logits_output = res.logits_output + top_logprobs_nums = batch.top_logprobs_nums + token_ids_logprobs = batch.token_ids_logprobs + accepted_indices = res.accepted_indices + assert len(accepted_indices) == len(logits_output.next_token_logits) + + temperatures = batch.sampling_info.temperatures + num_draft_tokens = batch.spec_info.draft_token_num + # acceptance indices are the indices in a "flattened" batch. + # dividing it to num_draft_tokens will yield the actual batch index. + temperatures = temperatures[accepted_indices // num_draft_tokens] + if SGLANG_RETURN_ORIGINAL_LOGPROB: + logprobs = torch.nn.functional.log_softmax( + logits_output.next_token_logits, dim=-1 + ) + else: + logprobs = torch.nn.functional.log_softmax( + logits_output.next_token_logits / temperatures, dim=-1 + ) + batch_next_token_ids = res.verified_id + num_tokens_per_req = [accept + 1 for accept in res.accept_length_per_req_cpu] + + # We should repeat top_logprobs_nums to match num_tokens_per_req. + top_logprobs_nums_repeat_interleaved = [] + token_ids_logprobs_repeat_interleaved = [] + for num, num_tokens in zip(top_logprobs_nums, num_tokens_per_req): + top_logprobs_nums_repeat_interleaved.extend([num] * num_tokens) + for token_ids, num_tokens in zip(token_ids_logprobs, num_tokens_per_req): + token_ids_logprobs_repeat_interleaved.extend([token_ids] * num_tokens) + + # Extract logprobs + if any(x > 0 for x in top_logprobs_nums): + ( + logits_output.next_token_top_logprobs_val, + logits_output.next_token_top_logprobs_idx, + ) = get_top_logprobs( + logprobs, + top_logprobs_nums_repeat_interleaved, + ) + + if any(x is not None for x in token_ids_logprobs): + ( + logits_output.next_token_token_ids_logprobs_val, + logits_output.next_token_token_ids_logprobs_idx, + ) = get_token_ids_logprobs( + logprobs, + token_ids_logprobs_repeat_interleaved, + ) + + logits_output.next_token_logprobs = logprobs[ + torch.arange(len(batch_next_token_ids), device=batch.sampling_info.device), + batch_next_token_ids, + ] + + # Add output logprobs to the request + pt = 0 + next_token_logprobs = logits_output.next_token_logprobs.tolist() + verified_ids = batch_next_token_ids.tolist() + for req, num_tokens in zip(batch.reqs, num_tokens_per_req, strict=True): + for _ in range(num_tokens): + if req.return_logprob: + req.output_token_logprobs_val.append(next_token_logprobs[pt]) + req.output_token_logprobs_idx.append(verified_ids[pt]) + if req.top_logprobs_num > 0: + req.output_top_logprobs_val.append( + res.logits_output.next_token_top_logprobs_val[pt] + ) + req.output_top_logprobs_idx.append( + res.logits_output.next_token_top_logprobs_idx[pt] + ) + pt += 1 + + def forward_draft_extend( + self, + batch: ScheduleBatch, + hidden_states: torch.Tensor, + next_token_ids: torch.Tensor, + seq_lens_cpu: Optional[torch.Tensor], + ): + """Run draft model extend. This API modifies the states of the batch. + + Args: + batch: The batch to run. + hidden_states: Hidden states from the target model forward + next_token_ids: Next token ids generated from the target forward. + """ + batch.spec_info = EagleDraftInput( + hidden_states=hidden_states, + verified_id=next_token_ids, + num_tokens_per_batch=1, + num_tokens_for_logprob_per_batch=1, + ) + batch.return_hidden_states = False + batch.spec_info.prepare_for_extend(batch) + batch.spec_info.capture_hidden_mode = CaptureHiddenMode.LAST + model_worker_batch = batch.get_model_worker_batch( + seq_lens_cpu_cache=seq_lens_cpu + ) + forward_batch = ForwardBatch.init_new( + model_worker_batch, self.mtp_model_runner(0) + ) + forward_batch.return_logprob = False + forward_batch.return_hidden_states_before_norm = True + topk_p_list = [] + topk_index_list = [] + for step in range(self.speculative_num_steps): + logits_output, _ = self.mtp_model_runner(step).forward(forward_batch) + if self.enable_nan_detection: + detect_nan(logits_output) + probs = torch.softmax(logits_output.next_token_logits, dim=-1) + topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + topk_p_list.append(topk_p) + topk_index_list.append(topk_index) + pt = 0 + if forward_batch.extend_seq_lens is not None: + for i, extend_len in enumerate(forward_batch.extend_seq_lens): + input_ids = forward_batch.input_ids[pt : pt + extend_len] + forward_batch.input_ids[pt : pt + extend_len] = torch.cat( + (input_ids[1:], topk_index[i].reshape(1)) + ) + pt += extend_len + + assert isinstance(forward_batch.spec_info, EagleDraftInput) + assert forward_batch.spec_info is batch.spec_info + forward_batch.spec_info.topk_p = torch.cat(topk_p_list, dim=1) + forward_batch.spec_info.topk_index = torch.cat(topk_index_list, dim=1) + has_finished, unfinished_req_index = False, [] + for i, req in enumerate(batch.reqs): + if req.finished(): + has_finished = True + else: + unfinished_req_index.append(i) + if has_finished: + unfinished_index_device = torch.tensor( + unfinished_req_index, + dtype=torch.int64, + device=batch.spec_info.topk_p.device, + ) + batch.spec_info.filter_batch( + unfinished_index_device, has_been_filtered=False + ) + + def forward_draft_extend_after_decode(self, batch: ScheduleBatch): + assert isinstance(batch.spec_info, EagleDraftInput) + # Backup fields that will be modified in-place + seq_lens_backup = batch.seq_lens.clone() + seq_lens_cpu_backup = batch.seq_lens_cpu.clone() + req_pool_indices_backup = batch.req_pool_indices + accept_length_backup = batch.spec_info.accept_length + return_logprob_backup = batch.return_logprob + + input_is_idle = batch.forward_mode.is_idle() + + if not input_is_idle and batch.spec_info.verified_id.numel() == 0: + batch = batch.copy() + batch.prepare_for_idle() + hidden_size = ( + self.model_config.hidden_size * 3 + if self.speculative_algorithm.is_eagle3() + else self.model_config.hidden_size + ) + batch.spec_info = EagleDraftInput.create_idle_input( + device=self.device, + hidden_size=hidden_size, + dtype=self.model_config.dtype, + topk=self.topk, + capture_hidden_mode=CaptureHiddenMode.LAST, + ) + + batch.spec_info.num_tokens_per_batch = self.speculative_num_steps + 1 + batch.spec_info.num_tokens_for_logprob_per_batch = 1 + batch.spec_info.prepare_extend_after_decode( + batch, + self.speculative_num_steps, + ) + batch.forward_mode = ( + ForwardMode.DRAFT_EXTEND + if not batch.forward_mode.is_idle() + else ForwardMode.IDLE + ) + + batch.return_hidden_states = False + model_worker_batch = batch.get_model_worker_batch() + assert model_worker_batch.capture_hidden_mode == CaptureHiddenMode.LAST + forward_batch = ForwardBatch.init_new( + model_worker_batch, self.mtp_model_runner(0) + ) + forward_batch.return_hidden_states_before_norm = True + if forward_batch.seq_lens_cpu is not None: + forward_batch.seq_lens_sum = forward_batch.seq_lens_cpu.sum().item() + else: + forward_batch.seq_lens_sum = batch.seq_lens.sum().item() + topk_p_list = [] + topk_index_list = [] + # Run + for step in range(self.speculative_num_steps): + can_cuda_graph = len( + self.cuda_graph_runner_for_draft_extend_list + ) and self.cuda_graph_runner_for_draft_extend_list[step].can_run( + forward_batch + ) + if can_cuda_graph: + logits_output = self.cuda_graph_runner_for_draft_extend_list[ + step + ].replay(forward_batch) + else: + forward_batch.can_run_dp_cuda_graph = False + if not forward_batch.forward_mode.is_idle(): + self.mtp_model_runner(step).attn_backend.init_forward_metadata( + forward_batch + ) + logits_output, _ = self.mtp_model_runner(step).forward( + forward_batch, skip_attn_backend_init=True + ) + + if self.enable_nan_detection: + detect_nan(logits_output) + probs = torch.softmax(logits_output.next_token_logits, dim=-1) + topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + topk_p_list.append(topk_p) + topk_index_list.append(topk_index) + pt = 0 + if forward_batch.extend_seq_lens is not None: + for i, extend_len in enumerate(forward_batch.extend_seq_lens): + input_ids = forward_batch.input_ids[pt : pt + extend_len] + forward_batch.input_ids[pt : pt + extend_len] = torch.cat( + (input_ids[1:], topk_index[i].reshape(1)) + ) + pt += extend_len + + forward_batch.spec_info.topk_p = torch.cat(topk_p_list, dim=1) + forward_batch.spec_info.topk_index = torch.cat(topk_index_list, dim=1) + + # Restore backup. + # This is because `seq_lens` can be modified in `prepare_extend_after_decode` + batch.forward_mode = ( + ForwardMode.DECODE if not input_is_idle else ForwardMode.IDLE + ) + batch.seq_lens = seq_lens_backup + batch.seq_lens_cpu = seq_lens_cpu_backup + batch.req_pool_indices = req_pool_indices_backup + batch.spec_info.accept_length = accept_length_backup + batch.return_logprob = return_logprob_backup diff --git a/python/sglang/srt/speculative/mtp_worker_v2.py b/python/sglang/srt/speculative/mtp_worker_v2.py new file mode 100644 index 000000000..981f82209 --- /dev/null +++ b/python/sglang/srt/speculative/mtp_worker_v2.py @@ -0,0 +1,750 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import contextlib +import logging +from typing import List, Optional, Tuple + +import torch + +from sglang.srt.environ import envs +from sglang.srt.layers.moe.utils import speculative_moe_backend_context +from sglang.srt.managers.schedule_batch import ModelWorkerBatch +from sglang.srt.managers.scheduler import GenerationBatchResult +from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch +from sglang.srt.server_args import ServerArgs +from sglang.srt.speculative.base_spec_worker import BaseDraftWorker, BaseSpecWorker +from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput +from sglang.srt.speculative.eagle_info_v2 import ( + assign_extend_cache_locs, + fill_accepted_out_cache_loc, + fill_new_verified_id, +) +from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient +from sglang.srt.speculative.mtp_draft_extend_cuda_graph_runner import ( + MTPMultiStepDraftExtendCudaGraphRunner, +) +from sglang.srt.speculative.mtp_utils import ( + assign_hidden_states_pool_triton, + rotate_input_ids_triton, +) +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.speculative.spec_utils import ( + detect_nan, + draft_tp_context, + select_top_k_tokens, +) +from sglang.srt.utils.common import empty_context, fast_topk, next_power_of_2 + +logger = logging.getLogger(__name__) + + +def _get_plan_stream( + device: str, +) -> Tuple[any, contextlib.AbstractContextManager]: + if envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get(): + plan_stream = torch.get_device_module(device).Stream() + plan_stream_ctx = torch.get_device_module(device).stream(plan_stream) + return plan_stream, plan_stream_ctx + else: + return None, contextlib.nullcontext() + + +class MTPDraftWorker(BaseDraftWorker): + def __init__( + self, + server_args: ServerArgs, + gpu_id: int, + tp_rank: int, + dp_rank: int, + moe_ep_rank: int, + nccl_port: int, + target_worker: TpModelWorker, + ): + # copy args + self.server_args = server_args + self.gpu_id = gpu_id + self.tp_rank = tp_rank + self.dp_rank = dp_rank + self.moe_ep_rank = moe_ep_rank + self.nccl_port = nccl_port + self.target_worker = target_worker + self.draft_extend_attn_backend_list = [] + self.model_config = target_worker.model_config + + # Args for easy access + self.device = server_args.device + self.topk = server_args.speculative_eagle_topk + self.speculative_num_steps = server_args.speculative_num_steps + self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + self.speculative_algorithm = SpeculativeAlgorithm.from_string( + server_args.speculative_algorithm + ) + + # Set constant + EagleDraftInput.ALLOC_LEN_PER_DECODE = max( + self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens + ) + + # Do not capture cuda graph in `TpModelWorker` init, + # will capture later with init_cuda_graphs() + backup_disable_cuda_graph = server_args.disable_cuda_graph + server_args.disable_cuda_graph = True + + # Share the allocator with a target worker. + # Draft and target worker own their own KV cache pools. + self.req_to_token_pool, self.token_to_kv_pool_allocator = ( + target_worker.get_memory_pool() + ) + with empty_context(), speculative_moe_backend_context(): + # Init draft worker + self.draft_worker = TpModelWorker( + server_args=server_args, + gpu_id=gpu_id, + tp_rank=tp_rank, + pp_rank=0, # FIXME + dp_rank=dp_rank, + moe_ep_rank=moe_ep_rank, + nccl_port=nccl_port, + is_draft_worker=True, + req_to_token_pool=self.req_to_token_pool, + token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, + is_mtp_worker=True, + ) + + # Alias for better readability + # self.draft_runner = self.draft_worker.model_runner + self.draft_runner_list = self.draft_worker.model_runner_list + + self.init_lm_head() + + # Used for KV Cache reversion + self.req_to_hidden_states_pool = torch.empty( + ( + self.req_to_token_pool.size, + self.speculative_num_steps - 1, + self.model_config.hidden_size, + ), + dtype=self.model_config.dtype, + device=self.device, + ) + + # Init attention backend and cuda graphs + for i in range(self.speculative_num_steps): + self.draft_runner_list[i].server_args.disable_cuda_graph = ( + backup_disable_cuda_graph + ) + self.draft_tp_context = ( + draft_tp_context if server_args.enable_dp_attention else empty_context + ) + with self.draft_tp_context( + self.draft_runner_list[0].tp_group + ), speculative_moe_backend_context(): + self.init_attention_backend() + self.init_cuda_graphs() + + self.tree_mask_mode = TreeMaskMode.FULL_MASK + + self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device) + + def mtp_model_runner(self, step: int): + return self.draft_runner_list[step] + + def init_lm_head(self): + embed, head = self.target_worker.model_runner.model.get_embed_and_head() + # Share the embedding and lm_head + for i in range(self.speculative_num_steps): + self.draft_runner_list[i].model.set_embed_and_head(embed, head) + + def init_attention_backend(self): + # Create attn backends + self.draft_extend_attn_backend_list = [] + for step in range(self.speculative_num_steps): + from sglang.srt.layers.attention.flashattention_backend import ( + FlashAttentionBackend, + ) + + self.draft_extend_attn_backend_list.append( + FlashAttentionBackend( + model_runner=self.draft_runner_list[step], + skip_prefill=False, + speculative_step_id=step, + ) + ) + self.draft_runner_list[step].attn_backend = ( + self.draft_extend_attn_backend_list[-1] + ) + + def init_cuda_graphs(self): + """Capture cuda graphs.""" + self.cuda_graph_runner = None + self.cuda_graph_runner_for_draft_extend = None + + if self.server_args.disable_cuda_graph: + return + + self.cuda_graph_runner_for_draft_extend = ( + MTPMultiStepDraftExtendCudaGraphRunner(self) + ) + + def reset_cuda_graph_buffers(self, forward_batch, batch_result): + if self.cuda_graph_runner_for_draft_extend: + self.cuda_graph_runner_for_draft_extend.reset_buffers( + forward_batch, batch_result + ) + + def draft(self, model_worker_batch: ModelWorkerBatch): + draft_input: EagleDraftInput = model_worker_batch.spec_info + forward_batch, can_cuda_graph = draft_input.prepare_for_v2_draft( + self.req_to_token_pool, + model_worker_batch, + self.cuda_graph_runner, + self.draft_runner_list[0], + self.topk, + self.speculative_num_steps, + ) + + # Run draft + parent_list, top_scores_index, draft_tokens = self.draft_forward(forward_batch) + + if model_worker_batch.forward_mode.is_idle(): + return EagleVerifyInput.create_idle_input( + self.topk, + self.speculative_num_steps, + self.speculative_num_draft_tokens, + ) + + # Build tree mask + # Directly write to cuda graph buffers for verify attn + tree_mask_buf, position_buf = ( + self.target_worker.model_runner.attn_backend.get_verify_buffers_to_fill_after_draft() + ) + ( + tree_mask, + position, + retrive_index, + retrive_next_token, + retrive_next_sibling, + draft_tokens, + ) = build_tree_kernel_efficient( + draft_input.verified_id, + parent_list, + top_scores_index, + draft_tokens, + model_worker_batch.seq_lens, + model_worker_batch.seq_lens_sum, + self.topk, + self.speculative_num_steps, + self.speculative_num_draft_tokens, + self.tree_mask_mode, + tree_mask_buf, + position_buf, + ) + + return EagleVerifyInput( + draft_token=draft_tokens, + custom_mask=tree_mask, + positions=position, + retrive_index=retrive_index, + retrive_next_token=retrive_next_token, + retrive_next_sibling=retrive_next_sibling, + retrive_cum_len=None, + spec_steps=self.speculative_num_steps, + topk=self.topk, + draft_token_num=self.speculative_num_draft_tokens, + capture_hidden_mode=None, + seq_lens_sum=None, + seq_lens_cpu=None, + ) + + def draft_forward(self, forward_batch: ForwardBatch): + # Parse args + spec_info: EagleDraftInput = forward_batch.spec_info + topk_p, topk_index, hidden_states = ( + spec_info.topk_p, + spec_info.topk_index, + spec_info.hidden_states, + ) + + # Return values + score_list: List[torch.Tensor] = [] + token_list: List[torch.Tensor] = [] + parents_list: List[torch.Tensor] = [] + + # Forward multiple steps + scores = None + _, hidden_states, scores, tree_info = select_top_k_tokens( + 0, topk_p, topk_index, hidden_states, scores, self.topk + ) + if self.speculative_num_steps == 1: + score_list.append(tree_info[0]) + token_list.append(tree_info[1]) + parents_list.append(tree_info[2]) + else: + for i in range(self.speculative_num_steps): + score_list.append(tree_info[0][:, :, i].unsqueeze(-1)) + token_index = tree_info[1][:, i].unsqueeze(-1) + token_list.append(token_index) + if i == 0: + parents_list.append(tree_info[2]) + else: + parents_list.append( + torch.full( + (tree_info[2].size(0), 1), + i, + dtype=torch.long, + device="cuda", + ) + ) + + # Organize the results + score_list = torch.cat(score_list, dim=1).flatten( + 1 + ) # b, n, topk; n= 1 + (num_steps-1) * self.topk + ss_token_list = torch.cat( + token_list, dim=1 + ) # b, (self.topk + (num_steps-1) * self.topk) + top_scores = torch.topk( + score_list, self.speculative_num_draft_tokens - 1, dim=-1 + ) + top_scores_index = top_scores.indices + top_scores_index = torch.sort(top_scores_index).values + draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1) + + if len(parents_list) > 1: + parent_list = torch.cat(parents_list[:-1], dim=1) + else: + batch_size = parents_list[0].shape[0] + parent_list = torch.empty(batch_size, 0, device=parents_list[0].device) + + return parent_list, top_scores_index, draft_tokens + + def draft_extend(self): + pass + + def _draft_extend_for_prefill( + self, + batch: ModelWorkerBatch, + target_hidden_states: torch.Tensor, + next_token_ids: torch.Tensor, + ): + """ + Run draft model extend to correctly fill the KV cache. + + Args: + batch: The batch to run. + target_hidden_states: Hidden states from the target model forward + next_token_ids: Next token ids generated from the target forward. + """ + # Construct spec_info + next_draft_input = EagleDraftInput( + hidden_states=target_hidden_states, + verified_id=next_token_ids, + new_seq_lens=batch.seq_lens, + # draft mode is same with decode mode, only 1 num token per batch + num_tokens_per_batch=1, + num_tokens_for_logprob_per_batch=1, + ) + + batch.spec_info = next_draft_input + + # Run forward + forward_batch = ForwardBatch.init_new(batch, self.draft_runner_list[0]) + forward_batch.return_hidden_states_before_norm = True + + # Construct input_ids + if not batch.forward_mode.is_idle(): + rotate_input_ids_triton( + forward_batch.input_ids, + forward_batch.extend_start_loc, + forward_batch.extend_seq_lens, + next_token_ids, + ) + + topk_p_list = [] + topk_index_list = [] + for step in range(self.speculative_num_steps): + logits_output, _ = self.draft_runner_list[step].forward(forward_batch) + probs = torch.softmax(logits_output.next_token_logits, dim=-1) + topk_p, topk_index = fast_topk(probs, self.topk, dim=-1) + topk_p_list.append(topk_p) + topk_index_list.append(topk_index) + if forward_batch.extend_seq_lens is not None: + rotate_input_ids_triton( + forward_batch.input_ids, + forward_batch.extend_start_loc, + forward_batch.extend_seq_lens, + topk_index, + ) + next_draft_input.topk_p = torch.cat(topk_p_list, dim=1) + next_draft_input.topk_index = torch.cat(topk_index_list, dim=1) + # next_draft_input.hidden_states = logits_output.hidden_states + + # Update req_to_hidden_states_pool for KV Cache reversion + if forward_batch.extend_seq_lens is not None: + assign_hidden_states_pool_triton( + target_hidden_states, + forward_batch.req_pool_indices, + self.req_to_hidden_states_pool, + self.speculative_num_steps - 1, + forward_batch.batch_size, + forward_batch.extend_seq_lens, + forward_batch.extend_start_loc, + ) + return next_draft_input + + def _draft_extend_for_decode( + self, batch: ModelWorkerBatch, batch_result: GenerationBatchResult + ): + # Batch 2: Draft extend + draft_input = EagleDraftInput( + hidden_states=batch_result.logits_output.hidden_states, + num_tokens_per_batch=self.speculative_num_steps + 1, + num_tokens_for_logprob_per_batch=1, + ) + + # Prepare for draft extend in a separate stream + # Notice that here we use batch_result.next_token_ids as the input ids + with self.plan_stream_ctx: + forward_batch = draft_input.prepare_for_extend_to_fill_draft_kvcache( + batch, + batch_result.next_token_ids, + self.speculative_num_draft_tokens, + self.draft_runner_list[0], + self.cuda_graph_runner_for_draft_extend, + ) + forward_batch.return_hidden_states_before_norm = True + + if self.plan_stream: + torch.get_device_module(self.device).current_stream().wait_stream( + self.plan_stream + ) + # Run draft extend batch in the main compute stream + can_cuda_graph = ( + self.cuda_graph_runner_for_draft_extend + and self.cuda_graph_runner_for_draft_extend.can_run(forward_batch) + ) + ret_topk_p_list = [] + ret_topk_index_list = [] + next_token_ids_backup = batch_result.next_token_ids.clone() + + if can_cuda_graph: + self.reset_cuda_graph_buffers(forward_batch, batch_result) + else: + logger.warning_once( + f"can't use cuda graph for draft extend! may have correctness issue!" + ) + select_index = ( + torch.arange(len(batch.seq_lens), device=self.device) + * self.speculative_num_draft_tokens + + batch_result.accept_lens + - 1 + ) + + for step in range(self.speculative_num_steps): + # log_info_on_rank0(logger, f"step: {step}, forward_batch.input_ids: {forward_batch.input_ids}") + if can_cuda_graph: + draft_logits_output = ( + self.cuda_graph_runner_for_draft_extend.get_runner(step).replay( + forward_batch, init_state=(step == 0) + ) + ) + ret_topk_p, ret_topk_index = ( + draft_logits_output.topk_p, + draft_logits_output.topk_index, + ) + else: + draft_logits_output, _ = self.draft_runner_list[step].forward( + forward_batch, skip_attn_backend_init=True + ) + probs = torch.softmax( + draft_logits_output.next_token_logits[select_index], dim=-1 + ) + ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1) + if forward_batch.extend_seq_lens is not None: + rotate_input_ids_triton( + forward_batch.input_ids, + forward_batch.extend_start_loc, + forward_batch.extend_seq_lens, + ret_topk_index, + select_index, + ) + ret_topk_p_list.append(ret_topk_p) + ret_topk_index_list.append(ret_topk_index) + + # Update req_to_hidden_states_pool for KV Cache reversion + if ( + self.cuda_graph_runner_for_draft_extend is not None + and forward_batch.extend_seq_lens is not None + ): + last_cuda_graph_runner = ( + self.cuda_graph_runner_for_draft_extend.get_last_runner() + ) + assign_hidden_states_pool_triton( + last_cuda_graph_runner.hidden_states, + last_cuda_graph_runner.req_pool_indices, + self.req_to_hidden_states_pool, + self.speculative_num_steps - 1, + forward_batch.batch_size, + last_cuda_graph_runner.extend_seq_lens, + last_cuda_graph_runner.extend_start_loc, + ) + + # Reorganize the spec info for the next batch + # draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[ + # select_index + # ] + # draft_logits_output.hidden_states = draft_logits_output.hidden_states[ + # select_index + # ] + batch_result.next_token_ids = next_token_ids_backup + # Construct the return values + next_draft_input = batch_result.next_draft_input + ( + next_draft_input.topk_p, + next_draft_input.topk_index, + next_draft_input.hidden_states, + ) = ( + torch.cat(ret_topk_p_list, dim=1).clone(), + torch.cat(ret_topk_index_list, dim=1).clone(), + None, + ) + + +class MTPWorkerV2(BaseSpecWorker): + def __init__( + self, + server_args: ServerArgs, + gpu_id: int, + tp_rank: int, + dp_rank: Optional[int], + moe_ep_rank: int, + nccl_port: int, + target_worker: TpModelWorker, + ): + # Parse arguments + self.server_args = server_args + self.topk = server_args.speculative_eagle_topk + self.speculative_num_steps = server_args.speculative_num_steps + self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + self.enable_nan_detection = server_args.enable_nan_detection + self.gpu_id = gpu_id + self.device = server_args.device + self._target_worker = target_worker + self.page_size = server_args.page_size + self.speculative_algorithm = SpeculativeAlgorithm.from_string( + server_args.speculative_algorithm + ) + + self.req_to_token_pool, self.token_to_kv_pool_allocator = ( + target_worker.get_memory_pool() + ) + + # Override the context length of the draft model to be the same as the target model. + server_args.context_length = target_worker.model_runner.model_config.context_len + + self._draft_worker = MTPDraftWorker( + server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker + ) + + # Some dummy tensors + self.num_new_pages_per_topk = torch.empty( + (), dtype=torch.int64, device=self.device + ) + self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device) + + self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device) + + @property + def target_worker(self): + return self._target_worker + + @property + def draft_worker(self): + return self._draft_worker + + def clear_cache_pool(self): + # allocator and kv cache pool are shared with target worker, which are cleared in scheduler + pass + + def forward_batch_generation(self, model_worker_batch: ModelWorkerBatch): + if ( + model_worker_batch.forward_mode.is_extend() + or model_worker_batch.is_extend_in_batch + ): + # Target prefill + model_worker_batch.capture_hidden_mode = CaptureHiddenMode.FULL + batch_output = self.target_worker.forward_batch_generation( + model_worker_batch + ) + + # Draft prefill + model_worker_batch.capture_hidden_mode = CaptureHiddenMode.LAST + batch_output.next_draft_input = self.draft_worker._draft_extend_for_prefill( + model_worker_batch, + batch_output.logits_output.hidden_states, + batch_output.next_token_ids, + ) + return batch_output + else: + if model_worker_batch.spec_info is None: + model_worker_batch.spec_info = EagleDraftInput.create_idle_input( + device=self.device, + hidden_size=self.target_worker.model_config.hidden_size, + dtype=self.target_worker.model_config.dtype, + topk=self.topk * self.speculative_num_steps, + capture_hidden_mode=CaptureHiddenMode.LAST, + ) + draft_input: EagleDraftInput = model_worker_batch.spec_info + verify_input: EagleVerifyInput = self.draft_worker.draft(model_worker_batch) + assert verify_input.is_verify_input() + model_worker_batch.spec_info = verify_input + batch_output = self.verify(model_worker_batch) + self.draft_worker._draft_extend_for_decode(model_worker_batch, batch_output) + return batch_output + + def verify( + self, + batch: ModelWorkerBatch, + ): + # Since batch.seq_lens is allocated in another stream, we need + # record_stream() to prevent pytorch gc and reuse the gpu memory + # while forward_stream is still running. + batch.seq_lens.record_stream( + torch.get_device_module(self.device).current_stream() + ) + + # Parse args + verify_input: EagleVerifyInput = batch.spec_info + bs = len(batch.seq_lens) + + # Batch 1: Target verify + # Prepare for target verify in a separate stream + with self.plan_stream_ctx: + verify_forward_batch, can_run_cuda_graph = ( + verify_input.prepare_for_v2_verify( + self.req_to_token_pool, + batch, + self.target_worker, + ) + ) + + # Correct some buffers due to the overlap plan + if self.plan_stream: + torch.get_device_module(self.device).current_stream().wait_stream( + self.plan_stream + ) + + # Some values such as custom_mask and position depend on the output of draft, + # so the previous plan step used the wrong values. Here, we need to run the related + # computation again to update them to the correct values. + self.target_worker.model_runner.attn_backend.update_verify_buffers_to_fill_after_draft( + verify_input, + ( + self.target_worker.model_runner.graph_runner.bs + if can_run_cuda_graph + else None + ), + ) + # Run target verify batch in the main compute stream + forward_batch_output = self.target_worker.forward_batch_generation( + model_worker_batch=None, + forward_batch=verify_forward_batch, + is_verify=True, + skip_attn_backend_init=True, + ) + logits_output = forward_batch_output.logits_output + + # Sample + if self.enable_nan_detection: + detect_nan(logits_output) + ( + predict, + accept_length, + accept_index, + ) = verify_input.sample(batch, logits_output) + new_seq_lens = batch.seq_lens + accept_length + verify_done = torch.get_device_module(self.device).Event() + verify_done.record() + + if not batch.forward_mode.is_idle(): + all_verified_id = predict[accept_index] + verified_id = torch.empty_like(accept_length, dtype=torch.int32) + fill_new_verified_id[(bs,)]( + all_verified_id, + accept_length, + verified_id, + self.speculative_num_draft_tokens, + ) + else: + verified_id = torch.empty((0,), device=self.device, dtype=torch.int32) + + # Construct the next draft input + next_draft_input = EagleDraftInput( + verified_id=verified_id, + new_seq_lens=new_seq_lens, + verify_done=verify_done, + ) + return GenerationBatchResult( + logits_output=logits_output, + next_token_ids=predict, + can_run_cuda_graph=can_run_cuda_graph, + next_draft_input=next_draft_input, + accept_lens=accept_length, + ) + + def move_accepted_tokens_to_target_kvcache( + self, + batch: ModelWorkerBatch, + accept_index: torch.Tensor, + accept_length: torch.Tensor, + ): + """ + Move accepted tokens to the target KV cache. + + Args: + batch: The batch to run. + accept_index: The index of the accepted tokens. + accept_length: The length of the accepted tokens. + """ + bs = len(batch.seq_lens) + size = bs * self.speculative_num_draft_tokens + + tgt_cache_loc = torch.zeros( + size, + dtype=torch.int64, + device=self.device, + ) + accepted_out_cache_loc = torch.zeros( + size, dtype=torch.int64, device=self.device + ) + assign_extend_cache_locs[(bs,)]( + batch.req_pool_indices, + self.req_to_token_pool.req_to_token, + batch.seq_lens, + batch.seq_lens + accept_length, + tgt_cache_loc, + self.req_to_token_pool.req_to_token.shape[1], + next_power_of_2(bs), + ) + fill_accepted_out_cache_loc[(size,)]( + accept_index, + batch.out_cache_loc, + accepted_out_cache_loc, + next_power_of_2(size), + ) + self.token_to_kv_pool_allocator.get_kvcache().move_kv_cache( + tgt_cache_loc, accepted_out_cache_loc + ) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 2fb41b407..cf2569b19 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -4,7 +4,7 @@ import logging import os import time from contextlib import contextmanager -from typing import TYPE_CHECKING, List +from typing import TYPE_CHECKING, List, Optional import torch import triton @@ -19,6 +19,8 @@ from sglang.srt.distributed.parallel_state import ( from sglang.srt.environ import envs from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import Req +from sglang.srt.mem_cache.common import get_last_loc +from sglang.srt.server_args import ServerArgs, get_global_server_args from sglang.srt.utils import is_cuda, is_hip, is_npu, next_power_of_2 _is_cuda = is_cuda() @@ -48,6 +50,14 @@ TREE_TRAVERSE_TIME_THRESHOLD = 1 # TODO: set this properly TREE_SPEC_KERNEL_AVAILABLE = _is_cuda # This kernel is only available for CUDA now +def spec_need_hidden_states(server_args: Optional[ServerArgs] = None) -> bool: + if server_args is None: + server_args = get_global_server_args() + + # TODO(lsyin): also skip when 1) step = 1 or 2) standalone draft model + return not server_args.enable_mtp + + @triton.jit def create_extend_after_decode_spec_info( verified_id, @@ -465,13 +475,14 @@ def select_top_k_tokens( if i == 0: # The first step after extend input_ids = topk_index.flatten() - hidden_states = hidden_states.repeat_interleave(topk, dim=0) + if hidden_states is not None: + hidden_states = hidden_states.repeat_interleave(topk, dim=0) scores = topk_p # shape: (b, topk) tree_info = ( topk_p.unsqueeze(1), # shape: (b, 1, topk) topk_index, # shape: (b, topk) - torch.arange(-1, topk, dtype=torch.long, device=hidden_states.device) + torch.arange(-1, topk, dtype=torch.long, device=input_ids.device) .unsqueeze(0) .repeat(topk_p.shape[0], 1), # shape: (b, topk + 1) ) @@ -695,3 +706,39 @@ def detect_nan(logits_output: LogitsProcessorOutput): if torch.any(torch.isnan(logits)): logger.error("Detected errors during sampling! NaN in the logits.") raise ValueError("Detected errors during sampling! NaN in the logits.") + + +# Disable torch.compile for this function because it will be +# even slower. +# @torch.compile(dynamic=True) +def get_last_loc_large_page_size_large_top_k( + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + speculative_num_steps: int, + topk: int, + page_size: int, +): + prefix_lens = seq_lens + last_page_lens = prefix_lens % page_size + num_new_pages_per_topk = ( + last_page_lens + speculative_num_steps + page_size - 1 + ) // page_size + seq_lens = prefix_lens // page_size * page_size + num_new_pages_per_topk * ( + page_size * topk + ) + extend_lens = seq_lens - prefix_lens + last_loc = get_last_loc( + req_to_token, + req_pool_indices, + prefix_lens, + ) + + return ( + prefix_lens, + seq_lens, + last_loc, + num_new_pages_per_topk, + extend_lens, + last_page_lens, + ) diff --git a/test/registered/function_call/test_function_call_parser.py b/test/registered/function_call/test_function_call_parser.py index 9c99e9f7e..e695259f3 100644 --- a/test/registered/function_call/test_function_call_parser.py +++ b/test/registered/function_call/test_function_call_parser.py @@ -10,6 +10,7 @@ from sglang.srt.function_call.glm4_moe_detector import Glm4MoeDetector from sglang.srt.function_call.json_array_parser import JsonArrayParser from sglang.srt.function_call.kimik2_detector import KimiK2Detector from sglang.srt.function_call.llama32_detector import Llama32Detector +from sglang.srt.function_call.mimo_detector import MiMoDetector from sglang.srt.function_call.mistral_detector import MistralDetector from sglang.srt.function_call.pythonic_detector import PythonicDetector from sglang.srt.function_call.qwen3_coder_detector import Qwen3CoderDetector @@ -2246,6 +2247,446 @@ class TestGlm4MoeDetector(unittest.TestCase): check_single_todos(result, expected_output) +class TestMiMoDetector(unittest.TestCase): + def setUp(self): + # Create sample tools for testing + self.tools = [ + Tool( + type="function", + function=Function( + name="get_current_weather", + description="Get the current weather", + parameters={ + "properties": { + "city": {"type": "string", "description": "The city name"}, + "state": { + "type": "string", + "description": "The state code", + }, + "unit": { + "type": "string", + "enum": ["fahrenheit", "celsius"], + }, + }, + "required": ["city", "state"], + }, + ), + ), + Tool( + type="function", + function=Function( + name="calculate_area", + description="Calculate area of a shape", + parameters={ + "properties": { + "shape": {"type": "string"}, + "dimensions": {"type": "object"}, + "precision": {"type": "integer"}, + } + }, + ), + ), + ] + self.detector = MiMoDetector() + + def test_has_tool_call(self): + """Test detection of tool call markers.""" + self.assertTrue(self.detector.has_tool_call("test")) + self.assertFalse(self.detector.has_tool_call("No tool call here")) + + def test_detect_and_parse_no_tools(self): + """Test parsing text without tool calls.""" + model_output = "This is a test response without any tool calls" + result = self.detector.detect_and_parse(model_output, tools=[]) + self.assertEqual(result.normal_text, model_output) + self.assertEqual(result.calls, []) + + def test_detect_and_parse_single_tool(self): + """Test parsing a single tool call.""" + model_output = """ + +Dallas +TX +fahrenheit + +""" + + result = self.detector.detect_and_parse(model_output, tools=self.tools) + + self.assertEqual(result.normal_text, "") + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_current_weather") + + params = json.loads(result.calls[0].parameters) + self.assertEqual(params["city"], "Dallas") + self.assertEqual(params["state"], "TX") + self.assertEqual(params["unit"], "fahrenheit") + + def test_detect_and_parse_with_content(self): + """Test parsing tool call with surrounding text.""" + model_output = """Sure! Let me check the weather for you. + +Dallas +TX +fahrenheit + +""" + + result = self.detector.detect_and_parse(model_output, tools=self.tools) + + self.assertEqual(result.normal_text, "Sure! Let me check the weather for you.") + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_current_weather") + + def test_detect_and_parse_multiline_param(self): + """Test parsing tool call with multiline parameter values.""" + model_output = """ + +rectangle +{"width": 10, "height": 20} +2 + +""" + + result = self.detector.detect_and_parse(model_output, tools=self.tools) + + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "calculate_area") + + params = json.loads(result.calls[0].parameters) + self.assertEqual(params["shape"], "rectangle") + self.assertEqual(params["dimensions"], {"width": 10, "height": 20}) + self.assertEqual(params["precision"], 2) + + def test_detect_and_parse_parallel_tools(self): + """Test parsing multiple tool calls.""" + model_output = """ + +Dallas +TX +fahrenheit + + + + +Orlando +FL +fahrenheit + +""" + + result = self.detector.detect_and_parse(model_output, tools=self.tools) + + self.assertEqual(result.normal_text, "") + self.assertEqual(len(result.calls), 2) + + # First call + self.assertEqual(result.calls[0].name, "get_current_weather") + params1 = json.loads(result.calls[0].parameters) + self.assertEqual(params1["city"], "Dallas") + self.assertEqual(params1["state"], "TX") + + # Second call + self.assertEqual(result.calls[1].name, "get_current_weather") + params2 = json.loads(result.calls[1].parameters) + self.assertEqual(params2["city"], "Orlando") + self.assertEqual(params2["state"], "FL") + + def test_parse_streaming_simple(self): + """Test basic streaming parsing.""" + chunks = [ + "Sure! ", + "Let me check ", + "the weather.", + "", + "\n", + "\nDallas", + "\nTX", + "\n", + "\n", + ] + + accumulated_text = "" + accumulated_calls = [] + tool_calls_by_index = {} + + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, tools=self.tools) + accumulated_text += result.normal_text + + # Track calls by tool_index to handle streaming properly + for call in result.calls: + if call.tool_index is not None: + if call.tool_index not in tool_calls_by_index: + tool_calls_by_index[call.tool_index] = { + "name": "", + "parameters": "", + } + + if call.name: + tool_calls_by_index[call.tool_index]["name"] = call.name + if call.parameters: + tool_calls_by_index[call.tool_index][ + "parameters" + ] += call.parameters + + self.assertEqual(accumulated_text, "Sure! Let me check the weather.") + self.assertEqual(len(tool_calls_by_index), 1) + + # Get the complete tool call + tool_call = tool_calls_by_index[0] + self.assertEqual(tool_call["name"], "get_current_weather") + + # Parse the accumulated parameters + params = json.loads(tool_call["parameters"]) + self.assertEqual(params["city"], "Dallas") + self.assertEqual(params["state"], "TX") + + def test_parse_streaming_incomplete(self): + """Test streaming with incomplete tool call.""" + # Send incomplete tool call + chunks = [ + "", + "\n", + "\nDallas", + "\n", + # Missing , , + ] + + tool_calls_by_index = {} + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, tools=self.tools) + + # Track calls by tool_index to handle streaming properly + for call in result.calls: + if call.tool_index is not None: + if call.tool_index not in tool_calls_by_index: + tool_calls_by_index[call.tool_index] = { + "name": "", + "parameters": "", + } + + if call.name: + tool_calls_by_index[call.tool_index]["name"] = call.name + if call.parameters: + tool_calls_by_index[call.tool_index][ + "parameters" + ] += call.parameters + + # Should have no complete tool calls yet (buffered) + self.assertEqual(len(tool_calls_by_index), 0) + + # Now complete it + result = self.detector.parse_streaming_increment( + "TX\n\n", tools=self.tools + ) + + # Update the accumulated parameters + for call in result.calls: + if call.tool_index is not None: + if call.tool_index not in tool_calls_by_index: + tool_calls_by_index[call.tool_index] = { + "name": "", + "parameters": "", + } + if call.name: + tool_calls_by_index[call.tool_index]["name"] = call.name + if call.parameters: + tool_calls_by_index[call.tool_index][ + "parameters" + ] += call.parameters + + # Now should have complete tool call + self.assertEqual(len(tool_calls_by_index), 1) + final_params = json.loads(tool_calls_by_index[0]["parameters"]) + self.assertEqual(final_params["city"], "Dallas") + self.assertEqual(final_params["state"], "TX") + + def test_edge_case_no_parameters(self): + """Test tool call without parameters.""" + model_output = """ + + +""" + + result = self.detector.detect_and_parse(model_output, tools=self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_current_weather") + self.assertEqual(json.loads(result.calls[0].parameters), {}) + + def test_edge_case_special_chars_in_value(self): + """Test parameter with special characters in value.""" + model_output = """ + +Dallas->TX + +""" + + result = self.detector.detect_and_parse(model_output, tools=self.tools) + self.assertEqual(len(result.calls), 1) + + params = json.loads(result.calls[0].parameters) + self.assertEqual(params["city"], "Dallas->TX") + + def test_extract_tool_calls_type_conversion(self): + """Test parameter type conversion based on tool schema.""" + test_tool = Tool( + type="function", + function=Function( + name="test_types", + parameters={ + "type": "object", + "properties": { + "int_param": {"type": "integer"}, + "float_param": {"type": "float"}, + "bool_param": {"type": "boolean"}, + "str_param": {"type": "string"}, + "obj_param": {"type": "object"}, + }, + }, + ), + ) + + model_output = """ + +42 +3.14 +true +hello world +{"key": "value"} + +""" + + result = self.detector.detect_and_parse(model_output, tools=[test_tool]) + + self.assertEqual(len(result.calls), 1) + params = json.loads(result.calls[0].parameters) + self.assertEqual(params["int_param"], 42) + self.assertEqual(params["float_param"], 3.14) + self.assertEqual(params["bool_param"], True) + self.assertEqual(params["str_param"], "hello world") + self.assertEqual(params["obj_param"], {"key": "value"}) + + def test_parse_streaming_incremental(self): + """Test that streaming is truly incremental with very small chunks.""" + # Simulate more realistic token-based chunks where is a single token + chunks = [ + "I'll check the weather.", + "", + "\n\n", + "", + "Dallas", + "\n", + "", + "TX", + "\n", + "\n", + "", + ] + + accumulated_text = "" + tool_calls = [] + chunks_count = 0 + + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, self.tools) + accumulated_text += result.normal_text + chunks_count += 1 + for tool_call_chunk in result.calls: + if ( + hasattr(tool_call_chunk, "tool_index") + and tool_call_chunk.tool_index is not None + ): + while len(tool_calls) <= tool_call_chunk.tool_index: + tool_calls.append({"name": "", "parameters": ""}) + tc = tool_calls[tool_call_chunk.tool_index] + if tool_call_chunk.name: + tc["name"] = tool_call_chunk.name + if tool_call_chunk.parameters: + tc["parameters"] += tool_call_chunk.parameters + + self.assertGreater(chunks_count, 3) + + # Verify the accumulated results + self.assertIn("I'll check the weather.", accumulated_text) + self.assertEqual(len(tool_calls), 1) + self.assertEqual(tool_calls[0]["name"], "get_current_weather") + + params = json.loads(tool_calls[0]["parameters"]) + self.assertEqual(params, {"city": "Dallas", "state": "TX"}) + + def test_parse_streaming_multiple_tools(self): + """Test streaming with multiple tool calls.""" + model_output = """ + +Dallas +TX + + +Some text in between. + + +circle +{"radius": 5} + +""" + + # Simulate streaming by chunks + chunk_size = 20 + chunks = [ + model_output[i : i + chunk_size] + for i in range(0, len(model_output), chunk_size) + ] + + accumulated_text = "" + tool_calls = [] + chunks_count = 0 + + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, self.tools) + accumulated_text += result.normal_text + chunks_count += 1 + for tool_call_chunk in result.calls: + if ( + hasattr(tool_call_chunk, "tool_index") + and tool_call_chunk.tool_index is not None + ): + while len(tool_calls) <= tool_call_chunk.tool_index: + tool_calls.append({"name": "", "parameters": ""}) + tc = tool_calls[tool_call_chunk.tool_index] + if tool_call_chunk.name: + tc["name"] = tool_call_chunk.name + if tool_call_chunk.parameters: + tc["parameters"] += tool_call_chunk.parameters + + self.assertIn("Some text in between.", accumulated_text) + self.assertEqual(len(tool_calls), 2) + self.assertEqual(tool_calls[0]["name"], "get_current_weather") + self.assertEqual(tool_calls[1]["name"], "calculate_area") + + # Verify parameters + params1 = json.loads(tool_calls[0]["parameters"]) + self.assertEqual(params1, {"city": "Dallas", "state": "TX"}) + + params2 = json.loads(tool_calls[1]["parameters"]) + self.assertEqual(params2, {"shape": "circle", "dimensions": {"radius": 5}}) + + def test_html_entity_decoding(self): + """Test that HTML entities in parameter values are decoded.""" + model_output = """ + +Dallas & Fort Worth +TX + +""" + + result = self.detector.detect_and_parse(model_output, tools=self.tools) + + self.assertEqual(len(result.calls), 1) + params = json.loads(result.calls[0].parameters) + self.assertEqual(params["city"], "Dallas & Fort Worth") + + class TestJsonArrayParser(unittest.TestCase): def setUp(self): # Create sample tools for testing