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