1 Commits

Author SHA1 Message Date
82330a31b1 add await to redis write 2026-04-17 07:19:15 +00:00
9 changed files with 1572 additions and 933 deletions

View File

@@ -1,7 +1,8 @@
#FROM vllm/vllm-openai:v0.19.0-cu130
#vllm says version 0.20.2rc1.dev9+g01d4d1ad3
FROM vllm/vllm-openai:nightly
#FROM vllm/vllm-openai:glm51-cu130
FROM vllm/vllm-openai:cu130-nightly-x86_64
# Fix the broken ass nightly build that forgot to include pandas
RUN pip install --no-cache-dir pandas
# Install LMCache for KV cache offloading / sharing across nodes
# Build with system CUDA 13.0 for Blackwell (B200)
@@ -14,21 +15,21 @@ RUN apt-get update && apt-get install -y git \
libnvjitlink-dev-13-0 && \
git clone https://github.com/biondizzle/LMCache.git /tmp/lmcache && \
cd /tmp/lmcache && \
git checkout dream-build && \
git checkout feat/redis-ttl && \
CUDA_HOME=/usr/local/cuda \
TORCH_CUDA_ARCH_LIST="10.0" \
pip install --no-cache-dir --no-build-isolation . && \
rm -rf /tmp/lmcache && export CACHE_BUSTER=3
rm -rf /tmp/lmcache && export CACHE_BUSTER=1
# Copy over nemotron reasonong parser
COPY ./super_v3_reasoning_parser.py /opt/super_v3_reasoning_parser.py
# Make sure we have patch to make MTP work on GLM
COPY indexer.py /usr/local/lib/python3.12/dist-packages/vllm/v1/attention/backends/mla/indexer.py
# Copy over deepseek tool call parser with MTP fixes
COPY deepseekv32_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/deepseekv32_tool_parser.py
# Make sure we have the latest up to date chat template
COPY glm_5.1_chat_template.jinja /opt/chat_template.jinja
# Copy over minimax tool call parser with kwargs fixes
COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/minimax_tool_parser.py
# GLM 5.1 LMCache config
COPY lmcache-config-glm-51.yaml /opt/lmcache-config-glm-51.yaml
# DEEPSEEK v4 LMCache config
COPY lmcache-config-dsv4.yaml /opt/lmcache-config-dsv4.yaml
# Copy over minimax parsers with kwargs fixes
COPY minimax_tool_parser.py /usr/local/lib/python3.12/dist-packages/vllm/tool_parsers/minimax_tool_parser.py
COPY minimax_m2_parser.py /usr/local/lib/python3.12/dist-packages/vllm/parser/minimax_m2_parser.py

616
deepseekv32_tool_parser.py Normal file
View File

@@ -0,0 +1,616 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
DeepSeek-V3.2 Tool Call Parser — re-parse-and-diff version.
Adapted from the GLM-4 streaming fix to make the streaming path robust
against multi-token deltas produced by MTP speculative decoding.
Instead of maintaining incremental state that advances one token at a
time, the streaming path re-parses the *entire* current_text on every
call, finds all <DSMLinvoke> regions (complete and in-progress),
builds a JSON arguments string for each, and diffs against what was
previously sent. This makes the parser agnostic to how many tokens
arrive per step.
Key changes vs. the upstream buffer-until-complete parser:
1. _extract_content() handles partial tag overlaps so content text
is never swallowed or duplicated when a tag boundary lands inside
a multi-token chunk.
2. _extract_invoke_regions() finds both complete and incomplete
invoke blocks, enabling streaming of partial arguments.
3. _build_args_json_so_far() constructs the JSON arguments string
incrementally from complete + partial <DSMLparameter> tags.
4. _compute_args_diff() emits only the newly-added characters.
Drop-in replacement: same class name, same interface.
"""
import json
import uuid
from collections.abc import Sequence
from typing import Any
import regex as re
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
FunctionCall,
ToolCall,
)
from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
from vllm.logger import init_logger
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import (
Tool,
ToolParser,
)
logger = init_logger(__name__)
def partial_tag_overlap(text: str, tag: str) -> int:
"""Length of the longest prefix of *tag* that matches a suffix of *text*.
E.g. text ending in ``"<tool_"`` returns 6 when tag is ``"<tool_call>"``.
Returns 0 when there is no overlap.
"""
max_check = min(len(tag) - 1, len(text))
for k in range(max_check, 0, -1):
if text.endswith(tag[:k]):
return k
return 0
class DeepSeekV32ToolParser(ToolParser):
"""
Re-parse-and-diff tool parser for DeepSeek-V3.2 DSML format.
On every streaming call the parser re-parses ``current_text`` to
find ``<DSMLinvoke>`` regions, builds the JSON arguments string
for each tool call, and diffs against what was previously sent to
emit only new content. This is robust against multi-token deltas
from MTP / EAGLE speculative decoding.
Example tool call format::
<DSMLfunction_calls>
<DSMLinvoke name="get_weather">
<DSMLparameter name="location" string="true">杭州</DSMLparameter>
<DSMLparameter name="date" string="true">2024-01-16</DSMLparameter>
</DSMLinvoke>
</DSMLfunction_calls>
"""
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None):
super().__init__(tokenizer, tools)
# ----- Tag constants -----
self.tool_call_start_token: str = "<DSMLfunction_calls>"
self.tool_call_end_token: str = "</DSMLfunction_calls>"
self.invoke_end_token: str = "</DSMLinvoke>"
self.param_end_token: str = "</DSMLparameter>"
# Alias expected by ToolParser base / adjust_request
self.tool_calls_start_token = self.tool_call_start_token
# ----- Compiled regexes -----
# Matches a complete <DSMLfunction_calls>…</DSMLfunction_calls>
self.tool_call_complete_regex = re.compile(
r"<DSMLfunction_calls>(.*?)</DSMLfunction_calls>", re.DOTALL
)
# Opening tag of an invoke block — captures the function name.
self.invoke_start_regex = re.compile(
r'<DSMLinvoke\s+name="([^"]+)"\s*>', re.DOTALL
)
# Complete invoke block.
self.invoke_complete_regex = re.compile(
r'<DSMLinvoke\s+name="([^"]+)"\s*>(.*?)</DSMLinvoke>',
re.DOTALL,
)
# Complete parameter tag — captures (name, string_attr, value).
self.parameter_complete_regex = re.compile(
r'<DSMLparameter\s+name="([^"]+)"\s+string="(true|false)"\s*>'
r"(.*?)"
r"</DSMLparameter>",
re.DOTALL,
)
# Just the opening header of a parameter tag (for partial params).
self.parameter_header_regex = re.compile(
r'<DSMLparameter\s+name="([^"]+)"\s+string="(true|false)"\s*>',
re.DOTALL,
)
# ----- Streaming state (reset per request) -----
self._sent_content_idx: int = 0
self._tool_call_ids: list[str] = []
self.streamed_args_for_tool: list[str] = []
self.prev_tool_call_arr: list[dict[str, Any]] = []
self.current_tool_id: int = -1
if not self.model_tokenizer:
raise ValueError(
"The model tokenizer must be passed to the ToolParser "
"constructor during construction."
)
logger.debug(
"Successfully initialized %s", self.__class__.__name__
)
# ------------------------------------------------------------------
# Request adjustment
# ------------------------------------------------------------------
def adjust_request(
self, request: ChatCompletionRequest | ResponsesRequest
) -> ChatCompletionRequest | ResponsesRequest:
request = super().adjust_request(request)
if request.tools and request.tool_choice != "none":
# Ensure DSML tokens are not stripped during decoding.
request.skip_special_tokens = False
return request
# ------------------------------------------------------------------
# Static / utility helpers
# ------------------------------------------------------------------
@staticmethod
def _tools_enabled(request: ChatCompletionRequest) -> bool:
"""Check whether tool calling is active for this request."""
try:
tools = getattr(request, "tools", None)
tool_choice = getattr(request, "tool_choice", None)
return bool(tools) and tool_choice != "none"
except Exception:
logger.exception("Failed to determine if tools are enabled.")
return False
def _generate_tool_call_id(self) -> str:
return f"call_{uuid.uuid4().hex[:24]}"
@staticmethod
def _json_escape_string_content(s: str) -> str:
"""JSON-escape a string value (without surrounding quotes)."""
if not s:
return ""
return json.dumps(s, ensure_ascii=False)[1:-1]
# ------------------------------------------------------------------
# Type conversion helpers
# ------------------------------------------------------------------
def _convert_param_value_checked(self, value: str, param_type: str) -> Any:
"""Convert a raw string value to the type indicated by *param_type*.
Raises on failure so the caller can try the next candidate type.
"""
if value.lower() == "null":
return None
param_type = param_type.lower()
if param_type in ("string", "str", "text"):
return value
elif param_type in ("integer", "int"):
return int(value)
elif param_type in ("number", "float"):
val = float(value)
return val if val != int(val) else int(val)
elif param_type in ("boolean", "bool"):
normed = value.strip().lower()
if normed not in ("false", "0", "true", "1"):
raise ValueError(f"Invalid boolean value: {value!r}")
return normed in ("true", "1")
elif param_type in ("object", "array"):
return json.loads(value)
else:
return json.loads(value)
def _convert_param_value(self, value: str, param_type: str | list[str]) -> Any:
"""Try each candidate type in turn; fall back to the raw string."""
if not isinstance(param_type, list):
param_type = [param_type]
for current_type in param_type:
try:
return self._convert_param_value_checked(value, current_type)
except Exception:
continue
return value
def _get_param_schema_type(
self, func_name: str, param_name: str
) -> str | list[str]:
"""Look up the JSON-schema type for a parameter, defaulting to
``"string"``."""
if self.tools:
for tool in self.tools:
if (
hasattr(tool, "function")
and tool.function.name == func_name
and hasattr(tool.function, "parameters")
):
schema = tool.function.parameters
if isinstance(schema, dict) and "properties" in schema:
prop = schema["properties"].get(param_name, {})
if isinstance(prop, dict):
return prop.get("type", "string")
break
return "string"
def _convert_with_schema(
self, func_name: str, param_name: str, value: str
) -> Any:
"""Convert *value* using the tool schema for *func_name*.*param_name*."""
param_type = self._get_param_schema_type(func_name, param_name)
return self._convert_param_value(value, param_type)
def _is_string_type(self, func_name: str, param_name: str) -> bool:
"""Return True if the schema says this parameter is a string."""
ptype = self._get_param_schema_type(func_name, param_name)
if isinstance(ptype, list):
return "string" in ptype
return ptype in ("string", "str", "text")
# ------------------------------------------------------------------
# Non-streaming extraction (unchanged logic, shared helpers)
# ------------------------------------------------------------------
def extract_tool_calls(
self,
model_output: str,
request: ChatCompletionRequest,
) -> ExtractedToolCallInformation:
"""Extract tool calls from complete model output (non-streaming)."""
if self.tool_call_start_token not in model_output:
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
try:
tool_calls: list[ToolCall] = []
for fc_block in self.tool_call_complete_regex.findall(model_output):
for invoke_name, invoke_body in self.invoke_complete_regex.findall(
fc_block
):
# Parse all parameters in this invoke.
raw_params: dict[str, str] = {}
for pname, _str_attr, pval in (
self.parameter_complete_regex.findall(invoke_body)
):
raw_params[pname] = pval
# Convert types via schema.
converted: dict[str, Any] = {}
for pname, pval in raw_params.items():
converted[pname] = self._convert_with_schema(
invoke_name, pname, pval
)
tool_calls.append(
ToolCall(
type="function",
function=FunctionCall(
name=invoke_name,
arguments=json.dumps(
converted, ensure_ascii=False
),
),
)
)
if not tool_calls:
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
first_idx = model_output.find(self.tool_call_start_token)
content = model_output[:first_idx] if first_idx > 0 else None
return ExtractedToolCallInformation(
tools_called=True, tool_calls=tool_calls, content=content
)
except Exception:
logger.exception("Error extracting tool calls from complete output")
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
# ------------------------------------------------------------------
# Streaming helpers — re-parse-and-diff
# ------------------------------------------------------------------
def _reset_streaming_state(self) -> None:
self._sent_content_idx = 0
self._tool_call_ids.clear()
self.streamed_args_for_tool.clear()
self.prev_tool_call_arr.clear()
self.current_tool_id = -1
def _extract_content(self, current_text: str) -> str | None:
"""Return any non-tool-call text that hasn't been sent yet.
Walks *current_text* from ``_sent_content_idx``, collecting text
outside ``<DSMLfunction_calls>`` regions. Uses
``partial_tag_overlap`` to avoid emitting bytes that might turn
out to be the start of the function-calls tag once the next
chunk arrives.
"""
content_segments: list[str] = []
pos = self._sent_content_idx
while pos < len(current_text):
start = current_text.find(self.tool_call_start_token, pos)
if start == -1:
# No (more) tool-call regions — send the tail, minus
# any suffix that could be the beginning of the tag.
tail = current_text[pos:]
overlap = partial_tag_overlap(tail, self.tool_call_start_token)
sendable = tail[: len(tail) - overlap] if overlap else tail
if sendable:
content_segments.append(sendable)
pos = len(current_text) - overlap
break
# Text between previous position and the tag start is content.
if start > pos:
content_segments.append(current_text[pos:start])
# Skip past the tool-call region.
end = current_text.find(self.tool_call_end_token, start)
if end != -1:
pos = end + len(self.tool_call_end_token)
else:
# Region still open — park cursor at start, stop.
pos = start
break
if content_segments:
self._sent_content_idx = pos
return "".join(content_segments)
if pos > self._sent_content_idx:
self._sent_content_idx = pos
return None
def _extract_invoke_regions(
self, text: str
) -> list[tuple[str, str, bool]]:
"""Find all invoke blocks inside the function_calls region.
Returns a list of ``(func_name, inner_text, is_complete)``
tuples. *inner_text* is everything between the invoke open
tag and the close tag (or the end of available text for the
last, potentially incomplete, invoke).
"""
results: list[tuple[str, str, bool]] = []
fc_start = text.find(self.tool_call_start_token)
if fc_start == -1:
return results
region_start = fc_start + len(self.tool_call_start_token)
fc_end = text.find(self.tool_call_end_token, region_start)
region = text[region_start:fc_end] if fc_end != -1 else text[region_start:]
pos = 0
while pos < len(region):
inv_match = self.invoke_start_regex.search(region, pos)
if not inv_match:
break
func_name = inv_match.group(1)
body_start = inv_match.end()
inv_end_pos = region.find(self.invoke_end_token, body_start)
if inv_end_pos != -1:
# Complete invoke block.
body = region[body_start:inv_end_pos]
results.append((func_name, body, True))
pos = inv_end_pos + len(self.invoke_end_token)
else:
# Incomplete — still being generated.
body = region[body_start:]
overlap = partial_tag_overlap(body, self.invoke_end_token)
if overlap:
body = body[:-overlap]
results.append((func_name, body, False))
break
return results
def _build_args_json_so_far(
self,
func_name: str,
inner_text: str,
is_complete: bool,
) -> str:
"""Build a JSON arguments string from the parameters found so far.
Handles both fully-closed ``<DSMLparameter>`` tags and the
single trailing partial parameter whose value is still being
streamed.
"""
# ---- Collect all fully-closed parameters ----
complete_params = self.parameter_complete_regex.findall(inner_text)
parts: list[str] = []
for param_name, string_attr, param_value in complete_params:
key_json = json.dumps(param_name, ensure_ascii=False)
if string_attr == "true":
val_json = json.dumps(param_value, ensure_ascii=False)
else:
converted = self._convert_with_schema(
func_name, param_name, param_value
)
val_json = json.dumps(converted, ensure_ascii=False)
parts.append(f"{key_json}: {val_json}")
# ---- Handle a trailing partial parameter ----
last_param_open = inner_text.rfind("<DSMLparameter")
last_param_close = inner_text.rfind(self.param_end_token)
has_partial = last_param_open != -1 and (
last_param_close == -1 or last_param_close < last_param_open
)
if has_partial:
partial_text = inner_text[last_param_open:]
header_match = self.parameter_header_regex.search(partial_text)
if header_match:
param_name = header_match.group(1)
string_attr = header_match.group(2)
partial_value = partial_text[header_match.end():]
# Strip any bytes that might be the beginning of the
# closing </DSMLparameter> tag.
overlap = partial_tag_overlap(
partial_value, self.param_end_token
)
if overlap:
partial_value = partial_value[:-overlap]
key_json = json.dumps(param_name, ensure_ascii=False)
if is_complete:
# Invoke is closed — treat whatever we have as final.
if string_attr == "true":
val_json = json.dumps(
partial_value, ensure_ascii=False
)
else:
converted = self._convert_with_schema(
func_name, param_name, partial_value
)
val_json = json.dumps(converted, ensure_ascii=False)
parts.append(f"{key_json}: {val_json}")
elif string_attr == "true" or self._is_string_type(
func_name, param_name
):
# Stream as an open JSON string (no closing quote).
escaped = self._json_escape_string_content(partial_value)
parts.append(f'{key_json}: "{escaped}')
else:
# Non-string — emit raw partial value.
parts.append(f"{key_json}: {partial_value}")
# ---- Assemble ----
if not parts:
return "{}" if is_complete else ""
joined = "{" + ", ".join(parts)
if is_complete:
joined += "}"
return joined
def _compute_args_diff(self, index: int, args_so_far: str) -> str | None:
"""Return only the characters in *args_so_far* that haven't been
sent yet, or ``None`` if there's nothing new."""
prev = self.streamed_args_for_tool[index]
if not args_so_far or len(args_so_far) <= len(prev):
return None
diff = args_so_far[len(prev):]
self.streamed_args_for_tool[index] = args_so_far
self.prev_tool_call_arr[index]["arguments"] = args_so_far
return diff
def _ensure_tool_state_for(self, index: int) -> None:
"""Grow the streaming-state arrays so *index* is valid."""
while len(self._tool_call_ids) <= index:
self._tool_call_ids.append(self._generate_tool_call_id())
while len(self.streamed_args_for_tool) <= index:
self.streamed_args_for_tool.append("")
while len(self.prev_tool_call_arr) <= index:
self.prev_tool_call_arr.append({})
# ------------------------------------------------------------------
# Main streaming entry point
# ------------------------------------------------------------------
def extract_tool_calls_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
request: ChatCompletionRequest,
) -> DeltaMessage | None:
"""Extract tool calls from streaming output using re-parse-and-diff.
On every call we:
1. Re-scan *current_text* for content outside tool-call regions.
2. Find all ``<DSMLinvoke>`` regions (complete + partial).
3. Build JSON args for each, diff against previous, emit deltas.
Because the entire text is re-parsed each time, the result is
correct regardless of how many tokens arrived in this step.
"""
# First chunk of a new stream — reset state.
if not previous_text:
self._reset_streaming_state()
# If tools aren't enabled, just forward content.
if not self._tools_enabled(request):
return DeltaMessage(content=delta_text) if delta_text else None
# 1. Extract any content outside tool-call regions.
content = self._extract_content(current_text)
# 2. Find all invoke regions.
regions = self._extract_invoke_regions(current_text)
tool_call_deltas: list[DeltaToolCall] = []
for i, (func_name, inner_text, is_complete) in enumerate(regions):
self._ensure_tool_state_for(i)
# Emit the tool name (once per tool call).
if "name" not in self.prev_tool_call_arr[i]:
self.prev_tool_call_arr[i]["name"] = func_name
tool_call_deltas.append(
DeltaToolCall(
index=i,
id=self._tool_call_ids[i],
type="function",
function=DeltaFunctionCall(
name=func_name,
arguments="",
),
)
)
# Build the JSON args so far and emit the diff.
args_so_far = self._build_args_json_so_far(
func_name, inner_text, is_complete
)
diff = self._compute_args_diff(i, args_so_far)
if diff:
tool_call_deltas.append(
DeltaToolCall(
index=i,
function=DeltaFunctionCall(arguments=diff),
)
)
if regions:
self.current_tool_id = len(regions) - 1
# 3. Return a delta if we have content or tool-call updates.
if content or tool_call_deltas:
return DeltaMessage(
content=content,
tool_calls=tool_call_deltas,
)
# Empty delta with token ids means EOS or closing tag — return
# non-None so the serving framework can finalize finish_reason.
if not delta_text and delta_token_ids and self.prev_tool_call_arr:
return DeltaMessage(content="")
return None

View File

@@ -1,119 +0,0 @@
[gMASK]<sop>
{%- if tools -%}
{%- macro tool_to_json(tool) -%}
{%- set ns_tool = namespace(first=true) -%}
{{ '{' -}}
{%- for k, v in tool.items() -%}
{%- if k != 'defer_loading' and k != 'strict' -%}
{%- if not ns_tool.first -%}{{- ', ' -}}{%- endif -%}
{%- set ns_tool.first = false -%}
"{{ k }}": {{ v | tojson(ensure_ascii=False) }}
{%- endif -%}
{%- endfor -%}
{{- '}' -}}
{%- endmacro -%}
<|system|>
# Tools
You may call one or more functions to assist with the user query.
You are provided with function signatures within <tools></tools> XML tags:
<tools>
{% for tool in tools %}
{%- if 'function' in tool -%}
{%- set tool = tool['function'] -%}
{%- endif -%}
{% if tool.defer_loading is not defined or not tool.defer_loading %}
{{ tool_to_json(tool) }}
{% endif %}
{% endfor %}
</tools>
For each function call, output the function name and arguments within the following XML format:
<tool_call>{function-name}<arg_key>{arg-key-1}</arg_key><arg_value>{arg-value-1}</arg_value><arg_key>{arg-key-2}</arg_key><arg_value>{arg-value-2}</arg_value>...</tool_call>{%- endif -%}
{%- macro visible_text(content) -%}
{%- if content is string -%}
{{- content }}
{%- elif content is iterable and content is not mapping -%}
{%- for item in content -%}
{%- if item is mapping and item.type == 'text' -%}
{{- item.text }}
{%- elif item is string -%}
{{- item }}
{%- endif -%}
{%- endfor -%}
{%- else -%}
{{- content }}
{%- endif -%}
{%- endmacro -%}
{%- set ns = namespace(last_user_index=-1, thinking_indices='') -%}
{%- for m in messages %}
{%- if m.role == 'user' %}
{%- set ns.last_user_index = loop.index0 -%}
{%- elif m.role == 'assistant' %}
{%- if m.reasoning_content is string %}
{%- set ns.thinking_indices = ns.thinking_indices ~ ',' ~ ns.last_user_index ~ ',' -%}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- set ns.has_thinking = false -%}
{%- for m in messages -%}
{%- if m.role == 'user' -%}<|user|>{{ visible_text(m.content) }}{% set ns.has_thinking = (',' ~ loop.index0 ~ ',') in ns.thinking_indices -%}
{%- elif m.role == 'assistant' -%}
<|assistant|>
{%- set content = visible_text(m.content) %}
{%- if m.reasoning_content is string %}
{%- set reasoning_content = m.reasoning_content %}
{%- elif '</think>' in content %}
{%- set reasoning_content = content.split('</think>')[0].split('<think>')[-1] %}
{%- set content = content.split('</think>')[-1] %}
{%- elif loop.index0 > ns.last_user_index and not (enable_thinking is defined and not enable_thinking) %}
{%- set reasoning_content = '' %}
{%- elif loop.index0 < ns.last_user_index and ns.has_thinking %}
{%- set reasoning_content = '' %}
{%- endif %}
{%- if ((clear_thinking is defined and not clear_thinking) or loop.index0 > ns.last_user_index) and reasoning_content is defined -%}
{{ '<think>' + reasoning_content + '</think>'}}
{%- else -%}
{{ '</think>' }}
{%- endif -%}
{%- if content.strip() -%}
{{ content.strip() }}
{%- endif -%}
{% if m.tool_calls %}
{% for tc in m.tool_calls %}
{%- if tc.function %}
{%- set tc = tc.function %}
{%- endif %}
{{- '<tool_call>' + tc.name -}}
{% set _args = tc.arguments %}{% for k, v in _args.items() %}<arg_key>{{ k }}</arg_key><arg_value>{{ v | tojson(ensure_ascii=False) if v is not string else v }}</arg_value>{% endfor %}</tool_call>{% endfor %}
{% endif %}
{%- elif m.role == 'tool' -%}
{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
{{- '<|observation|>' -}}
{%- endif %}
{%- if m.content is string -%}
{{- '<tool_response>' + m.content + '</tool_response>' -}}
{%- elif m.content is iterable and m.content is not mapping and m.content and m.content.0.type == "tool_reference" -%}
{{- '<tool_response><tools>\n' -}}
{% for tr in m.content %}
{%- for tool in tools -%}
{%- if 'function' in tool -%}
{%- set tool = tool['function'] -%}
{%- endif -%}
{%- if tool.name == tr.name -%}
{{- tool_to_json(tool) + '\n' -}}
{%- endif -%}
{%- endfor -%}
{%- endfor -%}
{{- '</tools></tool_response>' -}}
{%- else -%}
{{- '<tool_response>' + visible_text(m.content) + '</tool_response>' -}}
{% endif -%}
{%- elif m.role == 'system' -%}
<|system|>{{ visible_text(m.content) }}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
<|assistant|>{{- '</think>' if (enable_thinking is defined and not enable_thinking) else '<think>' -}}
{%- endif -%}

View File

@@ -1,777 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
import torch
import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.deep_gemm import (
get_paged_mqa_logits_metadata,
has_deep_gemm,
)
from vllm.utils.math_utils import cdiv
from vllm.utils.platform_utils import num_compute_units
from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
AttentionMetadataBuilder,
CommonAttentionMetadata,
MultipleOf,
)
from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping
from vllm.v1.attention.backends.utils import (
split_decodes_and_prefills,
)
from vllm.v1.kv_cache_interface import AttentionSpec, MLAAttentionSpec
from vllm.v1.worker.cp_utils import get_total_cp_world_size
logger = init_logger(__name__)
@triton.jit
def _prepare_uniform_decode_kernel(
seq_lens_ptr,
decode_seq_lens_ptr,
block_table_ptr,
block_table_stride,
expanded_block_table_ptr,
expanded_bt_stride,
decode_lens_ptr,
max_decode_len,
BLOCK_SIZE: tl.constexpr,
):
idx = tl.program_id(0)
req_id = idx // max_decode_len
local_idx = idx % max_decode_len
# Compute number of KVs attended to by this token.
seq_len = tl.load(seq_lens_ptr + req_id)
per_token_seq_len = seq_len - max_decode_len + local_idx + 1
tl.store(decode_seq_lens_ptr + idx, per_token_seq_len)
# Copy block table row.
src = block_table_ptr + req_id * block_table_stride
dst = expanded_block_table_ptr + idx * expanded_bt_stride
for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE):
off = i + tl.arange(0, BLOCK_SIZE)
mask = off < expanded_bt_stride
src_block = tl.load(src + off, mask=mask)
tl.store(dst + off, src_block, mask=mask)
# All reqs now have decode_len = 1.
tl.store(decode_lens_ptr + idx, 1)
def split_indexer_prefill_chunks(
seq_lens_cpu: torch.Tensor,
query_lens_cpu: torch.Tensor,
workspace_size: int,
max_logits_bytes: int,
request_offset: int = 0,
) -> list[tuple[slice, slice]]:
"""
Split prefill requests into chunks for the sparse indexer, respecting:
- N constraint: total_seq_lens <= workspace_size (existing O(N) workspace)
- Logits constraint: M * N * 4 <= max_logits_bytes
When a single request-level chunk still exceeds the logits budget,
sub-chunks on the query dimension (M) to bound peak memory.
Returns list of (req_slice, query_slice) tuples.
"""
chunks: list[tuple[slice, slice]] = []
n = len(seq_lens_cpu)
max_logits_elems = max_logits_bytes // 4
end = 0
while end < n:
start, chunk_m, chunk_n = end, 0, 0
while end < n:
q, s = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
new_m, new_n = chunk_m + q, chunk_n + s
if new_n <= workspace_size and new_m * new_n <= max_logits_elems:
chunk_m, chunk_n = new_m, new_n
end += 1
else:
break
# A single request can exceed the budget, requiring sub-chunking
# on the query dimension.
if end == start:
chunk_m, chunk_n = query_lens_cpu[end].item(), seq_lens_cpu[end].item()
end += 1
req_slice = slice(start + request_offset, end + request_offset)
max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m
for q_off in range(0, chunk_m, max_q):
sub_m = min(max_q, chunk_m - q_off)
chunks.append((req_slice, slice(q_off, q_off + sub_m)))
return chunks
class DeepseekV32IndexerBackend(AttentionBackend):
@staticmethod
def get_name() -> str:
return "DEEPSEEK_V32_INDEXER"
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
return [1, 64] if current_platform.is_rocm() else [64]
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
return [32, 64, 128]
@staticmethod
def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]:
return DeepseekV32IndexerMetadataBuilder
@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
assert num_kv_heads == 1
return (num_blocks, block_size, head_size)
@staticmethod
def get_kv_cache_stride_order(
include_num_layers_dimension: bool = False,
) -> tuple[int, ...]:
if include_num_layers_dimension:
# DeepseekV32Indexer kernels do not support cross-layer
# KV cache layout. Identity permutation keeps num_layers
# first, signaling incompatibility.
return (0, 1, 2, 3)
return (0, 1, 2)
class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
@staticmethod
def get_name() -> str:
return "DEEPSEEK_V4_INDEXER"
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
return [256]
@dataclass
class DeepseekV32IndexerPrefillChunkMetadata:
block_table: torch.Tensor
cu_seqlen_ks: torch.Tensor
cu_seqlen_ke: torch.Tensor
cu_seq_lens: torch.Tensor
token_to_seq: torch.Tensor
total_seq_lens: int
token_start: int
token_end: int
num_reqs: int
skip_kv_gather: bool = False
@dataclass
class DeepseekV32IndexerPrefillMetadata:
chunks: list[DeepseekV32IndexerPrefillChunkMetadata]
@dataclass
class DeepSeekV32IndexerDecodeMetadata:
block_table: torch.Tensor
# seq_lens: per-token effective context lengths.
# - flatten path / plain decode: 1D (batch_size,)
# - native MTP path: 2D (B, next_n) where [b,j] = L_b - next_n + j + 1
# Both fp8_fp4_paged_mqa_logits and the topk kernels accept both shapes.
seq_lens: torch.Tensor
decode_lens: torch.Tensor
requires_padding: bool
schedule_metadata: torch.Tensor
@dataclass
class DeepseekV32IndexerMetadata:
# FIXME (zyongye)
# hacky way to access the data now, need to be in chunked meta
seq_lens: torch.Tensor
max_seq_len: int
slot_mapping: torch.Tensor
# New for MLA (compared to FlashAttention)
# For handling prefill decode split
num_decodes: int
num_decode_tokens: int
num_prefills: int
num_prefill_tokens: int
decode: DeepSeekV32IndexerDecodeMetadata | None = None
prefill: DeepseekV32IndexerPrefillMetadata | None = None
def get_max_prefill_buffer_size(vllm_config: VllmConfig):
max_model_len = vllm_config.model_config.max_model_len
# NOTE(Chen): 40 is a magic number for controlling the prefill buffer size.
# Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes.
# The flashmla_sparse backend uses a workspace size of 5 * max_model_len.
# The memory usage of the workspace there is 576 * 2 bytes; so we size this as
# (576 * 2 // 132) * 5 = 40 to maximize this workspace size while still fitting
# within the flashmla_sparse workspace.
# For DeepSeek-V3.2, the max_model_len is 163840.
# 40 * 163840 * 132 = 865075200 bytes = 825 MB
return max_model_len * 40
class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
reorder_batch_threshold: int = 1
natively_supported_next_n_fp4: list[int] = [1, 2]
# TODO (matt): integrate kernel with next_n = 4 support
@classmethod
def get_cudagraph_support(
cls,
vllm_config: VllmConfig,
kv_cache_spec: AttentionSpec,
) -> AttentionCGSupport:
return AttentionCGSupport.UNIFORM_BATCH
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
scheduler_config = self.vllm_config.scheduler_config
# NOTE(Chen):an estimated max size of flattened_kv. Need to double check.
self.max_prefill_buffer_size = get_max_prefill_buffer_size(self.vllm_config)
self.num_speculative_tokens = (
self.vllm_config.speculative_config.num_speculative_tokens
if self.vllm_config.speculative_config
else 0
)
self.use_fp4_indexer_cache = (
self.vllm_config.attention_config.use_fp4_indexer_cache
)
assert (
current_platform.is_device_capability_family(100)
or not self.use_fp4_indexer_cache
), (
"use_fp4_indexer_cache requires Blackwell datacenter GPUs "
"(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and "
"earlier architectures are not supported."
)
next_n = self.num_speculative_tokens + 1
self.reorder_batch_threshold += self.num_speculative_tokens
# NOTE(zyongye) fp4 indexer cache only natively supports next_n in
# natively_supported_next_n_fp4; for other next_n values we fall back
# to the flattening path. Outside the SM100 datacenter family the FP8
# paged MQA logits kernel has the same [1, 2] constraint (deepgemm
# smxx_fp8_fp4_paged_mqa_logits.hpp:233), so flatten there too.
self.use_flattening = (
self.use_fp4_indexer_cache
or not current_platform.is_device_capability_family(100)
) and next_n not in self.natively_supported_next_n_fp4
sm_count = num_compute_units(self.device.index)
self.num_sms = sm_count
self.offsets_buffer = torch.arange(
next_n, device=self.device, dtype=torch.int32
)
self.decode_lens_buffer = torch.zeros(
(scheduler_config.max_num_batched_tokens,),
dtype=torch.int32,
device=self.device,
)
if not self.use_flattening and next_n > 1:
# Native MTP: 2D buffer for per-token seq_lens.
self.decode_seq_lens_buffer = torch.zeros(
(scheduler_config.max_num_seqs, next_n),
dtype=torch.int32,
device=self.device,
)
else:
# Flattening or no MTP: 1D buffer for expanded per-token seq_lens.
self.decode_seq_lens_buffer = torch.zeros(
(scheduler_config.max_num_batched_tokens,),
dtype=torch.int32,
device=self.device,
)
self.arange_buffer = torch.arange(
max(
scheduler_config.max_num_seqs * next_n,
scheduler_config.max_num_batched_tokens,
),
dtype=torch.int32,
device=self.device,
)
max_num_blocks_per_req = cdiv(
self.vllm_config.model_config.max_model_len,
self.kv_cache_spec.block_size * get_total_cp_world_size(),
)
self.expanded_block_table_buffer = torch.zeros(
(
scheduler_config.max_num_batched_tokens,
max_num_blocks_per_req,
),
dtype=torch.int32,
device=self.device,
)
# See: DeepGMM/csrc/apis/attention.hpp
self.scheduler_metadata_buffer = torch.empty(
(self.num_sms + 1, 2), dtype=torch.int32, device=self.device
)
# KV compression. Default to 1 for no compression.
self.compress_ratio = 1
# Get compress_ratio for DeepseekV4 support
if isinstance(self.kv_cache_spec, MLAAttentionSpec):
self.compress_ratio = self.kv_cache_spec.compress_ratio
# Pre-allocate buffers for CUDA graph compatibility when
if self.compress_ratio > 1:
# compress_ratio > 1 (DeepseekV4)
# Compressed slot mapping output buffer
self.compressed_slot_mapping_buffer = torch.zeros(
(scheduler_config.max_num_batched_tokens,),
dtype=torch.int64,
device=self.device,
)
# Buffer for compressed seq_lens in decode path
self.expanded_seq_lens_buffer = torch.zeros(
(scheduler_config.max_num_batched_tokens,),
dtype=torch.int32,
device=self.device,
)
def _prepare_decode_tensors(
self,
seq_lens: torch.Tensor,
block_table: torch.Tensor,
decode_lens: torch.Tensor,
decode_lens_cpu: torch.Tensor,
query_start_loc: torch.Tensor,
num_decodes: int,
num_decode_tokens: int,
use_native: bool,
next_n: int,
max_decode_len: int,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]:
"""Expand seq_lens/block_table/decode_lens for the decode kernels.
Flatten path (not use_native, max_decode_len > 1):
Each multi-token decode request is expanded into individual
single-token entries so the kernel always sees next_n=1.
Native path (use_native or max_decode_len == 1):
Plain decode or spec-decode with 2D per-token context lengths.
Returns (seq_lens, block_table, decode_lens, batch_size, requires_padding).
seq_lens is 1D (batch_size,) for flatten/plain, 2D (B, next_n) for native MTP.
"""
min_decode_len = int(decode_lens_cpu.min().item())
if not use_native and max_decode_len > 1:
assert self.decode_seq_lens_buffer.dim() == 1
if min_decode_len == max_decode_len:
# Uniform decode lengths.
num_decode_tokens = num_decodes * max_decode_len
_prepare_uniform_decode_kernel[(num_decode_tokens,)](
seq_lens,
self.decode_seq_lens_buffer,
block_table,
block_table.stride(0),
self.expanded_block_table_buffer,
self.expanded_block_table_buffer.stride(0),
self.decode_lens_buffer,
max_decode_len,
BLOCK_SIZE=1024,
)
self.decode_seq_lens_buffer[num_decode_tokens:] = 0
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
return seq_lens, block_table, decode_lens, num_decode_tokens, False
else:
# Variable decode lengths.
# Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is
# padding) and decode_lens [3, 1, 4, 0] in the below example comments.
# The context lengths are therefore
# [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0].
# 3 + 1 + 4 + 0 = 8
actual_expanded = int(decode_lens_cpu.sum().item())
# Fuse expanded_base and expanded_starts into a single
# repeat_interleave:
# seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1
# where context_start[b] = seq_lens[b] - decode_lens[b].
# Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8]
# expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4]
# result = [8, 9, 10, 7, 9, 10, 11, 12]
expanded_offsets = torch.repeat_interleave(
seq_lens - decode_lens - query_start_loc,
decode_lens,
output_size=actual_expanded,
)
# [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space
self.decode_seq_lens_buffer[:actual_expanded] = (
expanded_offsets + self.arange_buffer[:actual_expanded] + 1
)
self.decode_seq_lens_buffer[actual_expanded:] = 0
seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens]
# Give each of the flattened entries the same block table row as the
# original request.
self.expanded_block_table_buffer[:actual_expanded] = (
torch.repeat_interleave(
block_table, decode_lens, dim=0, output_size=actual_expanded
)
)
if actual_expanded < num_decode_tokens:
self.expanded_block_table_buffer[
actual_expanded:num_decode_tokens, 0
] = 0
block_table = self.expanded_block_table_buffer[:num_decode_tokens]
# All reqs now have decode_len=1
self.decode_lens_buffer[:num_decode_tokens] = 1
decode_lens = self.decode_lens_buffer[:num_decode_tokens]
return seq_lens, block_table, decode_lens, num_decode_tokens, False
else:
# Native path: plain decode (next_n==1) or spec decode
# with 2D per-token context lengths (next_n > 1).
#
# When decode_lens are not truly uniform (e.g. some requests have
# decode_len < next_n due to padding or short prefills), the simple
# reshape in sparse_attn_indexer won't work. Use pack_seq_triton
# (requires_padding) instead.
requires_padding = min_decode_len != max_decode_len
if use_native and next_n > 1:
assert self.decode_seq_lens_buffer.dim() == 2
# (B, max_decode_len): token j attends to
# L - max_decode_len + j + 1 KV tokens.
self.decode_seq_lens_buffer[:num_decodes, :max_decode_len] = (
seq_lens.unsqueeze(1)
- max_decode_len
+ 1
+ self.offsets_buffer[:max_decode_len]
)
seq_lens = self.decode_seq_lens_buffer[:num_decodes, :max_decode_len]
return seq_lens, block_table, decode_lens, num_decodes, requires_padding
def build(
self,
common_prefix_len: int,
common_attn_metadata: CommonAttentionMetadata,
fast_build: bool = False,
) -> DeepseekV32IndexerMetadata:
num_reqs = common_attn_metadata.num_reqs
num_tokens = common_attn_metadata.num_actual_tokens
query_start_loc = common_attn_metadata.query_start_loc
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
seq_lens = common_attn_metadata.seq_lens
slot_mapping = common_attn_metadata.slot_mapping
block_table = common_attn_metadata.block_table_tensor
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
split_decodes_and_prefills(
common_attn_metadata,
decode_threshold=self.reorder_batch_threshold,
require_uniform=not self.use_flattening,
)
)
assert num_decodes + num_prefills == num_reqs
assert num_decode_tokens + num_prefill_tokens == num_tokens
compressed_slot_mapping = slot_mapping
compressed_seq_lens = seq_lens
if self.compress_ratio > 1:
compressed_slot_mapping = get_compressed_slot_mapping(
num_tokens,
query_start_loc,
seq_lens,
block_table,
self.kv_cache_spec.storage_block_size,
self.compress_ratio,
out=self.compressed_slot_mapping_buffer,
)
compressed_seq_lens = seq_lens // self.compress_ratio
prefill_metadata = None
if num_prefills > 0:
# This CPU value is an upper bound for async-spec extend rows. It
# is safe for chunking/allocation because CUDA metadata below is
# built from exact device seq_lens and gather ignores the tail.
assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
compressed_seq_lens_cpu = (
seq_lens_cpu // self.compress_ratio
if self.compress_ratio > 1
else seq_lens_cpu
)
prefill_query_lens_cpu = torch.diff(
query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1]
)
max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
# Upper bound is exact for prefill rows (the `[num_decodes:]`
# slice below).
assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
chunk_specs = split_indexer_prefill_chunks(
compressed_seq_lens_cpu[num_decodes:],
prefill_query_lens_cpu,
self.max_prefill_buffer_size,
max_logits_bytes,
request_offset=num_decodes,
)
chunks = []
for req_slice, query_slice in chunk_specs:
metadata = build_prefill_chunk_metadata(
req_slice.start,
req_slice.stop,
query_start_loc,
query_start_loc_cpu,
seq_lens,
compressed_seq_lens,
compressed_seq_lens_cpu,
common_attn_metadata.block_table_tensor,
self.compress_ratio,
query_slice=query_slice,
skip_kv_gather=query_slice.start > 0,
)
# Skip when total_seq_lens is 0 (i.e., no compressed token).
if metadata is not None:
chunks.append(metadata)
prefill_metadata = DeepseekV32IndexerPrefillMetadata(chunks)
decode_metadata = None
if num_decodes > 0:
torch.diff(
common_attn_metadata.query_start_loc[: num_decodes + 1],
out=self.decode_lens_buffer[:num_decodes],
)
decode_lens = self.decode_lens_buffer[:num_decodes]
decode_lens_cpu = torch.diff(
common_attn_metadata.query_start_loc_cpu[: num_decodes + 1]
)
seq_lens = common_attn_metadata.seq_lens[:num_decodes]
block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...]
max_decode_len = int(decode_lens_cpu.max().item())
next_n = 1 + self.num_speculative_tokens
use_native = not self.use_flattening and max_decode_len <= next_n
seq_lens, block_table, decode_lens, batch_size, requires_padding = (
self._prepare_decode_tensors(
seq_lens=seq_lens,
block_table=block_table,
decode_lens=decode_lens,
decode_lens_cpu=decode_lens_cpu,
query_start_loc=common_attn_metadata.query_start_loc[:num_decodes],
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
use_native=use_native,
next_n=next_n,
max_decode_len=max_decode_len,
)
)
# For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores
# compressed tokens. Convert uncompressed seq_lens to compressed.
if self.compress_ratio > 1:
# True iff seq_lens aliases decode_seq_lens_buffer (flatten or
# native wrote it); False iff it aliases common_attn_metadata.
seq_lens_is_local_view = (use_native and next_n > 1) or (
not use_native and max_decode_len > 1
)
if seq_lens_is_local_view:
seq_lens //= self.compress_ratio
else:
# Copy to avoid mutating shared state; keeps CG address stable.
self.expanded_seq_lens_buffer[:num_decodes] = (
seq_lens // self.compress_ratio
)
self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0
seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens]
# Non-MTP: deep_gemm paged MQA logits requires 2D context_lens
# (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream
# kernels see the same (B, next_n) layout as the MTP path.
if seq_lens.dim() == 1:
seq_lens = seq_lens.unsqueeze(-1)
seq_lens = seq_lens.contiguous()
# DeepGEMM is required for the paged MQA logits on CUDA devices
if current_platform.is_cuda() and has_deep_gemm():
self.scheduler_metadata_buffer[:] = get_paged_mqa_logits_metadata(
seq_lens,
self.kv_cache_spec.storage_block_size,
self.num_sms,
)
decode_metadata = DeepSeekV32IndexerDecodeMetadata(
block_table=block_table,
seq_lens=seq_lens,
decode_lens=decode_lens,
requires_padding=requires_padding,
schedule_metadata=self.scheduler_metadata_buffer,
)
attn_metadata = DeepseekV32IndexerMetadata(
seq_lens=common_attn_metadata.seq_lens,
max_seq_len=common_attn_metadata.max_seq_len,
slot_mapping=compressed_slot_mapping,
num_decodes=num_decodes,
num_decode_tokens=num_decode_tokens,
num_prefills=num_prefills,
num_prefill_tokens=num_prefill_tokens,
prefill=prefill_metadata,
decode=decode_metadata,
)
return attn_metadata
def build_prefill_chunk_metadata(
start_idx: int,
end_idx: int,
query_start_loc: torch.Tensor,
query_start_loc_cpu: torch.Tensor,
uncompressed_seq_lens: torch.Tensor,
compressed_seq_lens: torch.Tensor,
compressed_seq_lens_cpu: torch.Tensor,
block_table: torch.Tensor,
compress_ratio: int,
query_slice: slice | None = None,
skip_kv_gather: bool = False,
) -> DeepseekV32IndexerPrefillChunkMetadata | None:
total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item()
if total_seq_lens == 0:
return None
num_reqs = end_idx - start_idx
device = block_table.device
token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device)
cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device)
# Assigning to slice avoids cpu sync.
cu_seq_lens[:1] = 0
torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:])
query_start_loc = (
query_start_loc[start_idx : end_idx + 1] - query_start_loc[start_idx]
)
total_query_len = int(
(query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item()
)
if query_slice is not None:
qs_start = query_slice.start
qs_stop = query_slice.stop
else:
qs_start = 0
qs_stop = total_query_len
output_query_len = qs_stop - qs_start
cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device)
cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device)
_build_prefill_chunk_metadata_kernel[(num_reqs,)](
query_start_loc,
uncompressed_seq_lens[start_idx:end_idx],
cu_seq_lens,
token_to_seq,
cu_seq_len_ks,
cu_seq_len_ke,
qs_start,
qs_stop,
BLOCK_SIZE=1024,
COMPRESS_RATIO=compress_ratio,
)
token_start = query_start_loc_cpu[start_idx].item()
if query_slice is not None:
token_end = token_start + qs_stop
token_start = token_start + qs_start
skip_kv_gather = skip_kv_gather or qs_start > 0
else:
token_end = query_start_loc_cpu[end_idx].item()
return DeepseekV32IndexerPrefillChunkMetadata(
cu_seqlen_ks=cu_seq_len_ks,
cu_seqlen_ke=cu_seq_len_ke,
cu_seq_lens=cu_seq_lens,
token_to_seq=token_to_seq,
total_seq_lens=total_seq_lens,
block_table=block_table[start_idx:end_idx],
token_start=token_start,
token_end=token_end,
num_reqs=num_reqs,
skip_kv_gather=skip_kv_gather,
)
@triton.jit
def _build_prefill_chunk_metadata_kernel(
# Inputs
query_start_loc_ptr,
uncompressed_seq_lens_ptr,
cu_compressed_seq_lens_ptr,
# Outputs
token_to_seq_ptr,
cu_compressed_seq_len_ks_ptr,
cu_compressed_seq_len_ke_ptr,
query_slice_start,
query_slice_stop,
BLOCK_SIZE: tl.constexpr,
COMPRESS_RATIO: tl.constexpr,
):
batch_idx = tl.program_id(0)
query_start = tl.load(query_start_loc_ptr + batch_idx)
query_end = tl.load(query_start_loc_ptr + batch_idx + 1)
query_len = query_end - query_start
seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx)
seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1)
compressed_seq_len = seq_end - seq_start
uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx)
start_pos = uncompressed_seq_len - query_len
for i in range(0, query_len, BLOCK_SIZE):
offset = i + tl.arange(0, BLOCK_SIZE)
abs_pos = query_start + offset
mask = (
(offset < query_len)
& (abs_pos >= query_slice_start)
& (abs_pos < query_slice_stop)
)
out_pos = abs_pos - query_slice_start
# Compute cu_seq_len_ks
tl.store(cu_compressed_seq_len_ks_ptr + out_pos, seq_start, mask=mask)
# Compute cu_seq_len_ke
seq_len_per_token = (start_pos + 1 + offset) // COMPRESS_RATIO
tl.store(
cu_compressed_seq_len_ke_ptr + out_pos,
seq_start + seq_len_per_token,
mask=mask,
)
# Compute token_to_seq
for i in range(0, compressed_seq_len, BLOCK_SIZE):
offset = i + tl.arange(0, BLOCK_SIZE)
mask = offset < compressed_seq_len
tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask)

View File

@@ -1,13 +0,0 @@
chunk_size: 256
local_cpu: true
max_local_cpu_size: 512.0
enable_lazy_memory_allocator: true
lazy_memory_initial_ratio: 0.2
remote_url: "redis://10.66.0.100:6379"
remote_serde: naive
remote_ttl: 1800
save_decode_cache: true
use_layerwise: true
save_unfull_chunk: true
blocking_timeout_secs: 60
cache_policy: LRU

View File

@@ -1,10 +0,0 @@
chunk_size: 256
local_cpu: true
max_local_cpu_size: 256.0
save_decode_cache: true
enable_lazy_memory_allocator: true
lazy_memory_initial_ratio: 1.0
use_gpu_connector_v3: true
remote_url: "redis://10.66.0.100:6379"
remote_serde: naive
remote_ttl: 1800

61
minimax_m2_parser.py Normal file
View File

@@ -0,0 +1,61 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
MiniMax M2 Parser - A unified parser for MiniMax M2 models.
This parser combines the existing MiniMaxM2ReasoningParser and
MinimaxM2ToolParser into a single unified interface by delegating
to those implementations.
"""
from vllm.logger import init_logger
from vllm.parser.abstract_parser import DelegatingParser
from vllm.reasoning.minimax_m2_reasoning_parser import MiniMaxM2ReasoningParser
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import (
Tool,
)
from vllm.tool_parsers.minimax_m2_tool_parser import MinimaxM2ToolParser
logger = init_logger(__name__)
class MiniMaxM2Parser(DelegatingParser):
"""
Unified parser for MiniMax M2 models that handles both reasoning
extraction and tool call parsing.
This parser delegates to the existing implementations:
- MiniMaxM2ReasoningParser for reasoning extraction
- MinimaxM2ToolParser for tool call parsing
MiniMax M2 models have two special behaviors:
1. Reasoning: They don't generate <think> start token, only </think> end
token. All content before </think> is reasoning, content after is the
actual response.
2. Tool Calls: They use <minimax:tool_call>...</minimax:tool_call> tags
with <invoke name="...">...</invoke> and <parameter name="...">...</parameter>
syntax.
"""
# Class-level parser classes for compatibility
reasoning_parser_cls = MiniMaxM2ReasoningParser
tool_parser_cls = MinimaxM2ToolParser
def __init__(
self,
tokenizer: TokenizerLike,
tools: list[Tool] | None = None,
*args,
**kwargs,
):
super().__init__(tokenizer, *args, **kwargs)
# Initialize the underlying parsers
self._reasoning_parser = MiniMaxM2ReasoningParser(tokenizer, *args, **kwargs)
self._tool_parser = MinimaxM2ToolParser(tokenizer, tools)
logger.debug(
"vLLM Successfully initialized parser %s!", self.__class__.__name__
)

852
minimax_tool_parser.py Normal file
View File

@@ -0,0 +1,852 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
from collections.abc import Sequence
from typing import Any
import regex as re
from vllm.entrypoints.chat_utils import make_tool_call_id
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
from vllm.entrypoints.openai.engine.protocol import (
DeltaFunctionCall,
DeltaMessage,
DeltaToolCall,
ExtractedToolCallInformation,
FunctionCall,
ToolCall,
)
from vllm.logger import init_logger
from vllm.tokenizers import TokenizerLike
from vllm.tool_parsers.abstract_tool_parser import (
Tool,
ToolParser,
)
from vllm.tool_parsers.utils import extract_intermediate_diff
logger = init_logger(__name__)
class MinimaxToolParser(ToolParser):
def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None, **kwargs):
super().__init__(tokenizer, tools)
# Initialize streaming state for tracking tool call progress
self.streaming_state: dict[str, Any] = {
"current_tool_index": -1, # Index of current tool being processed
"tool_ids": [], # List of tool call IDs
"sent_tools": [], # List of tools that have been sent
}
# Define tool call tokens and patterns
self.tool_call_start_token = "<tool_calls>"
self.tool_call_end_token = "</tool_calls>"
self.tool_call_regex = re.compile(
r"<tool_calls>(.*?)</tool_calls>|<tool_calls>(.*)", re.DOTALL
)
self.thinking_tag_pattern = r"<think>(.*?)</think>"
self.tool_name_pattern = re.compile(r'"name":\s*"([^"]+)"')
self.tool_args_pattern = re.compile(r'"arguments":\s*')
# Buffer for handling partial tool calls during streaming
self.pending_buffer = ""
self.in_thinking_tag = False
if not self.model_tokenizer:
raise ValueError(
"The model tokenizer must be passed to the ToolParser "
"constructor during construction."
)
# Get token IDs for tool call start/end tokens
self.tool_call_start_token_id = self.vocab.get(self.tool_call_start_token)
self.tool_call_end_token_id = self.vocab.get(self.tool_call_end_token)
if self.tool_call_start_token_id is None or self.tool_call_end_token_id is None:
logger.warning(
"Minimax Tool parser could not locate tool call start/end "
"tokens in the tokenizer. Falling back to string matching."
)
def preprocess_model_output(self, model_output: str) -> str:
"""
Preprocess model output by removing tool calls from thinking tags.
Args:
model_output: Raw model output string
Returns:
Preprocessed model output with tool calls removed from thinking tags
"""
def remove_tool_calls_from_think(match):
think_content = match.group(1)
cleaned_content = re.sub(
r"<tool_calls>.*?</tool_calls>", "", think_content, flags=re.DOTALL
)
return f"<think>{cleaned_content}</think>"
return re.sub(
self.thinking_tag_pattern,
remove_tool_calls_from_think,
model_output,
flags=re.DOTALL,
)
def _clean_duplicate_braces(self, args_text: str) -> str:
"""
Clean duplicate closing braces from arguments text.
Args:
args_text: Raw arguments text
Returns:
Cleaned arguments text with proper JSON formatting
"""
args_text = args_text.strip()
if not args_text:
return args_text
try:
json.loads(args_text)
return args_text
except json.JSONDecodeError:
pass
while args_text.endswith("}}"):
candidate = args_text[:-1]
try:
json.loads(candidate)
return candidate
except json.JSONDecodeError:
args_text = candidate
return args_text
def _clean_delta_braces(self, delta_text: str) -> str:
"""
Clean delta text by removing excessive closing braces.
Args:
delta_text: Delta text to clean
Returns:
Cleaned delta text
"""
if not delta_text:
return delta_text
delta_stripped = delta_text.strip()
if delta_stripped and all(c in "}\n\r\t " for c in delta_stripped):
brace_count = delta_stripped.count("}")
if brace_count > 1:
return "}\n" if delta_text.endswith("\n") else "}"
return delta_text
def extract_tool_calls(
self,
model_output: str,
request: ChatCompletionRequest,
) -> ExtractedToolCallInformation:
"""
Extract tool calls from model output for non-streaming mode.
Args:
model_output: Complete model output
request: Chat completion request
Returns:
ExtractedToolCallInformation containing tool calls and content
"""
processed_output = self.preprocess_model_output(model_output)
if self.tool_call_start_token not in processed_output:
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
try:
function_call_tuples = self.tool_call_regex.findall(processed_output)
raw_function_calls = []
for match in function_call_tuples:
tool_call_content = match[0] if match[0] else match[1]
if tool_call_content.strip():
lines = tool_call_content.strip().split("\n")
for line in lines:
line = line.strip()
if line and line.startswith("{") and line.endswith("}"):
try:
parsed_call = json.loads(line)
raw_function_calls.append(parsed_call)
except json.JSONDecodeError:
continue
tool_calls = []
for function_call in raw_function_calls:
if "name" in function_call and "arguments" in function_call:
tool_calls.append(
ToolCall(
type="function",
function=FunctionCall(
name=function_call["name"],
arguments=json.dumps(
function_call["arguments"], ensure_ascii=False
),
),
)
)
processed_pos = processed_output.find(self.tool_call_start_token)
if processed_pos != -1:
processed_content = processed_output[:processed_pos].strip()
if processed_content:
lines = processed_content.split("\n")
for line in reversed(lines):
line = line.strip()
if line:
pos = model_output.find(line)
if pos != -1:
content = model_output[: pos + len(line)]
break
else:
content = ""
else:
content = ""
else:
content = model_output
return ExtractedToolCallInformation(
tools_called=len(tool_calls) > 0,
tool_calls=tool_calls,
content=content.strip() if content.strip() else None,
)
except Exception:
logger.exception(
"An unexpected error occurred during tool call extraction."
)
return ExtractedToolCallInformation(
tools_called=False, tool_calls=[], content=model_output
)
def _update_thinking_state(self, text: str) -> None:
"""
Update the thinking tag state based on text content.
Args:
text: Text to analyze for thinking tags
"""
open_count = text.count("<think>")
close_count = text.count("</think>")
self.in_thinking_tag = open_count > close_count or (
open_count == close_count and text.endswith("</think>")
)
def _is_potential_tag_start(self, text: str) -> bool:
"""
Check if text might be the start of a tool call tag.
Args:
text: Text to check
Returns:
True if text could be the start of a tool call tag
"""
for tag in [self.tool_call_start_token, self.tool_call_end_token]:
if any(
tag.startswith(text[-i:])
for i in range(1, min(len(text) + 1, len(tag)))
):
return True
return False
def _should_buffer_content(self, delta_text: str) -> bool:
"""
Determine if content should be buffered for later processing.
Args:
delta_text: Delta text to check
Returns:
True if content should be buffered
"""
if self.in_thinking_tag:
return False
return bool(
self.pending_buffer
or self.tool_call_start_token in delta_text
or self.tool_call_end_token in delta_text
or delta_text.startswith("<")
)
def _split_content_for_buffering(self, delta_text: str) -> tuple[str, str]:
"""
Split delta text into safe content and potential tag content.
Args:
delta_text: Delta text to split
Returns:
Tuple of (safe_content, potential_tag_content)
"""
if self.in_thinking_tag:
return delta_text, ""
for tag in [self.tool_call_start_token, self.tool_call_end_token]:
for i in range(1, len(tag)):
tag_prefix = tag[:i]
pos = delta_text.rfind(tag_prefix)
if pos != -1 and tag.startswith(delta_text[pos:]):
return delta_text[:pos], delta_text[pos:]
return delta_text, ""
def _process_buffer(self, new_content: str) -> str:
"""
Process buffered content and return output content.
Args:
new_content: New content to add to buffer
Returns:
Processed output content
"""
self.pending_buffer += new_content
output_content = ""
if self.in_thinking_tag:
output_content = self.pending_buffer
self.pending_buffer = ""
return output_content
while self.pending_buffer:
start_pos = self.pending_buffer.find(self.tool_call_start_token)
end_pos = self.pending_buffer.find(self.tool_call_end_token)
if start_pos != -1 and (end_pos == -1 or start_pos < end_pos):
tag_pos, tag_len = start_pos, len(self.tool_call_start_token)
elif end_pos != -1:
tag_pos, tag_len = end_pos, len(self.tool_call_end_token)
else:
if self._is_potential_tag_start(self.pending_buffer):
break
output_content += self.pending_buffer
self.pending_buffer = ""
break
output_content += self.pending_buffer[:tag_pos]
self.pending_buffer = self.pending_buffer[tag_pos + tag_len :]
return output_content
def _reset_streaming_state(self) -> None:
"""Reset the streaming state to initial values."""
self.streaming_state = {
"current_tool_index": -1,
"tool_ids": [],
"sent_tools": [],
}
def _advance_to_next_tool(self) -> None:
"""Advance to the next tool in the streaming sequence."""
self.streaming_state["current_tool_index"] = (
int(self.streaming_state["current_tool_index"]) + 1
)
def _set_current_tool_index(self, index: int) -> None:
"""
Set the current tool index.
Args:
index: Tool index to set
"""
self.streaming_state["current_tool_index"] = index
def _get_current_tool_index(self) -> int:
"""
Get the current tool index.
Returns:
Current tool index
"""
return int(self.streaming_state["current_tool_index"])
def _get_next_unsent_tool_index(self, tool_count: int) -> int:
"""
Get the index of the next unsent tool.
Args:
tool_count: Total number of tools
Returns:
Index of next unsent tool, or -1 if all tools sent
"""
sent_tools = list(self.streaming_state["sent_tools"])
for i in range(tool_count):
if i < len(sent_tools):
if not sent_tools[i]["sent_name"]:
return i
else:
return i
return -1
def _ensure_state_arrays(self, tool_count: int) -> None:
"""
Ensure state arrays have sufficient capacity for tool_count tools.
Args:
tool_count: Number of tools to prepare for
"""
sent_tools = list(self.streaming_state["sent_tools"])
tool_ids = list(self.streaming_state["tool_ids"])
while len(sent_tools) < tool_count:
sent_tools.append(
{
"sent_name": False,
"sent_arguments": "",
"id": make_tool_call_id(),
}
)
while len(tool_ids) < tool_count:
tool_ids.append(None)
self.streaming_state["sent_tools"] = sent_tools
self.streaming_state["tool_ids"] = tool_ids
def _detect_tools_in_text(self, text: str) -> int:
"""
Detect the number of tools in text by counting name patterns.
Args:
text: Text to analyze
Returns:
Number of tools detected
"""
matches = self.tool_name_pattern.findall(text)
return len(matches)
def _find_tool_boundaries(self, text: str) -> list[tuple[int, int]]:
"""
Find the boundaries of tool calls in text.
Args:
text: Text to analyze
Returns:
List of (start, end) positions for tool calls
"""
boundaries = []
i = 0
while i < len(text):
if text[i] == "{":
start = i
depth = 0
has_name = False
has_arguments = False
while i < len(text):
if text[i] == "{":
depth += 1
elif text[i] == "}":
depth -= 1
if depth == 0:
end = i + 1
segment = text[start:end]
if '"name"' in segment and '"arguments"' in segment:
boundaries.append((start, end))
break
if not has_name and '"name"' in text[start : i + 1]:
has_name = True
if not has_arguments and '"arguments"' in text[start : i + 1]:
has_arguments = True
i += 1
if depth > 0 and has_name:
boundaries.append((start, i))
else:
i += 1
return boundaries
def _extract_tool_args(self, tool_content: str, args_match: re.Match[str]) -> str:
"""
Extract tool arguments from tool content.
Args:
tool_content: Tool call content
args_match: Regex match for arguments pattern
Returns:
Extracted arguments as string
"""
args_start_pos = args_match.end()
remaining_content = tool_content[args_start_pos:]
if remaining_content.strip().startswith("{"):
depth = 0
for i, char in enumerate(remaining_content):
if char == "{":
depth += 1
elif char == "}":
depth -= 1
if depth == 0:
return remaining_content[: i + 1]
else:
args_end = remaining_content.find("}")
if args_end > 0:
return remaining_content[:args_end].strip()
return remaining_content.rstrip("}").strip()
def _get_current_tool_content(
self, text: str, tool_index: int
) -> tuple[str | None, str | None]:
"""
Get the content of a specific tool by index.
Args:
text: Text containing tool calls
tool_index: Index of tool to extract
Returns:
Tuple of (tool_name, tool_arguments) or (None, None) if not found
"""
boundaries = self._find_tool_boundaries(text)
if tool_index >= len(boundaries):
return None, None
start, end = boundaries[tool_index]
tool_content = text[start:end]
name_match = self.tool_name_pattern.search(tool_content)
name = name_match.group(1) if name_match else None
args_match = self.tool_args_pattern.search(tool_content)
if args_match:
try:
args_text = self._extract_tool_args(tool_content, args_match)
return name, args_text
except Exception:
remaining_content = tool_content[args_match.end() :]
args_text = remaining_content.rstrip("}").strip()
return name, args_text
return name, None
def _handle_tool_name_streaming(
self, tool_content: str, tool_count: int
) -> DeltaMessage | None:
"""
Handle streaming of tool names.
Args:
tool_content: Content containing tool calls
tool_count: Total number of tools
Returns:
DeltaMessage with tool name or None if no tool to stream
"""
next_idx = self._get_next_unsent_tool_index(tool_count)
if next_idx == -1:
return None
boundaries = self._find_tool_boundaries(tool_content)
if next_idx >= len(boundaries):
return None
tool_name, _ = self._get_current_tool_content(tool_content, next_idx)
if not tool_name:
return None
self._set_current_tool_index(next_idx)
sent_tools = list(self.streaming_state["sent_tools"])
tool_ids = list(self.streaming_state["tool_ids"])
tool_id = sent_tools[next_idx]["id"]
tool_ids[next_idx] = tool_id
sent_tools[next_idx]["sent_name"] = True
self.streaming_state["sent_tools"] = sent_tools
self.streaming_state["tool_ids"] = tool_ids
return DeltaMessage(
tool_calls=[
DeltaToolCall(
index=next_idx,
type="function",
id=tool_id,
function=DeltaFunctionCall(name=tool_name).model_dump(
exclude_none=True
),
)
]
)
def _handle_tool_args_streaming(
self, tool_content: str, tool_count: int
) -> DeltaMessage | None:
"""
Handle streaming of tool arguments.
Args:
tool_content: Content containing tool calls
tool_count: Total number of tools
Returns:
DeltaMessage with tool arguments or None if no arguments to stream
"""
current_idx = self._get_current_tool_index()
if current_idx < 0 or current_idx >= tool_count:
return None
tool_name, tool_args = self._get_current_tool_content(tool_content, current_idx)
if not tool_name or tool_args is None:
return None
sent_tools = list(self.streaming_state["sent_tools"])
if not sent_tools[current_idx]["sent_name"]:
return None
clean_args = self._clean_duplicate_braces(tool_args)
sent_args = sent_tools[current_idx]["sent_arguments"]
if clean_args != sent_args:
if sent_args and clean_args.startswith(sent_args):
args_delta = extract_intermediate_diff(clean_args, sent_args)
if args_delta:
args_delta = self._clean_delta_braces(args_delta)
sent_tools[current_idx]["sent_arguments"] = clean_args
self.streaming_state["sent_tools"] = sent_tools
if clean_args.endswith("}"):
self._advance_to_next_tool()
return DeltaMessage(
tool_calls=[
DeltaToolCall(
index=current_idx,
function=DeltaFunctionCall(
arguments=args_delta
).model_dump(exclude_none=True),
)
]
)
elif not sent_args and clean_args:
clean_args_delta = self._clean_delta_braces(clean_args)
sent_tools[current_idx]["sent_arguments"] = clean_args
self.streaming_state["sent_tools"] = sent_tools
if clean_args.endswith("}"):
self._advance_to_next_tool()
return DeltaMessage(
tool_calls=[
DeltaToolCall(
index=current_idx,
function=DeltaFunctionCall(
arguments=clean_args_delta
).model_dump(exclude_none=True),
)
]
)
return None
def _is_end_tool_calls(self, current_text: str) -> bool:
if self.tool_call_end_token not in current_text:
return False
end_token_positions = []
search_start = 0
while True:
pos = current_text.find(self.tool_call_end_token, search_start)
if pos == -1:
break
end_token_positions.append(pos)
search_start = pos + 1
think_regions = []
for match in re.finditer(
self.thinking_tag_pattern, current_text, flags=re.DOTALL
):
think_regions.append((match.start(), match.end()))
for pos in end_token_positions:
in_think = any(
pos >= t_start and pos < t_end for t_start, t_end in think_regions
)
if not in_think:
return True
return False
def extract_tool_calls_streaming(
self,
previous_text: str,
current_text: str,
delta_text: str,
previous_token_ids: Sequence[int],
current_token_ids: Sequence[int],
delta_token_ids: Sequence[int],
request: ChatCompletionRequest,
) -> DeltaMessage | None:
self._update_thinking_state(current_text)
if self.in_thinking_tag:
return DeltaMessage(content=delta_text)
if self._should_buffer_content(delta_text):
buffered_output = self._process_buffer(delta_text)
return DeltaMessage(content=buffered_output) if buffered_output else None
if self._is_end_tool_calls(current_text):
return DeltaMessage(content=delta_text)
safe_content, potential_tag = self._split_content_for_buffering(delta_text)
if potential_tag:
self.pending_buffer += potential_tag
return DeltaMessage(content=safe_content) if safe_content else None
processed_current_text = self.preprocess_model_output(current_text)
if self.tool_call_start_token not in processed_current_text:
if (
self.tool_call_end_token in delta_text
and self.tool_call_start_token in current_text
):
return None
if delta_text.strip() == "" and self.tool_call_start_token in current_text:
return None
if (
self._get_current_tool_index() != -1
and self.tool_call_end_token in current_text
):
self._reset_streaming_state()
return DeltaMessage(content=delta_text)
if (
self.tool_call_start_token_id is not None
and self.tool_call_start_token_id in delta_token_ids
and len(delta_token_ids) == 1
):
return None
original_tool_start = self._find_tool_start_outside_thinking(current_text)
if original_tool_start is None:
return None
content_before_tools = self._extract_content_before_tools(
current_text, delta_text, original_tool_start
)
if content_before_tools:
return DeltaMessage(content=content_before_tools)
try:
tool_content = self._extract_tool_content(current_text, original_tool_start)
current_tools_count = self._detect_tools_in_text(tool_content)
if current_tools_count == 0:
return None
if self._get_current_tool_index() == -1:
self._reset_streaming_state()
self._ensure_state_arrays(current_tools_count)
return self._handle_tool_name_streaming(
tool_content, current_tools_count
) or self._handle_tool_args_streaming(tool_content, current_tools_count)
except Exception:
logger.exception(
"An unexpected error occurred ", "during streaming tool call handling."
)
return None
def _find_tool_start_outside_thinking(self, current_text: str) -> int | None:
"""
Find the start position of tool calls outside of thinking tags.
Args:
current_text: Current text to search
Returns:
Position of tool call start or None if not found
"""
search_start = 0
while True:
pos = current_text.find(self.tool_call_start_token, search_start)
if pos == -1:
return None
think_regions = [
(m.start(), m.end())
for m in re.finditer(
r"<think>(.*?)</think>", current_text, flags=re.DOTALL
)
]
in_think = any(
pos >= t_start and pos < t_end for t_start, t_end in think_regions
)
if not in_think:
return pos
search_start = pos + 1
def _extract_content_before_tools(
self, current_text: str, delta_text: str, tool_start: int
) -> str | None:
"""
Extract content that appears before tool calls.
Args:
current_text: Current text
delta_text: Delta text
tool_start: Start position of tools
Returns:
Content before tools or None
"""
if tool_start > 0:
delta_start_pos = len(current_text) - len(delta_text)
if delta_start_pos < tool_start:
content_part = delta_text
if delta_start_pos + len(delta_text) > tool_start:
content_part = delta_text[: tool_start - delta_start_pos]
return content_part if content_part else None
return None
def _extract_tool_content(self, current_text: str, tool_start: int) -> str:
"""
Extract tool content from current text starting at tool_start.
Args:
current_text: Current text
tool_start: Start position of tool calls
Returns:
Extracted tool content
"""
tool_content_start = tool_start + len(self.tool_call_start_token)
tool_content = current_text[tool_content_start:]
end_pos = tool_content.find(self.tool_call_end_token)
if end_pos != -1:
tool_content = tool_content[:end_pos]
return tool_content

View File

@@ -0,0 +1,28 @@
from vllm.reasoning.abs_reasoning_parsers import ReasoningParserManager
from vllm.reasoning.deepseek_r1_reasoning_parser import DeepSeekR1ReasoningParser
@ReasoningParserManager.register_module("super_v3")
class SuperV3ReasoningParser(DeepSeekR1ReasoningParser):
def extract_reasoning(self, model_output, request):
reasoning_content, final_content = super().extract_reasoning(
model_output, request
)
if (
hasattr(request, "chat_template_kwargs")
and request.chat_template_kwargs
and (
request.chat_template_kwargs.get("enable_thinking") is False
or request.chat_template_kwargs.get("force_nonempty_content") is True
)
and final_content is None
):
"""
The original `deepseek_r1` reasoning parser this inherits from will automatically put everything in the reasoning content when it cannot parse out reasoning. This was fine for the DeepSeek R1 model that was not intended to be used without reasoning.
1. Since the Nemotron 3 Nano and Super both have thinking off modes modulated by "enable_thinking=false" in the chat template kwargs, this change instead which will properly place the content in cases where there is no thinking enabled via config.
2. There are rare cases where the model will output only reasoning without an end-think token `</think>` (e.g. reasoning exceeds max length), which results in empty content returned. End users may want to unilaterally avoid such cases and always have a content response even if the model does not finish its reasoning.
"""
# Put all nonempty content into the content, rather than return content
reasoning_content, final_content = None, reasoning_content
return reasoning_content, final_content