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

Co-authored-by: 谢学扬 <xiexueyang@xiaomi.com>
Co-authored-by: tz <tangzhen3@xiaomi.com>
Co-authored-by: 李家乐 <lijiale10@xiaomi.com>
Co-authored-by: 张晨 <zhangchen50@xiaomi.com>
Co-authored-by: Shaohui Liu <liushaohui3@xiaomi.com>
Co-authored-by: 王晨 <wangchen77@xiaomi.com>
Co-authored-by: jiangzihan <jiangzihan@xiaomi.com>
Co-authored-by: xiexueyang <xyxie_wangyi@163.com>
Co-authored-by: Linghao Zhang <zhanglinghao@xiaomi.com>
Co-authored-by: ispobock <ispobaoke@gmail.com>
Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
Co-authored-by: JoyFuture <35593546+JoyFuture@users.noreply.github.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
Co-authored-by: root <root@bj9-ml-g8h20e-k8s-slave106-20251106.alicn.idc.xiaomi.com>

python/sglang/srt/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,
             )
         )
 

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

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

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

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

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

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

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,

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,

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:
+        <tool_call>
+        <function=execute_bash>
+        <parameter=command>pwd && ls</parameter>
+        </function>
+        </tool_call>
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.bot_token = "<tool_call>"
+        self.eot_token = "</tool_call>"
+        self.tool_call_regex = re.compile(r"<tool_call>(.*?)</tool_call>", re.DOTALL)
+        self.func_regex = re.compile(r"<function=([^>]+)>(.*?)</function>", re.DOTALL)
+        self.param_regex = re.compile(
+            r"<parameter=([^>]+)>(.*?)</parameter>", 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 <tool_call>...</tool_call>.
+
+        Structure:
+            tool_call_body contains: <function=name>...params...</function>
+        """
+        # Match complete <function=name>body</function> 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

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]

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,
                     ),
                 }

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)

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

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

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

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
 

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 (

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(

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

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
             ]

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

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)

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

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

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)
 

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

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

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

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]

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

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

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)

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, :]
+        )

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

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

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,
+    )

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("<tool_call>test</tool_call>"))
+        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 = """<tool_call>
+<function=get_current_weather>
+<parameter=city>Dallas</parameter>
+<parameter=state>TX</parameter>
+<parameter=unit>fahrenheit</parameter>
+</function>
+</tool_call>"""
+
+        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.<tool_call>
+<function=get_current_weather>
+<parameter=city>Dallas</parameter>
+<parameter=state>TX</parameter>
+<parameter=unit>fahrenheit</parameter>
+</function>
+</tool_call>"""
+
+        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 = """<tool_call>
+<function=calculate_area>
+<parameter=shape>rectangle</parameter>
+<parameter=dimensions>{"width": 10, "height": 20}</parameter>
+<parameter=precision>2</parameter>
+</function>
+</tool_call>"""
+
+        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 = """<tool_call>
+<function=get_current_weather>
+<parameter=city>Dallas</parameter>
+<parameter=state>TX</parameter>
+<parameter=unit>fahrenheit</parameter>
+</function>
+</tool_call>
+<tool_call>
+<function=get_current_weather>
+<parameter=city>Orlando</parameter>
+<parameter=state>FL</parameter>
+<parameter=unit>fahrenheit</parameter>
+</function>
+</tool_call>"""
+
+        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.",
+            "<tool_call>",
+            "\n<function=get_current_weather>",
+            "\n<parameter=city>Dallas</parameter>",
+            "\n<parameter=state>TX</parameter>",
+            "\n</function>",
+            "\n</tool_call>",
+        ]
+
+        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 = [
+            "<tool_call>",
+            "\n<function=get_current_weather>",
+            "\n<parameter=city>Dallas</parameter>",
+            "\n<parameter=state>",
+            # Missing </parameter>, </function>, </tool_call>
+        ]
+
+        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</parameter>\n</function>\n</tool_call>", 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 = """<tool_call>
+<function=get_current_weather>
+</function>
+</tool_call>"""
+
+        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 = """<tool_call>
+<function=get_current_weather>
+<parameter=city>Dallas->TX</parameter>
+</function>
+</tool_call>"""
+
+        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 = """<tool_call>
+<function=test_types>
+<parameter=int_param>42</parameter>
+<parameter=float_param>3.14</parameter>
+<parameter=bool_param>true</parameter>
+<parameter=str_param>hello world</parameter>
+<parameter=obj_param>{"key": "value"}</parameter>
+</function>
+</tool_call>"""
+
+        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 <tool_call> is a single token
+        chunks = [
+            "I'll check the weather.",
+            "<tool_call>",
+            "\n<function=get_current_weather>\n",
+            "<parameter=city>",
+            "Dallas",
+            "</parameter>\n",
+            "<parameter=state>",
+            "TX",
+            "</parameter>\n",
+            "</function>\n",
+            "</tool_call>",
+        ]
+
+        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 = """<tool_call>
+<function=get_current_weather>
+<parameter=city>Dallas</parameter>
+<parameter=state>TX</parameter>
+</function>
+</tool_call>
+Some text in between.
+<tool_call>
+<function=calculate_area>
+<parameter=shape>circle</parameter>
+<parameter=dimensions>{"radius": 5}</parameter>
+</function>
+</tool_call>"""
+
+        # 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 = """<tool_call>
+<function=get_current_weather>
+<parameter=city>Dallas &amp; Fort Worth</parameter>
+<parameter=state>TX</parameter>
+</function>
+</tool_call>"""
+
+        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