diff --git a/docs/contributing/model/basic.md b/docs/contributing/model/basic.md index 59e57e4ad14..0cc24baae92 100644 --- a/docs/contributing/model/basic.md +++ b/docs/contributing/model/basic.md @@ -136,7 +136,7 @@ The model should also be added to the `MODELS_CONFIG_MAP` dictionary in [vllm/mo For case (2), we recommend using as a reference the implementation of [`JambaForCausalLM`](../../../vllm/model_executor/models/jamba.py) (for an example of a model that uses Mamba-1 and attention together) or [`NemotronHForCausalLM`](../../../vllm/model_executor/models/nemotron_h.py) (for an example of a model that uses Mamba-2 and attention together). These models should follow the same instructions as case (1), but they should inherit protocol `IsHybrid` (instead of `IsAttentionFree`) and it is *not* necessary to add them to the `MODELS_CONFIG_MAP` (their runtime defaults will be inferred from the protocol). -For case (3), we recommend looking at the implementation of [`MiniMaxText01ForCausalLM`](../../../vllm/model_executor/models/minimax_text_01.py) or [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which use custom "mamba-like" layers `MiniMaxText01LinearAttention` and `ShortConv` respectively. +For case (3), we recommend looking at the implementation of [`Lfm2ForCausalLM`](../../../vllm/model_executor/models/lfm2.py) as a reference, which uses a custom "mamba-like" layer `ShortConv`. Please follow the same guidelines as case (2) for implementing these models. We use "mamba-like" to refer to layers that possess a state that is updated in-place, rather than being appended-to (like KV cache for attention). For implementing new custom mamba-like layers, one should inherit from `MambaBase` and implement the methods `get_state_dtype`, `get_state_shape` to calculate the data types and state shapes at runtime, as well as `mamba_type` and `get_attn_backend`. @@ -144,5 +144,5 @@ It is also necessary to implement the "attention meta-data" class which handles Please see [`LinearAttentionMetadata`](../../../vllm/v1/attention/backends/linear_attn.py) or [`ShortConvAttentionMetadata`](../../../vllm/v1/attention/backends/short_conv_attn.py) for examples of this. It is also worth noting that we should update `MambaAttentionBackendEnum` in [`registry.py`](../../../vllm/v1/attention/backends/registry.py) when adding a new mamba backend. Finally, if one wants to support torch compile and CUDA graphs, it necessary to wrap the call to the mamba-like layer inside a custom op and register it. -Please see the calls to `direct_register_custom_op` in [vllm/model_executor/models/minimax_text_01.py](../../../vllm/model_executor/models/minimax_text_01.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this. +Please see the calls to `direct_register_custom_op` in [vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py](../../../vllm/model_executor/layers/mamba/linear/minimax_linear_attn.py) or [vllm/model_executor/layers/mamba/short_conv.py](../../../vllm/model_executor/layers/mamba/short_conv.py) for examples of this. The new custom op should then be added to the list `_attention_ops` in [vllm/config/compilation.py](../../../vllm/config/compilation.py) to ensure that piecewise CUDA graphs works as intended. diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index 1d10a94c712..10626a254b1 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -321,15 +321,6 @@ For Qwen2.5, the chat template in tokenizer_config.json has already included sup Flags: `--tool-call-parser hermes` -### MiniMax Models (`minimax_m1`) - -Supported models: - -* `MiniMaxAi/MiniMax-M1-40k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja)) -* `MiniMaxAi/MiniMax-M1-80k` (use with [examples/tool_chat_template_minimax_m1.jinja](../../examples/tool_chat_template_minimax_m1.jinja)) - -Flags: `--tool-call-parser minimax --chat-template examples/tool_chat_template_minimax_m1.jinja` - ### DeepSeek-V3 Models (`deepseek_v3`) Supported models: diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index e67bc197d32..264f5a72195 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -441,7 +441,6 @@ th { | `MiMoV2ForCausalLM` | MiMoV2Pro | `XiaomiMiMo/MiMo-V2.5-Pro`, etc. | | ✅︎ | | `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ | | `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ | -| `MiniMaxForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01-hf`, etc. | | | | `MiniMaxM2ForCausalLM` | MiniMax-M2, MiniMax-M2.1 | `MiniMaxAI/MiniMax-M2`, etc. | ✅︎ | ✅︎ | | `MistralForCausalLM` | Ministral-3, Mistral, Mistral-Instruct | `mistralai/Ministral-3-3B-Instruct-2512`, `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ | | `MistralLarge3ForCausalLM` | Mistral-Large-3-675B-Base-2512, Mistral-Large-3-675B-Instruct-2512 | `mistralai/Mistral-Large-3-675B-Base-2512`, `mistralai/Mistral-Large-3-675B-Instruct-2512`, etc. | ✅︎ | ✅︎ | @@ -487,8 +486,6 @@ th { | `TeleChat2ForCausalLM` | TeleChat2 | `Tele-AI/TeleChat2-3B`, `Tele-AI/TeleChat2-7B`, `Tele-AI/TeleChat2-35B`, etc. | ✅︎ | ✅︎ | | `TeleChat3ForCausalLM` | TeleChat3 | `Tele-AI/TeleChat3-36B-Thinking`, `Tele-AI/TeleChat3-Coder-36B-Thinking`, etc. | ✅︎ | ✅︎ | | `TeleFLMForCausalLM` | TeleFLM | `CofeAI/FLM-2-52B-Instruct-2407`, `CofeAI/Tele-FLM`, etc. | ✅︎ | ✅︎ | -| `MiniMaxM1ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-M1-40k`, `MiniMaxAI/MiniMax-M1-80k`, etc. | | | -| `MiniMaxText01ForCausalLM` | MiniMax-Text | `MiniMaxAI/MiniMax-Text-01`, etc. | | | | `Zamba2ForCausalLM` | Zamba2 | `Zyphra/Zamba2-7B-instruct`, `Zyphra/Zamba2-2.7B-instruct`, `Zyphra/Zamba2-1.2B-instruct`, etc. | | | !!! note @@ -595,7 +592,6 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `MiMoV2OmniForCausalLM` | MiMo-V2.5-Omni | T + IE+ + VE+ + A+ | `XiaomiMiMo/MiMo-V2.5-Omni` | | ✅︎ | | `MiniCPMO` | MiniCPM-O | T + IE+ + VE+ + AE+ | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ | | `MiniCPMV` | MiniCPM-V | T + IE+ + VE+ | `openbmb/MiniCPM-V-2` (see note), `openbmb/MiniCPM-Llama3-V-2_5`, `openbmb/MiniCPM-V-2_6`, `openbmb/MiniCPM-V-4`, `openbmb/MiniCPM-V-4_5`, etc. | ✅︎ | | -| `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + IE+ | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ | | `Mistral3ForConditionalGeneration` | Mistral3 (HF Transformers) | T + I+ | `mistralai/Mistral-Small-3.1-24B-Instruct-2503`, etc. | ✅︎ | ✅︎ | | `MolmoForCausalLM` | Molmo | T + I+ | `allenai/Molmo-7B-D-0924`, `allenai/Molmo-7B-O-0924`, etc. | ✅︎ | ✅︎ | | `Molmo2ForConditionalGeneration` | Molmo2 | T + I+ / V | `allenai/Molmo2-4B`, `allenai/Molmo2-8B`, `allenai/Molmo2-O-7B`, `allenai/MolmoWeb-4B`^, `allenai/MolmoWeb-8B`^ | ✅︎ | ✅︎ | diff --git a/docs/usage/v1_guide.md b/docs/usage/v1_guide.md index 74d7e3eb2b0..eca23a11bc8 100644 --- a/docs/usage/v1_guide.md +++ b/docs/usage/v1_guide.md @@ -128,7 +128,7 @@ Models that use Mamba-2 and Mamba-1 layers (e.g., `Mamba2ForCausalLM`, `MambaFor Hybrid models that combine Mamba-2 and Mamba-1 layers with standard attention layers are also supported (e.g., `BambaForCausalLM`, `Zamba2ForCausalLM`, `NemotronHForCausalLM`, `FalconH1ForCausalLM` and `GraniteMoeHybridForCausalLM`, `JambaForCausalLM`, `Plamo2ForCausalLM`). -Hybrid models with mechanisms different to Mamba are also supported (e.g, `MiniMaxText01ForCausalLM`, `MiniMaxM1ForCausalLM`, `Lfm2ForCausalLM`). +Hybrid models with mechanisms different to Mamba are also supported (e.g, `Lfm2ForCausalLM`). Please note that prefix caching is not yet supported for any of the above models. diff --git a/examples/generate/multimodal/vision_language_offline.py b/examples/generate/multimodal/vision_language_offline.py index 1b3741a3e42..e837625908c 100644 --- a/examples/generate/multimodal/vision_language_offline.py +++ b/examples/generate/multimodal/vision_language_offline.py @@ -1481,39 +1481,6 @@ def run_minicpmv(questions: list[str], modality: str) -> ModelRequestData: return run_minicpmv_base(questions, modality, "openbmb/MiniCPM-V-2_6") -def run_minimax_vl_01(questions: list[str], modality: str) -> ModelRequestData: - assert modality == "image" - - model_name = "MiniMaxAI/MiniMax-VL-01" - - engine_args = EngineArgs( - model=model_name, - max_num_seqs=2, - limit_mm_per_prompt={modality: 1}, - trust_remote_code=True, - tensor_parallel_size=8, - ) - - tokenizer = AutoTokenizer.from_pretrained(model_name) - messages = [ - [ - { - "role": "user", - "content": [{"type": "image"}, {"type": "text", "text": question}], - } - ] - for question in questions - ] - prompts = tokenizer.apply_chat_template( - messages, add_generation_prompt=True, tokenize=False - ) - - return ModelRequestData( - engine_args=engine_args, - prompts=prompts, - ) - - # Mistral-3 HF-format def run_mistral3(questions: list[str], modality: str) -> ModelRequestData: assert modality == "image" @@ -2485,7 +2452,6 @@ model_example_map = { "mantis": run_mantis, "minicpmo": run_minicpmo, "minicpmv": run_minicpmv, - "minimax_vl_01": run_minimax_vl_01, "mistral3": run_mistral3, "molmo": run_molmo, "molmo2": run_molmo2, diff --git a/examples/tool_chat_template_minimax_m1.jinja b/examples/tool_chat_template_minimax_m1.jinja deleted file mode 100644 index 2d5bbf4de56..00000000000 --- a/examples/tool_chat_template_minimax_m1.jinja +++ /dev/null @@ -1,91 +0,0 @@ -{{ '' -}} -{%- if custom_tools is defined %} - {%- set tools = custom_tools %} -{%- endif %} -{%- if not tools is defined %} - {%- set tools = none %} -{%- endif %} - -{#- Extract system message #} -{% set ns = namespace(system_prompt='') -%} -{%- if messages[0]['role'] == 'system' %} - {%- if messages[0]['content'] is string %} - {%- set ns.system_prompt = messages[0]['content']|trim %} - {%- else %} - {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %} - {%- endif %} - {%- set messages = messages[1:] %} -{%- else %} - {%- if tools is not none %} - {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %} - {%- else %} - {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %} - {%- endif %} -{%- endif %} - -{#- System message #} -{%- if ns.system_prompt != '' %} -{{ 'system ai_setting=assistant\n' + ns.system_prompt + '\n' -}} -{%- endif %} - -{#- Tools configuration #} -{%- if tools is not none %} -{{ 'system tool_setting=tools\nYou are provided with these tools:\n\n' -}} -{%- for tool in tools %} -{{ tool | tojson ~ '\n' -}} -{%- endfor %} -{{ '\n\nIf you need to call tools, please respond with XML tags, and provide tool-name and json-object of arguments, following the format below:\n\n{"name": , "arguments": }\n...\n\n' -}} -{%- endif %} - -{#- Process messages #} -{%- for message in messages %} - {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %} - {%- if message['role'] == 'user' %} -{{ 'user name=user\n' -}} -{%- if message['content'] is string %} -{{ message['content']|trim -}} -{%- else %} -{%- for content in message['content'] %} -{%- if content['type'] == 'text' %} -{{ content['text']|trim -}} -{%- endif %} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- elif message['role'] == 'assistant' %} -{{ 'ai name=assistant\n' -}} -{%- if message['content'] is string %} -{{ message['content']|trim -}} -{%- else %} -{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %} -{{ content['text']|trim -}} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- endif %} - {%- elif 'tool_calls' in message %} -{{ 'ai name=assistant\n\n' -}} -{%- for tool_call in message.tool_calls %} -{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}} -{%- endfor %} -{{ '\n' -}} - {%- elif message.role == "tool" or message.role == "ipython" %} -{{ 'tool name=tools\n' -}} -{%- if message.content is string %} -{{ 'tool result: ' + message.content + '\n\n' -}} -{%- else %} -{%- for content in message['content'] %} -{%- if content['type'] == 'text' %} -{{ 'tool result: ' + content['text'] + '\n\n' -}} -{%- elif content.get('name') %} -{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}} -{%- endif %} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- endif %} -{%- endfor %} - -{%- if add_generation_prompt %} -{{ 'ai name=assistant\n' -}} -{%- endif %} \ No newline at end of file diff --git a/rust/src/chat/src/renderer/hf/format.rs b/rust/src/chat/src/renderer/hf/format.rs index 2c990fb37ba..4c0fb68e595 100644 --- a/rust/src/chat/src/renderer/hf/format.rs +++ b/rust/src/chat/src/renderer/hf/format.rs @@ -386,7 +386,6 @@ mod tests { tool_chat_template_llama3.2_pythonic.jinja => String tool_chat_template_llama4_json.jinja => OpenAi tool_chat_template_llama4_pythonic.jinja => OpenAi - tool_chat_template_minimax_m1.jinja => OpenAi tool_chat_template_mistral.jinja => String tool_chat_template_mistral3.jinja => OpenAi tool_chat_template_mistral_parallel.jinja => String diff --git a/rust/src/chat/tests/templates/vllm_examples/tool_chat_template_minimax_m1.jinja b/rust/src/chat/tests/templates/vllm_examples/tool_chat_template_minimax_m1.jinja deleted file mode 100644 index 2d5bbf4de56..00000000000 --- a/rust/src/chat/tests/templates/vllm_examples/tool_chat_template_minimax_m1.jinja +++ /dev/null @@ -1,91 +0,0 @@ -{{ '' -}} -{%- if custom_tools is defined %} - {%- set tools = custom_tools %} -{%- endif %} -{%- if not tools is defined %} - {%- set tools = none %} -{%- endif %} - -{#- Extract system message #} -{% set ns = namespace(system_prompt='') -%} -{%- if messages[0]['role'] == 'system' %} - {%- if messages[0]['content'] is string %} - {%- set ns.system_prompt = messages[0]['content']|trim %} - {%- else %} - {%- set ns.system_prompt = messages[0]['content'][0]['text']|trim %} - {%- endif %} - {%- set messages = messages[1:] %} -{%- else %} - {%- if tools is not none %} - {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %} - {%- else %} - {%- set ns.system_prompt = "You are a helpful assistant created by Minimax based on MiniMax-M1 model." %} - {%- endif %} -{%- endif %} - -{#- System message #} -{%- if ns.system_prompt != '' %} -{{ 'system ai_setting=assistant\n' + ns.system_prompt + '\n' -}} -{%- endif %} - -{#- Tools configuration #} -{%- if tools is not none %} -{{ 'system tool_setting=tools\nYou are provided with these tools:\n\n' -}} -{%- for tool in tools %} -{{ tool | tojson ~ '\n' -}} -{%- endfor %} -{{ '\n\nIf you need to call tools, please respond with XML tags, and provide tool-name and json-object of arguments, following the format below:\n\n{"name": , "arguments": }\n...\n\n' -}} -{%- endif %} - -{#- Process messages #} -{%- for message in messages %} - {%- if not (message.role == 'ipython' or message.role == 'tool' or 'tool_calls' in message) %} - {%- if message['role'] == 'user' %} -{{ 'user name=user\n' -}} -{%- if message['content'] is string %} -{{ message['content']|trim -}} -{%- else %} -{%- for content in message['content'] %} -{%- if content['type'] == 'text' %} -{{ content['text']|trim -}} -{%- endif %} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- elif message['role'] == 'assistant' %} -{{ 'ai name=assistant\n' -}} -{%- if message['content'] is string %} -{{ message['content']|trim -}} -{%- else %} -{%- for content in message['content'] | selectattr('type', 'equalto', 'text') %} -{{ content['text']|trim -}} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- endif %} - {%- elif 'tool_calls' in message %} -{{ 'ai name=assistant\n\n' -}} -{%- for tool_call in message.tool_calls %} -{{ '{"name": "' + tool_call.function.name + '", "arguments": ' + tool_call.function.arguments | tojson + '}\n' -}} -{%- endfor %} -{{ '\n' -}} - {%- elif message.role == "tool" or message.role == "ipython" %} -{{ 'tool name=tools\n' -}} -{%- if message.content is string %} -{{ 'tool result: ' + message.content + '\n\n' -}} -{%- else %} -{%- for content in message['content'] %} -{%- if content['type'] == 'text' %} -{{ 'tool result: ' + content['text'] + '\n\n' -}} -{%- elif content.get('name') %} -{{ 'tool name: ' + content['name'] + '\ntool result: ' + content['text'] + '\n\n' -}} -{%- endif %} -{%- endfor %} -{%- endif %} -{{ '\n' -}} - {%- endif %} -{%- endfor %} - -{%- if add_generation_prompt %} -{{ 'ai name=assistant\n' -}} -{%- endif %} \ No newline at end of file diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index a9afe73cad6..b6945fb0aa3 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -810,29 +810,6 @@ VLM_TEST_SETTINGS = { hf_output_post_proc=model_utils.minicpmv_trunc_hf_output, patch_hf_runner=model_utils.minicpmv_26_patch_hf_runner, ), - "minimax_vl_01": VLMTestInfo( - models=["MiniMaxAI/MiniMax-VL-01"], - prompt_formatter=lambda img_prompt: f"user: {img_prompt} assistant:", # noqa: E501 - img_idx_to_prompt=lambda _: "", - test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), - max_model_len=8192, - max_num_seqs=4, - dtype="bfloat16", - hf_output_post_proc=model_utils.minimax_vl_01_hf_output, - patch_hf_runner=model_utils.minimax_vl_01_patch_hf_runner, - auto_cls=AutoModelForImageTextToText, - marks=[ - large_gpu_mark(min_gb=80), - # TODO: [ROCm] Fix pickle issue with ROCm spawn and tp>1 - pytest.mark.skipif( - current_platform.is_rocm(), - reason=( - "ROCm: Model too large for single GPU; " - "multi-GPU blocked by HF _LazyConfigMapping pickle issue with spawn" - ), - ), - ], - ), "molmo": VLMTestInfo( models=["allenai/Molmo-7B-D-0924"], test_type=(VLMTestType.IMAGE, VLMTestType.MULTI_IMAGE), diff --git a/tests/models/multimodal/generation/vlm_utils/model_utils.py b/tests/models/multimodal/generation/vlm_utils/model_utils.py index 62ea36061c9..e3f08bf9237 100644 --- a/tests/models/multimodal/generation/vlm_utils/model_utils.py +++ b/tests/models/multimodal/generation/vlm_utils/model_utils.py @@ -245,13 +245,6 @@ def minicpmv_trunc_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutpu return output_ids, output_str, out_logprobs -def minimax_vl_01_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutput: - output_ids, output_str, out_logprobs = hf_output - if output_str.endswith(""): - output_str = output_str.split("")[0] - return output_ids, output_str, out_logprobs - - def ultravox_trunc_hf_output(hf_output: RunnerOutput, model: str) -> RunnerOutput: output_ids, output_str, out_logprobs = hf_output @@ -1023,17 +1016,6 @@ def minicpmv_26_patch_hf_runner(hf_model: HfRunner) -> HfRunner: return hf_model -def minimax_vl_01_patch_hf_runner(hf_model: HfRunner) -> HfRunner: - orig_generate = hf_model.model.generate - - def _generate(self, *args, image_sizes=None, **kwargs): - return orig_generate(*args, decode_text=False, **kwargs) - - hf_model.model.generate = types.MethodType(_generate, hf_model.model) - - return hf_model - - def molmo_patch_hf_runner(hf_model: HfRunner) -> HfRunner: """Patches and returns an instance of the HfRunner to use for Molmo.""" hf_processor = hf_model.processor diff --git a/tests/models/multimodal/processing/test_minimax_vl_01.py b/tests/models/multimodal/processing/test_minimax_vl_01.py deleted file mode 100644 index 9b4c4f9531e..00000000000 --- a/tests/models/multimodal/processing/test_minimax_vl_01.py +++ /dev/null @@ -1,113 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import pytest -from PIL import Image - -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.parse import ImageSize -from vllm.multimodal.processing import BaseMultiModalProcessor - -from ....conftest import ImageTestAssets -from ...utils import build_model_context - - -@pytest.mark.parametrize("model_id", ["MiniMaxAI/MiniMax-VL-01"]) -@pytest.mark.parametrize("num_imgs", [1, 2]) -def test_processor_override( - image_assets: ImageTestAssets, - model_id: str, - num_imgs: int, -): - ctx = build_model_context( - model_id, - mm_processor_kwargs=None, - limit_mm_per_prompt={"image": num_imgs}, - ) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) - prompt = "" * num_imgs - image = Image.new("RGB", size=(364, 364)) - mm_data = {"image": [image] * num_imgs} - - processed_inputs = processor( - prompt, - mm_items=processor.info.parse_mm_data(mm_data), - hf_processor_mm_kwargs={}, - ) - image_placeholders = processed_inputs["mm_placeholders"]["image"] - - assert len(image_placeholders) == num_imgs - - -def _validate_image_prompt_replacements_one( - processor: BaseMultiModalProcessor, - num_imgs: int, - failed_size_excs: list[tuple[ImageSize, Exception]], - image_size: ImageSize, -) -> None: - prompt = "" * num_imgs - image = Image.new("RGB", size=image_size) - mm_data = {"image": [image] * num_imgs} - - try: - processed_inputs = processor( - prompt, - mm_items=processor.info.parse_mm_data(mm_data), - hf_processor_mm_kwargs={}, - ) - - image_placeholders = processed_inputs["mm_placeholders"]["image"] - assert len(image_placeholders) == num_imgs - - except Exception as exc: - failed_size_excs.append((image_size, exc)) - - -def _test_image_prompt_replacements( - processor, - *, - num_imgs: int, - image_sizes: list[ImageSize], -) -> None: - failed_size_excs = list[tuple[ImageSize, Exception]]() - - for size in image_sizes: - _validate_image_prompt_replacements_one( - processor, num_imgs, failed_size_excs, size - ) - - if failed_size_excs: - msg = "Found failing image sizes:" + "\n========\n".join( - f"[{size}]\n{exc}" for size, exc in failed_size_excs - ) - raise AssertionError(msg) - - -@pytest.mark.parametrize("model_id", ["MiniMaxAI/MiniMax-VL-01"]) -@pytest.mark.parametrize("num_imgs", [1, 2]) -def test_processor_prompt_replacements_regression(model_id, num_imgs): - ctx = build_model_context( - model_id, - mm_processor_kwargs=None, - limit_mm_per_prompt={"image": num_imgs}, - ) - processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) - - image_ratios = [ - (171, 152), - (184, 161), - (198, 176), - (333, 296), - (369, 328), - (488, 183), - (2560, 1669), - ] - image_sizes = [ - size for w, h in image_ratios for size in [ImageSize(w, h), ImageSize(h, w)] - ] - - _test_image_prompt_replacements( - processor, - num_imgs=num_imgs, - image_sizes=image_sizes, - ) diff --git a/tests/models/registry.py b/tests/models/registry.py index e865f8efe85..8f7ea822642 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -421,15 +421,6 @@ _TEXT_GENERATION_EXAMPLE_MODELS = { }, trust_remote_code=True, ), - "MiniMaxForCausalLM": _HfExamplesInfo("MiniMaxAI/MiniMax-Text-01-hf"), - "MiniMaxText01ForCausalLM": _HfExamplesInfo( - "MiniMaxAI/MiniMax-Text-01", - trust_remote_code=True, - revision="a59aa9cbc53b9fb8742ca4e9e1531b9802b6fdc3", - ), - "MiniMaxM1ForCausalLM": _HfExamplesInfo( - "MiniMaxAI/MiniMax-M1-40k", trust_remote_code=True - ), "MiniMaxM2ForCausalLM": _HfExamplesInfo( "MiniMaxAI/MiniMax-M2", trust_remote_code=True, @@ -1113,10 +1104,6 @@ _MULTIMODAL_EXAMPLE_MODELS = { "openbmb/MiniCPM-V-4_6", min_transformers_version="5.7.0", ), - "MiniMaxVL01ForConditionalGeneration": _HfExamplesInfo( - "MiniMaxAI/MiniMax-VL-01", - trust_remote_code=True, - ), "MiniMaxM3SparseForConditionalGeneration": _HfExamplesInfo( "MiniMaxAI/MiniMax-M3", trust_remote_code=True, diff --git a/tests/models/test_initialization.py b/tests/models/test_initialization.py index 476ad1c7c17..6632d50bc0f 100644 --- a/tests/models/test_initialization.py +++ b/tests/models/test_initialization.py @@ -98,11 +98,6 @@ def can_initialize( vllm_config.validate_block_size() return scheduler_kv_cache_config - if model_arch == "MiniMaxVL01ForConditionalGeneration": - pytest.skip( - "pickle error when loading `transformers.models.auto.CONFIG_MAPPING`" - ) - if model_arch == "MoonshotKimiaForCausalLM": pytest.skip( "Kimi-Audio requires SpeechToTextConfig " diff --git a/tests/tool_parsers/test_minimax_tool_parser.py b/tests/tool_parsers/test_minimax_tool_parser.py deleted file mode 100644 index 08b2104277b..00000000000 --- a/tests/tool_parsers/test_minimax_tool_parser.py +++ /dev/null @@ -1,1227 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# ruff: noqa: E501 - -import json -from typing import Any - -import pytest - -from vllm.entrypoints.openai.chat_completion.protocol import ( - ChatCompletionToolsParam, -) -from vllm.entrypoints.openai.engine.protocol import ( - FunctionCall, - ToolCall, -) -from vllm.tokenizers import get_tokenizer -from vllm.tool_parsers.minimax_tool_parser import MinimaxToolParser - -# Use a common model that is likely to be available -MODEL = "MiniMaxAi/MiniMax-M1-40k" - - -@pytest.fixture(scope="module") -def minimax_tokenizer(): - return get_tokenizer(tokenizer_name=MODEL) - - -@pytest.fixture -def minimax_tool_parser(minimax_tokenizer): - return MinimaxToolParser(minimax_tokenizer) - - -@pytest.fixture -def sample_tools(): - return [ - ChatCompletionToolsParam( - type="function", - function={ - "name": "get_current_weather", - "description": "Get the current weather", - "parameters": { - "type": "object", - "properties": { - "city": {"type": "string", "description": "The city name"}, - "state": {"type": "string", "description": "The state code"}, - "unit": {"type": "string", "enum": ["fahrenheit", "celsius"]}, - }, - "required": ["city", "state"], - }, - }, - ), - ChatCompletionToolsParam( - type="function", - function={ - "name": "calculate_area", - "description": "Calculate area of a shape", - "parameters": { - "type": "object", - "properties": { - "shape": {"type": "string"}, - "dimensions": {"type": "object"}, - "precision": {"type": "integer"}, - }, - }, - }, - ), - ] - - -def assert_tool_calls( - actual_tool_calls: list[ToolCall], expected_tool_calls: list[ToolCall] -): - assert len(actual_tool_calls) == len(expected_tool_calls) - - for actual_tool_call, expected_tool_call in zip( - actual_tool_calls, expected_tool_calls - ): - assert isinstance(actual_tool_call.id, str) - assert len(actual_tool_call.id) > 16 - - assert actual_tool_call.type == "function" - assert actual_tool_call.function == expected_tool_call.function - - -def test_extract_tool_calls_no_tools(minimax_tool_parser): - model_output = "This is a test" - extracted_tool_calls = minimax_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - assert not extracted_tool_calls.tools_called - assert extracted_tool_calls.tool_calls == [] - assert extracted_tool_calls.content == model_output - - -@pytest.mark.parametrize( - ids=[ - "single_tool_call", - "multiple_tool_calls", - "tool_call_with_content_before", - "tool_call_with_single_line_json", - "tool_call_incomplete_tag", - ], - argnames=["model_output", "expected_tool_calls", "expected_content"], - argvalues=[ - ( - """ -{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}} -""", - [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "Dallas", - "state": "TX", - "unit": "fahrenheit", - } - ), - ) - ) - ], - None, - ), - ( - """ -{"name": "get_current_weather", "arguments": {"city": "Dallas", "state": "TX", "unit": "fahrenheit"}} -{"name": "get_current_weather", "arguments": {"city": "Orlando", "state": "FL", "unit": "fahrenheit"}} -""", - [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "Dallas", - "state": "TX", - "unit": "fahrenheit", - } - ), - ) - ), - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "Orlando", - "state": "FL", - "unit": "fahrenheit", - } - ), - ) - ), - ], - None, - ), - ( - """I'll help you check the weather. -{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}} -""", - [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "Seattle", - "state": "WA", - "unit": "celsius", - } - ), - ) - ) - ], - "I'll help you check the weather.", - ), - ( - """ -{"name": "get_current_weather", "arguments": {"city": "New York", "state": "NY", "unit": "celsius"}} -""", - [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "New York", - "state": "NY", - "unit": "celsius", - } - ), - ) - ) - ], - None, - ), - ( - """ -{"name": "get_current_weather", "arguments": {"city": "Boston", "state": "MA"}}""", - [ - ToolCall( - function=FunctionCall( - name="get_current_weather", - arguments=json.dumps( - { - "city": "Boston", - "state": "MA", - } - ), - ) - ) - ], - None, - ), - ], -) -def test_extract_tool_calls( - minimax_tool_parser, model_output, expected_tool_calls, expected_content -): - extracted_tool_calls = minimax_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - assert extracted_tool_calls.tools_called - - assert_tool_calls(extracted_tool_calls.tool_calls, expected_tool_calls) - - assert extracted_tool_calls.content == expected_content - - -def test_preprocess_model_output_with_thinking_tags(minimax_tool_parser): - """Test that tool calls within thinking tags are removed during preprocessing.""" - model_output = """Let me think about this. -{"name": "fake_tool", "arguments": {"param": "value"}} - This should be removed. - -I'll help you with that. -{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA"}} -""" - - processed_output = minimax_tool_parser.preprocess_model_output(model_output) - - # The tool call within thinking tags should be removed - assert "fake_tool" not in processed_output - # But the thinking tag itself should remain - assert "" in processed_output - assert "" in processed_output - # The actual tool call outside thinking tags should remain - assert "get_current_weather" in processed_output - - -def test_extract_tool_calls_with_thinking_tags(minimax_tool_parser): - """Test tool extraction when thinking tags contain tool calls that should be ignored.""" - model_output = """I should use a tool. -{"name": "ignored_tool", "arguments": {"should": "ignore"}} - - -Let me help you with the weather. -{"name": "get_current_weather", "arguments": {"city": "Miami", "state": "FL", "unit": "fahrenheit"}} -""" - - extracted_tool_calls = minimax_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - assert len(extracted_tool_calls.tool_calls) == 1 - assert extracted_tool_calls.tool_calls[0].function.name == "get_current_weather" - - # Content extraction is based on the position of the first in the original model_output - # Since preprocessing removes tool calls within thinking tags, the actual first is the external one - expected_content = """I should use a tool. -{"name": "ignored_tool", "arguments": {"should": "ignore"}} - - -Let me help you with the weather.""" - assert extracted_tool_calls.content == expected_content - - -def test_extract_tool_calls_invalid_json(minimax_tool_parser): - """Test that invalid JSON in tool calls is handled gracefully.""" - model_output = """ -{"name": "valid_tool", "arguments": {"city": "Seattle"}} -{invalid json here} -{"name": "another_valid_tool", "arguments": {"param": "value"}} -""" - - extracted_tool_calls = minimax_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - # Should extract only the valid JSON tool calls - assert len(extracted_tool_calls.tool_calls) == 2 - assert extracted_tool_calls.tool_calls[0].function.name == "valid_tool" - assert extracted_tool_calls.tool_calls[1].function.name == "another_valid_tool" - - -def test_extract_tool_calls_missing_name_or_arguments(minimax_tool_parser): - """Test that tool calls missing name or arguments are filtered out.""" - model_output = """ -{"name": "valid_tool", "arguments": {"city": "Seattle"}} -{"name": "missing_args"} -{"arguments": {"city": "Portland"}} -{"name": "another_valid_tool", "arguments": {"param": "value"}} -""" - - extracted_tool_calls = minimax_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - - assert extracted_tool_calls.tools_called - # Should extract only the valid tool calls with both name and arguments - assert len(extracted_tool_calls.tool_calls) == 2 - assert extracted_tool_calls.tool_calls[0].function.name == "valid_tool" - assert extracted_tool_calls.tool_calls[1].function.name == "another_valid_tool" - - -def test_streaming_basic_functionality(minimax_tool_parser): - """Test basic streaming functionality.""" - # Reset streaming state - minimax_tool_parser.current_tool_name_sent = False - minimax_tool_parser.prev_tool_call_arr = [] - minimax_tool_parser.current_tool_id = -1 - minimax_tool_parser.streamed_args_for_tool = [] - - # Test with a simple tool call - current_text = """ -{"name": "get_current_weather", "arguments": {"city": "Seattle"}} -""" - - # First call should handle the initial setup - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text="", - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - # The result might be None or contain tool call information - # This depends on the internal state management - if result is not None and hasattr(result, "tool_calls") and result.tool_calls: - assert len(result.tool_calls) >= 0 - - -def test_streaming_with_content_before_tool_calls(minimax_tool_parser): - """Test streaming when there's content before tool calls.""" - # Reset streaming state - minimax_tool_parser.current_tool_name_sent = False - minimax_tool_parser.prev_tool_call_arr = [] - minimax_tool_parser.current_tool_id = -1 - minimax_tool_parser.streamed_args_for_tool = [] - - current_text = "I'll help you with that. " - - # When there's content before tool calls, it should be returned as content - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text="I'll help you", - current_text=current_text, - delta_text=" with that. ", - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - if result is not None and hasattr(result, "content"): - # Should contain some content - assert result.content is not None - - -def test_streaming_no_tool_calls(minimax_tool_parser): - """Test streaming when there are no tool calls.""" - current_text = "This is just regular text without any tool calls." - - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text="This is just regular text", - current_text=current_text, - delta_text=" without any tool calls.", - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - # Should return the delta text as content - assert result is not None - assert hasattr(result, "content") - assert result.content == " without any tool calls." - - -def test_streaming_with_thinking_tags(minimax_tool_parser): - """Test streaming with thinking tags that contain tool calls.""" - # Reset streaming state - minimax_tool_parser.current_tool_name_sent = False - minimax_tool_parser.prev_tool_call_arr = [] - minimax_tool_parser.current_tool_id = -1 - minimax_tool_parser.streamed_args_for_tool = [] - - current_text = """{"name": "ignored", "arguments": {}}{"name": "real_tool", "arguments": {"param": "value"}}""" - - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=current_text, - delta_text=current_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - # The preprocessing should remove tool calls from thinking tags - # and only process the real tool call - if result is not None and hasattr(result, "tool_calls") and result.tool_calls: - for tool_call in result.tool_calls: - assert tool_call.function.name != "ignored" - - -def test_extract_tool_calls_multiline_json_not_supported(minimax_tool_parser): - """Test that multiline JSON in tool calls is not currently supported.""" - model_output = """ -{ - "name": "get_current_weather", - "arguments": { - "city": "New York", - "state": "NY", - "unit": "celsius" - } -} -""" - - extracted_tool_calls = minimax_tool_parser.extract_tool_calls( - model_output, request=None - ) # type: ignore[arg-type] - - # Multiline JSON is currently not supported, should return no tools called - assert not extracted_tool_calls.tools_called - assert extracted_tool_calls.tool_calls == [] - assert extracted_tool_calls.content is None - - -def test_streaming_arguments_incremental_output(minimax_tool_parser): - """Test that streaming arguments are returned incrementally, not cumulatively.""" - # Reset streaming state - minimax_tool_parser.current_tool_name_sent = False - minimax_tool_parser.prev_tool_call_arr = [] - minimax_tool_parser.current_tool_id = -1 - minimax_tool_parser.streamed_args_for_tool = [] - - # Simulate progressive tool call building - stages = [ - # Stage 1: Function name complete - '\n{"name": "get_current_weather", "arguments": ', - # Stage 2: Arguments object starts with first key - '\n{"name": "get_current_weather", "arguments": {"city": ', - # Stage 3: First parameter value added - '\n{"name": "get_current_weather", "arguments": {"city": "Seattle"', - # Stage 4: Second parameter added - '\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA"', - # Stage 5: Third parameter added, arguments complete - '\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}', - # Stage 6: Tool calls closed - '\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n', - ] - - function_name_sent = False - previous_args_content = "" - - for i, current_text in enumerate(stages): - previous_text = stages[i - 1] if i > 0 else "" - delta_text = current_text[len(previous_text) :] if i > 0 else current_text - - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - print(f"Stage {i}: Current text: {repr(current_text)}") - print(f"Stage {i}: Delta text: {repr(delta_text)}") - - if result is not None and hasattr(result, "tool_calls") and result.tool_calls: - tool_call = result.tool_calls[0] - - # Check if function name is sent (should happen only once) - if tool_call.function and tool_call.function.name: - assert tool_call.function.name == "get_current_weather" - function_name_sent = True - print(f"Stage {i}: Function name sent: {tool_call.function.name}") - - # Check if arguments are sent incrementally - if tool_call.function and tool_call.function.arguments: - args_fragment = tool_call.function.arguments - print(f"Stage {i}: Got arguments fragment: {repr(args_fragment)}") - - # For incremental output, each fragment should be new content only - # The fragment should not contain all previous content - if i >= 2 and previous_args_content: # After we start getting arguments - # The new fragment should not be identical to or contain all previous content - assert args_fragment != previous_args_content, ( - f"Fragment should be incremental, not cumulative: {args_fragment}" - ) - - # If this is truly incremental, the fragment should be relatively small - # compared to the complete arguments so far - if len(args_fragment) > len(previous_args_content): - print( - "Warning: Fragment seems cumulative rather than incremental" - ) - - previous_args_content = args_fragment - - # Verify function name was sent at least once - assert function_name_sent, "Function name should have been sent" - - -def test_streaming_arguments_delta_only(minimax_tool_parser): - """Test that each streaming call returns only the delta (new part) of arguments.""" - # Reset streaming state - minimax_tool_parser.current_tool_name_sent = False - minimax_tool_parser.prev_tool_call_arr = [] - minimax_tool_parser.current_tool_id = -1 - minimax_tool_parser.streamed_args_for_tool = [] - - # Simulate two consecutive calls with growing arguments - call1_text = ( - '\n{"name": "test_tool", "arguments": {"param1": "value1"}}' - ) - call2_text = '\n{"name": "test_tool", "arguments": {"param1": "value1", "param2": "value2"}}' - - print(f"Call 1 text: {repr(call1_text)}") - print(f"Call 2 text: {repr(call2_text)}") - - # First call - should get the function name and initial arguments - result1 = minimax_tool_parser.extract_tool_calls_streaming( - previous_text="", - current_text=call1_text, - delta_text=call1_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - print(f"Result 1: {result1}") - if result1 and hasattr(result1, "tool_calls") and result1.tool_calls: - for i, tc in enumerate(result1.tool_calls): - print(f" Tool call {i}: {tc}") - - # Second call - should only get the delta (new part) of arguments - result2 = minimax_tool_parser.extract_tool_calls_streaming( - previous_text=call1_text, - current_text=call2_text, - delta_text=', "param2": "value2"}', - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - print(f"Result 2: {result2}") - if result2 and hasattr(result2, "tool_calls") and result2.tool_calls: - for i, tc in enumerate(result2.tool_calls): - print(f" Tool call {i}: {tc}") - - # Verify the second call only returns the delta - if result2 is not None and hasattr(result2, "tool_calls") and result2.tool_calls: - tool_call = result2.tool_calls[0] - if tool_call.function and tool_call.function.arguments: - args_delta = tool_call.function.arguments - print(f"Arguments delta from second call: {repr(args_delta)}") - - # Should only contain the new part, not the full arguments - # The delta should be something like ', "param2": "value2"}' or just '"param2": "value2"' - assert ( - ', "param2": "value2"}' in args_delta - or '"param2": "value2"' in args_delta - ), f"Expected delta containing param2, got: {args_delta}" - - # Should NOT contain the previous parameter data - assert '"param1": "value1"' not in args_delta, ( - f"Arguments delta should not contain previous data: {args_delta}" - ) - - # The delta should be relatively short (incremental, not cumulative) - expected_max_length = len(', "param2": "value2"}') + 10 # Some tolerance - assert len(args_delta) <= expected_max_length, ( - f"Delta seems too long (possibly cumulative): {args_delta}" - ) - - print("✓ Delta validation passed") - else: - print("No arguments in result2 tool call") - else: - print("No tool calls in result2 or result2 is None") - # This might be acceptable if no incremental update is needed - # But let's at least verify that result1 had some content - assert result1 is not None, "At least the first call should return something" - - -def test_streaming_openai_compatibility(minimax_tool_parser): - """Test that streaming behavior with buffering works correctly.""" - # Reset streaming state - minimax_tool_parser.current_tool_name_sent = False - minimax_tool_parser.prev_tool_call_arr = [] - minimax_tool_parser.current_tool_id = -1 - minimax_tool_parser.streamed_args_for_tool = [] - # Reset buffering state - minimax_tool_parser.pending_buffer = "" - minimax_tool_parser.in_thinking_tag = False - minimax_tool_parser.thinking_depth = 0 - - # Test scenario: simple buffering without complex tool call context - test_cases: list[dict[str, Any]] = [ - { - "stage": "Token: <", - "previous": "", - "current": "<", - "delta": "<", - "expected_content": None, # Should be buffered - }, - { - "stage": "Token: tool_calls>", - "previous": "<", - "current": "", - "delta": "tool_calls>", - "expected_content": None, # Complete tag, should not output - }, - { - "stage": "Regular content", - "previous": "Hello", - "current": "Hello world", - "delta": " world", - "expected_content": " world", # Normal content should pass through - }, - { - "stage": "Content with end tag start", - "previous": "Text", - "current": "Text content", - "delta": "calls>", - "expected_content": None, # Complete close tag, should not output - }, - ] - - for i, test_case in enumerate(test_cases): - print(f"\n--- Stage {i}: {test_case['stage']} ---") - print(f"Previous: {repr(test_case['previous'])}") - print(f"Current: {repr(test_case['current'])}") - print(f"Delta: {repr(test_case['delta'])}") - - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text=test_case["previous"], - current_text=test_case["current"], - delta_text=test_case["delta"], - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - print(f"Result: {result}") - - # Check expected content - if test_case["expected_content"] is None: - assert result is None or not getattr(result, "content", None), ( - f"Stage {i}: Expected no content, got {result}" - ) - print("✓ No content output as expected") - else: - assert result is not None and hasattr(result, "content"), ( - f"Stage {i}: Expected content, got {result}" - ) - assert result.content == test_case["expected_content"], ( - f"Stage {i}: Expected content {test_case['expected_content']}, got {result.content}" - ) - print(f"✓ Content matches: {repr(result.content)}") - - print("✓ Streaming test with buffering completed successfully") - - -def test_streaming_thinking_tag_buffering(minimax_tool_parser): - """Test that tool calls within thinking tags are properly handled during streaming.""" - # Reset streaming state - minimax_tool_parser.current_tool_name_sent = False - minimax_tool_parser.prev_tool_call_arr = [] - minimax_tool_parser.current_tool_id = -1 - minimax_tool_parser.streamed_args_for_tool = [] - # Reset buffering state - minimax_tool_parser.pending_buffer = "" - minimax_tool_parser.in_thinking_tag = False - minimax_tool_parser.thinking_depth = 0 - - # Test scenario: tool calls within thinking tags should be ignored - test_cases: list[dict[str, Any]] = [ - { - "stage": "Start thinking", - "previous": "", - "current": "I need to use a tool. ", - "delta": "I need to use a tool. ", - "expected_content": "I need to use a tool. ", # Should pass through as content - }, - { - "stage": "Tool call in thinking", - "previous": "I need to use a tool. ", - "current": 'I need to use a tool. \n{"name": "ignored_tool", "arguments": {"param": "value"}}\n', - "delta": '\n{"name": "ignored_tool", "arguments": {"param": "value"}}\n', - "expected_content": '\n{"name": "ignored_tool", "arguments": {"param": "value"}}\n', # should be preserved in thinking tags - }, - { - "stage": "Real tool call after thinking", - "previous": 'I need to use a tool. \n{"name": "ignored_tool", "arguments": {"param": "value"}}\n', - "current": 'I need to use a tool. \n{"name": "ignored_tool", "arguments": {"param": "value"}}\n\n', - "delta": "\n", - "expected_content": "\n", # Should output '\n' and suppress - }, - ] - - for i, test_case in enumerate(test_cases): - print(f"\n--- Stage {i}: {test_case['stage']} ---") - print(f"Previous: {repr(test_case['previous'])}") - print(f"Current: {repr(test_case['current'])}") - print(f"Delta: {repr(test_case['delta'])}") - - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text=test_case["previous"], - current_text=test_case["current"], - delta_text=test_case["delta"], - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - print(f"Result: {result}") - - # Check expected content - if "expected_content" in test_case: - if test_case["expected_content"] is None: - assert result is None or not getattr(result, "content", None), ( - f"Stage {i}: Expected no content, got {result}" - ) - else: - assert result is not None and hasattr(result, "content"), ( - f"Stage {i}: Expected content, got {result}" - ) - assert result.content == test_case["expected_content"], ( - f"Stage {i}: Expected content {test_case['expected_content']}, got {result.content}" - ) - print(f"✓ Content matches: {repr(result.content)}") - - # Check tool calls - if test_case.get("expected_tool_call"): - assert ( - result is not None - and hasattr(result, "tool_calls") - and result.tool_calls - ), f"Stage {i}: Expected tool call, got {result}" - - tool_call = result.tool_calls[0] - assert tool_call.function.name == "real_tool", ( - f"Expected real_tool, got {tool_call.function.name}" - ) - print(f"✓ Real tool call detected: {tool_call.function.name}") - - print("✓ Thinking tag buffering test completed successfully") - - -def reset_streaming_state(minimax_tool_parser): - """Helper function to properly reset the streaming state for MinimaxToolParser.""" - # Reset minimax-specific state - minimax_tool_parser._reset_streaming_state() - - # Reset base class state (these should still be reset for compatibility) - minimax_tool_parser.prev_tool_call_arr = [] - minimax_tool_parser.current_tool_id = -1 - minimax_tool_parser.current_tool_name_sent = False - minimax_tool_parser.streamed_args_for_tool = [] - - -def test_streaming_complex_scenario_with_multiple_tools(minimax_tool_parser): - """Test complex streaming scenario: tools inside tags and multiple tool calls in one group.""" - # Reset streaming state - reset_streaming_state(minimax_tool_parser) - - # Complex scenario: tools inside thinking tags and multiple tools in one group - test_stages: list[dict[str, Any]] = [ - { - "stage": "Initial content", - "previous": "", - "current": "Let me help you with this task.", - "delta": "Let me help you with this task.", - "expected_content": "Let me help you with this task.", - "expected_tool_calls": 0, - }, - { - "stage": "Start thinking tag", - "previous": "Let me help you with this task.", - "current": "Let me help you with this task.I need to analyze this situation first.", - "delta": "I need to analyze this situation first.", - "expected_content": "I need to analyze this situation first.", - "expected_tool_calls": 0, - }, - { - "stage": "Tool call inside thinking tag starts", - "previous": "Let me help you with this task.I need to analyze this situation first.", - "current": "Let me help you with this task.I need to analyze this situation first.", - "delta": "", - "expected_content": "", # Inside thinking tags, tool tags should be preserved as content - "expected_tool_calls": 0, - }, - { - "stage": "Complete tool call inside thinking tag", - "previous": "Let me help you with this task.I need to analyze this situation first.", - "current": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n', - "delta": '\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n', - "expected_content": '\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n', - "expected_tool_calls": 0, # Tools inside thinking tags should be ignored - }, - { - "stage": "End thinking tag", - "previous": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n', - "current": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n', - "delta": "", - "expected_content": "", - "expected_tool_calls": 0, - }, - { - "stage": "Multiple tools group starts", - "previous": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n', - "current": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n\nNow I need to get weather information and calculate area.', - "delta": "\nNow I need to get weather information and calculate area.", - "expected_content": "\nNow I need to get weather information and calculate area.", # should be filtered - "expected_tool_calls": 0, - }, - { - "stage": "First tool in group", - "previous": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n\nNow I need to get weather information and calculate area.', - "current": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n\nNow I need to get weather information and calculate area.\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}', - "delta": '\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}', - "expected_content": None, # No content should be output when tool call is in progress - "expected_tool_calls": 1, - "expected_tool_name": "get_current_weather", - }, - { - "stage": "Second tool in group", - "previous": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n\nNow I need to get weather information and calculate area.\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}', - "current": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n\nNow I need to get weather information and calculate area.\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}', - "delta": '\n{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}', - "expected_content": None, - "expected_tool_calls": 1, - "expected_tool_name": "calculate_area", - }, - { - "stage": "Complete tool calls group", - "previous": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n\nNow I need to get weather information and calculate area.\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}', - "current": 'Let me help you with this task.I need to analyze this situation first.\n{"name": "internal_analysis", "arguments": {"query": "analyze situation"}}\n\nNow I need to get weather information and calculate area.\n{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}}\n{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}}', - "delta": "", - "expected_content": None, - "expected_tool_calls": 0, - }, - ] - - tool_calls_count = 0 - - for i, test_case in enumerate(test_stages): - print(f"\n--- Stage {i}: {test_case['stage']} ---") - print( - f"Previous: {repr(test_case['previous'][:100])}{'...' if len(test_case['previous']) > 100 else ''}" - ) - print(f"Current: {repr(test_case['current'][-100:])}") - print(f"Delta: {repr(test_case['delta'])}") - - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text=test_case["previous"], - current_text=test_case["current"], - delta_text=test_case["delta"], - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - print(f"Result: {result}") - - # Check expected content - if test_case["expected_content"] is None: - assert result is None or not getattr(result, "content", None), ( - f"Stage {i}: Expected no content output, got {result}" - ) - print("✓ No content output as expected") - else: - assert result is not None and hasattr(result, "content"), ( - f"Stage {i}: Expected content output, got {result}" - ) - assert result.content == test_case["expected_content"], ( - f"Stage {i}: Expected content {repr(test_case['expected_content'])}, got {repr(result.content)}" - ) - print(f"✓ Content matches: {repr(result.content)}") - - # Check tool calls - expected_tool_calls = test_case["expected_tool_calls"] - actual_tool_calls = ( - len(result.tool_calls) - if result and hasattr(result, "tool_calls") and result.tool_calls - else 0 - ) - - if expected_tool_calls > 0: - assert actual_tool_calls >= expected_tool_calls, ( - f"Stage {i}: Expected at least {expected_tool_calls} tool calls, got {actual_tool_calls}" - ) - - if "expected_tool_name" in test_case: - # Find the tool call with the expected name - found_tool_call = None - for tool_call in result.tool_calls: - if tool_call.function.name == test_case["expected_tool_name"]: - found_tool_call = tool_call - break - - assert found_tool_call is not None, ( - f"Stage {i}: Expected tool name {test_case['expected_tool_name']} not found in tool calls: {[tc.function.name for tc in result.tool_calls]}" - ) - print(f"✓ Tool call correct: {found_tool_call.function.name}") - - # Ensure tools inside thinking tags are not called - assert found_tool_call.function.name != "internal_analysis", ( - f"Stage {i}: Tool 'internal_analysis' inside thinking tags should not be called" - ) - - tool_calls_count += actual_tool_calls - print(f"✓ Detected {actual_tool_calls} tool calls") - else: - assert actual_tool_calls == 0, ( - f"Stage {i}: Expected no tool calls, got {actual_tool_calls}" - ) - - # Verify overall results - print("\n=== Test Summary ===") - print(f"Total tool calls count: {tool_calls_count}") - assert tool_calls_count >= 2, ( - f"Expected at least 2 valid tool calls (outside thinking tags), but got {tool_calls_count}" - ) - - print("✓ Complex streaming test completed:") - print(" - ✓ Tools inside thinking tags correctly ignored") - print(" - ✓ Two tool groups outside thinking tags correctly parsed") - print(" - ✓ Content and tool call streaming correctly handled") - print(" - ✓ Buffering mechanism works correctly") - - -def test_streaming_character_by_character_output(minimax_tool_parser): - """Test character-by-character streaming output to simulate real streaming scenarios.""" - # Reset streaming state - reset_streaming_state(minimax_tool_parser) - - # Complete text that will be streamed character by character - complete_text = """I'll help you with the weather analysis. Let me think about this. -{"name": "internal_analysis", "arguments": {"type": "thinking"}} -This tool should be ignored. - -Now I'll get the weather information for you. -{"name": "get_current_weather", "arguments": {"city": "Seattle", "state": "WA", "unit": "celsius"}} -{"name": "calculate_area", "arguments": {"shape": "rectangle", "dimensions": {"width": 10, "height": 5}}} -Here are the results.""" - - print("\n=== Starting character-by-character streaming test ===") - print(f"Complete text length: {len(complete_text)} characters") - - # Track the streaming results - content_fragments = [] - tool_calls_detected = [] - - # Stream character by character - for i in range(1, len(complete_text) + 1): - current_text = complete_text[:i] - previous_text = complete_text[: i - 1] if i > 1 else "" - delta_text = complete_text[i - 1 : i] - - # Show progress every 50 characters - if i % 50 == 0 or i == len(complete_text): - print(f"Progress: {i}/{len(complete_text)} characters") - - # Call the streaming parser - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - # Collect results - if result is not None: - if hasattr(result, "content") and result.content: - content_fragments.append(result.content) - # Log important content fragments - if any( - keyword in result.content - for keyword in [ - "", - "", - "", - "", - ] - ): - print(f" Char {i}: Content fragment: {repr(result.content)}") - - if hasattr(result, "tool_calls") and result.tool_calls: - for tool_call in result.tool_calls: - tool_info = { - "character_position": i, - "function_name": tool_call.function.name - if tool_call.function - else None, - "arguments": tool_call.function.arguments - if tool_call.function - else None, - } - tool_calls_detected.append(tool_info) - print(f" Char {i}: Tool call detected: {tool_call.function.name}") - if tool_call.function.arguments: - print(f" Arguments: {repr(tool_call.function.arguments)}") - - # Verify results - print("\n=== Streaming Test Results ===") - print(f"Total content fragments: {len(content_fragments)}") - print(f"Total tool calls detected: {len(tool_calls_detected)}") - - # Reconstruct content from fragments - reconstructed_content = "".join(content_fragments) - print(f"Reconstructed content length: {len(reconstructed_content)}") - - # Verify thinking tags content is preserved - assert "" in reconstructed_content, ( - "Opening thinking tag should be preserved in content" - ) - assert "" in reconstructed_content, ( - "Closing thinking tag should be preserved in content" - ) - - # Verify that tool calls inside thinking tags are NOT extracted as actual tool calls - thinking_tool_calls = [ - tc for tc in tool_calls_detected if tc["function_name"] == "internal_analysis" - ] - assert len(thinking_tool_calls) == 0, ( - f"Tool calls inside thinking tags should be ignored, but found: {thinking_tool_calls}" - ) - - # Verify that real tool calls outside thinking tags ARE extracted - weather_tool_calls = [ - tc for tc in tool_calls_detected if tc["function_name"] == "get_current_weather" - ] - area_tool_calls = [ - tc for tc in tool_calls_detected if tc["function_name"] == "calculate_area" - ] - print(tool_calls_detected) - assert len(weather_tool_calls) > 0, ( - "get_current_weather tool call should be detected" - ) - assert len(area_tool_calls) > 0, "calculate_area tool call should be detected" - - # Verify tool call arguments are properly streamed - weather_args_found = any( - tc["arguments"] for tc in weather_tool_calls if tc["arguments"] - ) - area_args_found = any(tc["arguments"] for tc in area_tool_calls if tc["arguments"]) - - print(f"Weather tool call with arguments: {weather_args_found}") - print(f"Area tool call with arguments: {area_args_found}") - - # Verify content before and after tool calls - assert "I'll help you with the weather analysis." in reconstructed_content, ( - "Initial content should be preserved" - ) - assert "Here are the results." in reconstructed_content, ( - "Final content should be preserved" - ) - - # Verify that and tags are not included in the final content - # (they should be filtered out when not inside thinking tags) - content_outside_thinking = reconstructed_content - # Remove thinking tag content to check content outside - if "" in content_outside_thinking and "" in content_outside_thinking: - start_think = content_outside_thinking.find("") - end_think = content_outside_thinking.find("") + len("") - content_outside_thinking = ( - content_outside_thinking[:start_think] - + content_outside_thinking[end_think:] - ) - - # Outside thinking tags, tool_calls tags should be filtered - tool_calls_in_content = content_outside_thinking.count("") - assert tool_calls_in_content == 0, ( - f" tags should be filtered from content outside thinking tags, but found {tool_calls_in_content}" - ) - - print("\n=== Character-by-character streaming test completed successfully ===") - print("✓ Tool calls inside thinking tags correctly ignored") - print("✓ Tool calls outside thinking tags correctly detected") - print("✓ Content properly streamed and reconstructed") - print("✓ Tool call tags properly filtered from content") - print("✓ Character-level streaming works correctly") - - -def test_streaming_character_by_character_simple_tool_call(minimax_tool_parser): - """Test character-by-character streaming for a simple tool call scenario.""" - # Reset streaming state - reset_streaming_state(minimax_tool_parser) - - # Simple tool call text - simple_text = 'Let me check the weather. \n{"name": "get_weather", "arguments": {"city": "NYC"}}\n' - - print("\n=== Simple character-by-character test ===") - print(f"Text: {repr(simple_text)}") - - content_parts = [] - tool_name_sent = False - tool_args_sent = False - - for i in range(1, len(simple_text) + 1): - current_text = simple_text[:i] - previous_text = simple_text[: i - 1] if i > 1 else "" - delta_text = simple_text[i - 1 : i] - - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - if result: - if hasattr(result, "content") and result.content: - content_parts.append(result.content) - print( - f" Char {i} ({repr(delta_text)}): Content: {repr(result.content)}" - ) - - if hasattr(result, "tool_calls") and result.tool_calls: - for tool_call in result.tool_calls: - if tool_call.function and tool_call.function.name: - tool_name_sent = True - print(f" Char {i}: Tool name: {tool_call.function.name}") - if tool_call.function and tool_call.function.arguments: - tool_args_sent = True - print( - f" Char {i}: Tool args: {repr(tool_call.function.arguments)}" - ) - - # Verify basic expectations - reconstructed_content = "".join(content_parts) - print(f"Final reconstructed content: {repr(reconstructed_content)}") - - assert tool_name_sent, "Tool name should be sent during streaming" - assert tool_args_sent, "Tool arguments should be sent during streaming" - assert "Let me check the weather." in reconstructed_content, ( - "Initial content should be preserved" - ) - - print("✓ Simple character-by-character test passed") - - -def test_streaming_character_by_character_with_buffering(minimax_tool_parser): - """Test character-by-character streaming with edge cases that trigger buffering.""" - # Reset streaming state - reset_streaming_state(minimax_tool_parser) - - # Text that includes potential buffering scenarios - buffering_text = 'Hello world\n{"name": "test"}\ndone' - - print("\n=== Buffering character-by-character test ===") - print(f"Text: {repr(buffering_text)}") - - all_content = [] - - for i in range(1, len(buffering_text) + 1): - current_text = buffering_text[:i] - previous_text = buffering_text[: i - 1] if i > 1 else "" - delta_text = buffering_text[i - 1 : i] - - result = minimax_tool_parser.extract_tool_calls_streaming( - previous_text=previous_text, - current_text=current_text, - delta_text=delta_text, - previous_token_ids=[], - current_token_ids=[], - delta_token_ids=[], - request=None, - ) - - if result and hasattr(result, "content") and result.content: - all_content.append(result.content) - print(f" Char {i} ({repr(delta_text)}): {repr(result.content)}") - - final_content = "".join(all_content) - print(f"Final content: {repr(final_content)}") - - # The parser should handle the edge case where appears before - assert "Hello" in final_content, "Initial 'Hello' should be preserved" - assert "world" in final_content, ( - "Content after false closing tag should be preserved" - ) - assert "done" in final_content, "Final content should be preserved" - - print("✓ Buffering character-by-character test passed") diff --git a/tests/v1/attention/test_attention_backends_selection.py b/tests/v1/attention/test_attention_backends_selection.py index e3d2e9dc457..8486d216a12 100644 --- a/tests/v1/attention/test_attention_backends_selection.py +++ b/tests/v1/attention/test_attention_backends_selection.py @@ -6,10 +6,12 @@ from types import SimpleNamespace import pytest +from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import ( + MiniMaxText01LinearAttention, +) from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2 from vllm.model_executor.layers.mamba.short_conv import ShortConv -from vllm.model_executor.models.minimax_text_01 import MiniMaxText01LinearAttention from vllm.v1.attention.backends.linear_attn import LinearAttentionBackend from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionBackend from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionBackend diff --git a/vllm/model_executor/models/minimax_text_01.py b/vllm/model_executor/models/minimax_text_01.py deleted file mode 100644 index 890dbe590ae..00000000000 --- a/vllm/model_executor/models/minimax_text_01.py +++ /dev/null @@ -1,1000 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Inference-only MiniMaxText01 model.""" - -from collections.abc import Iterable -from itertools import islice -from typing import TYPE_CHECKING - -if TYPE_CHECKING: - pass - -import regex as re -import torch -from torch import nn -from transformers import MiniMaxConfig - -from vllm.compilation.decorators import support_torch_compile -from vllm.config import CacheConfig, VllmConfig -from vllm.distributed.parallel_state import ( - get_pp_group, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) -from vllm.forward_context import get_forward_context -from vllm.model_executor.layers.activation import SiluAndMul -from vllm.model_executor.layers.attention import Attention -from vllm.model_executor.layers.fused_moe import ( - FusedMoE, -) -from vllm.model_executor.layers.layernorm import RMSNorm -from vllm.model_executor.layers.linear import ( - MergedColumnParallelLinear, - QKVParallelLinear, - ReplicatedLinear, - RowParallelLinear, -) -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import ( - MiniMaxText01LinearAttention, -) -from vllm.model_executor.layers.mamba.mamba_utils import ( - MambaStateCopyFunc, - MambaStateCopyFuncCalculator, - MambaStateDtypeCalculator, - MambaStateShapeCalculator, -) -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.model_executor.layers.rotary_embedding import get_rope -from vllm.model_executor.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) -from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm.model_executor.models.utils import maybe_prefix -from vllm.sequence import IntermediateTensors -from vllm.v1.attention.backend import AttentionMetadata - -from .interfaces import HasInnerState, IsHybrid -from .utils import ( - AutoWeightsLoader, - PPMissingLayer, - is_pp_missing_parameter, - make_layers, -) - - -def replace_weight_name( - name: str, key: str = None, to: str = None, count: int = None, prefix: str = None -) -> str: - name = name.replace(key, to) if count is None else name.replace(key, to, count) - return name - - -def weight_loader_with_alias(alias: str): - def wrapper(func: callable): - def inner_func( - param: torch.Tensor, - loaded_weight: torch.Tensor, - *args, - prefix: str = None, - **kwargs, - ): - value = func(param, loaded_weight, *args, **kwargs) - return value - - return inner_func - - return wrapper - - -class MiniMaxText01MLP(nn.Module): - def __init__( - self, - hidden_size: int, - intermediate_size: int, - quant_config: QuantizationConfig | None = None, - layer_idx: int = None, - prefix: str = "mlp", - ) -> None: - super().__init__() - self.layer_idx = layer_idx - - self.gate_up_proj = MergedColumnParallelLinear( - hidden_size, - [intermediate_size] * 2, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.gate_up_proj", - ) - self.down_proj = RowParallelLinear( - intermediate_size, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.down_proj", - ) - self.act_fn = SiluAndMul() - return - - def forward(self, x: torch.Tensor) -> torch.Tensor: - gate_up, _ = self.gate_up_proj(x) - x = self.act_fn(gate_up) - x, _ = self.down_proj(x) - return x - - -class MiniMaxText01MoE(nn.Module): - def __init__( - self, - num_experts: int, - top_k: int, - hidden_size: int, - intermediate_size: int, - params_dtype: torch.dtype | None = None, - layer_idx: int = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "moe", - ) -> None: - super().__init__() - - self.layer_idx = layer_idx - self.tp_size = get_tensor_model_parallel_world_size() - self.num_total_experts = num_experts - self.top_k = top_k - self.hidden_size = hidden_size - self.intermediate_size = intermediate_size // self.tp_size - self.quant_config = quant_config - - if params_dtype is None: - params_dtype = torch.get_default_dtype() - self.params_dtype = params_dtype - - self.gate = ReplicatedLinear( - self.hidden_size, - self.num_total_experts, - bias=False, - params_dtype=torch.float32, - quant_config=None, - prefix=f"{prefix}.gate", - ) - self.gate.weight.weight_loader = MiniMaxText01MoE.gate_weight_loader - - self.experts = FusedMoE( - num_experts=self.num_total_experts, - top_k=self.top_k, - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size * self.tp_size, - params_dtype=self.params_dtype, - renormalize=True, - quant_config=self.quant_config, - tp_size=self.tp_size, - prefix=f"{prefix}.experts", - ) - return - - @staticmethod - def gate_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None: - assert param.size() == loaded_weight.size() - param.data.copy_(loaded_weight.to(torch.float32)) - return - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - num_tokens, hidden_size = hidden_states.shape - hidden_states = hidden_states.view(-1, self.hidden_size) - router_logits_fp32, _ = self.gate(hidden_states.to(torch.float32)) - final_hidden_states = self.experts( - hidden_states, router_logits_fp32.to(hidden_states.dtype) - ) - final_hidden = final_hidden_states.view(num_tokens, hidden_size) - return final_hidden - - -class MiniMaxText01Attention(nn.Module): - def __init__( - self, - hidden_size: int, - num_heads: int, - head_dim: int, - num_kv_heads: int, - max_position: int = 4096 * 32, - rope_parameters: dict | None = None, - sliding_window: int | None = None, - quant_config: QuantizationConfig | None = None, - layer_idx: int = None, - cache_config: CacheConfig | None = None, - prefix: str = "mha", - ) -> None: - super().__init__() - self.layer_idx = layer_idx - - self.hidden_size = hidden_size - tp_size = get_tensor_model_parallel_world_size() - self.total_num_heads = num_heads - assert self.total_num_heads % tp_size == 0 - self.num_heads = self.total_num_heads // tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= tp_size: - assert self.total_num_kv_heads % tp_size == 0 - else: - assert tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) - self.head_dim = head_dim - - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - self.scaling = self.head_dim**-0.5 - self.sliding_window = sliding_window - self.prefix = prefix - - self.qkv_proj = QKVParallelLinear( - hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.qkv_proj", - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=False, - quant_config=quant_config, - prefix=f"{prefix}.o_proj", - ) - self.attn = Attention( - self.num_heads, - self.head_dim, - self.scaling, - num_kv_heads=self.num_kv_heads, - cache_config=cache_config, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) - self.rotary_emb = get_rope( - head_size=self.head_dim, - max_position=max_position, - rope_parameters=rope_parameters, - is_neox_style=True, - dtype=torch.float32, - ) - return - - def forward( - self, - hidden_states: torch.Tensor, - output: torch.Tensor, - positions: torch.Tensor, - **kwargs, - ) -> None: - qkv, _ = self.qkv_proj(hidden_states) - q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) - attn_output = self.attn(q, k, v) - output[:], _ = self.o_proj(attn_output) - - -class MiniMaxText01DecoderLayer(nn.Module): - def __init__( - self, - config: MiniMaxConfig, - vllm_config: VllmConfig, - expert_num: int = 1, - layer_id: int = None, - linear_layer_id: int | None = None, - prefix: str = "decoder", - ) -> None: - self._ilayer = layer_id - self._irank = get_tensor_model_parallel_rank() - self.prefix = prefix - super().__init__() - - self.hidden_size = config.hidden_size - self.expert_num = expert_num - - head_dim = getattr(config, "head_dim", None) - if head_dim is None: - head_dim = config.hidden_size // config.num_attention_heads - rotary_dim = getattr(config, "rotary_dim", head_dim) - config.rope_parameters["partial_rotary_factor"] = rotary_dim / head_dim - if hasattr(config, "max_model_len") and isinstance(config.max_model_len, int): - max_position_embeddings = min( - config.max_position_embeddings, config.max_model_len - ) - if config.attention_type == 0: - self.self_attn = MiniMaxText01LinearAttention( - config, - vllm_config, - prefix=prefix, - ) - elif config.attention_type == 1: - self.self_attn = MiniMaxText01Attention( - hidden_size=self.hidden_size, - num_heads=config.num_attention_heads, - head_dim=head_dim, - num_kv_heads=config.num_key_value_heads, - max_position=max_position_embeddings, - rope_parameters=config.rope_parameters, - sliding_window=config.sliding_window, - quant_config=vllm_config.quant_config, - layer_idx=self._ilayer, - cache_config=vllm_config.cache_config, - prefix=prefix, - ) - else: - raise ValueError( - f"Unsupported attention_type {self.config.attention_type}: " - f"should be 0 (linear) or 1 (full)." - ) - - if expert_num == 1: - self.mlp = MiniMaxText01MLP( - hidden_size=self.hidden_size, - intermediate_size=config.intermediate_size, - quant_config=vllm_config.quant_config, - layer_idx=self._ilayer, - prefix=prefix, - ) - else: - self.block_sparse_moe = MiniMaxText01MoE( - num_experts=expert_num, - top_k=config.num_experts_per_tok, - hidden_size=config.hidden_size, - intermediate_size=config.intermediate_size, - layer_idx=self._ilayer, - quant_config=vllm_config.quant_config, - prefix=prefix, - ) - - self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = RMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - if config.attention_type == 0: - self.layernorm_attention_alpha = getattr( - config, - "layernorm_linear_attention_alpha", - getattr(config, "linear_attn_alpha_factor", 1), - ) - self.layernorm_attention_beta = getattr( - config, - "layernorm_linear_attention_beta", - getattr(config, "linear_attn_beta_factor", 1), - ) - else: - self.layernorm_attention_alpha = getattr( - config, - "layernorm_full_attention_alpha", - getattr(config, "full_attn_alpha_factor", 1), - ) - self.layernorm_attention_beta = getattr( - config, - "layernorm_full_attention_beta", - getattr(config, "full_attn_beta_factor", 1), - ) - self.layernorm_mlp_alpha = getattr( - config, "layernorm_mlp_alpha", getattr(config, "mlp_alpha_factor", 1) - ) - self.layernorm_mlp_beta = getattr( - config, "layernorm_mlp_beta", getattr(config, "mlp_beta_factor", 1) - ) - self.postnorm = getattr(config, "postnorm", False) - self.shared_moe = False - - shared_intermediate = getattr(config, "shared_intermediate_size", 0) - if isinstance(shared_intermediate, list): - shared_intermediate = ( - shared_intermediate[layer_id] - if layer_id < len(shared_intermediate) - else 0 - ) - if shared_intermediate > 0: - self.shared_moe = True - self.shared_mlp = MiniMaxText01MLP( - hidden_size=self.hidden_size, - intermediate_size=shared_intermediate, - quant_config=vllm_config.quant_config, - layer_idx=self._ilayer, - prefix=prefix, - ) - self.coefficient = ReplicatedLinear( - self.hidden_size, - 1, - bias=False, - quant_config=vllm_config.quant_config, - params_dtype=torch.float32, - ) - self.coefficient.weight.weight_loader = self.shared_moe_coefficient_loader - self.shared_moe_mode = getattr(config, "shared_moe_mode", "softmax") - return - - def forward( - self, - hidden_states: torch.Tensor, - positions: torch.Tensor, - attn_metadata: AttentionMetadata, - residual: torch.Tensor | None, - is_warmup: bool = False, - **kwargs, - ) -> tuple[torch.Tensor, torch.Tensor]: - layernorm_input = hidden_states - layernorm_output = self.input_layernorm(layernorm_input) - residual = layernorm_output if self.postnorm else layernorm_input - self_attention_output = torch.empty_like(layernorm_output) - self.self_attn( - hidden_states=layernorm_output, - output=self_attention_output, - positions=positions, - ) - - residual = residual * self.layernorm_attention_alpha - self_attention_output = self_attention_output * self.layernorm_attention_beta - - layernorm_input = residual + self_attention_output - layernorm_output = self.post_attention_layernorm(layernorm_input) - residual = layernorm_output if self.postnorm else layernorm_input - - if self.expert_num == 1: - hidden_states = self.mlp(layernorm_output) - else: - moe_layernorm_output = layernorm_output.clone() - moe_hidden_states = self.block_sparse_moe(moe_layernorm_output) - if self.shared_moe: - before_moe_dtype = layernorm_output.dtype - moe_hidden_fp32 = moe_hidden_states.to(torch.float32) - output_mlp = self.shared_mlp(layernorm_output).to(torch.float32) - - coef, _ = self.coefficient(layernorm_output.to(torch.float32)) - - if self.shared_moe_mode == "softmax": - coef = torch.nn.functional.softmax(coef, dim=-1) - hidden_states = moe_hidden_fp32 * (1 - coef) + output_mlp * coef - elif self.shared_moe_mode == "sigmoid": - coef = torch.nn.functional.sigmoid(coef) - hidden_states = moe_hidden_fp32 * (1 - coef) + output_mlp * coef - - hidden_states = hidden_states.to(before_moe_dtype) - else: - hidden_states = moe_hidden_states - - residual = residual * self.layernorm_mlp_alpha - hidden_states = hidden_states * self.layernorm_mlp_beta - - hidden_states = residual + hidden_states - - return hidden_states, None - - @staticmethod - def shared_moe_coefficient_loader( - param: torch.Tensor, loaded_weight: torch.Tensor - ) -> None: - assert param.size() == loaded_weight.size() - - param.data.copy_(loaded_weight.to(torch.float32)) - return - - -@support_torch_compile -class MiniMaxText01Model(nn.Module): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config: MiniMaxConfig = vllm_config.model_config.hf_config - scheduler_config = vllm_config.scheduler_config - self.config = config - self.CONCAT_FFN = True - - self.vocab_size = config.vocab_size - - self.decoder_attention_types = getattr( - config, "attn_type_list", False - ) or getattr(config, "decoder_attention_types", False) - # The HF format uses "layer_types" instead of "attn_type_list" - # where "linear_attention" is 0 and "full_attention" is 1 - if not self.decoder_attention_types and hasattr(config, "layer_types"): - self.decoder_attention_types = [] - for layer_type in config.layer_types: - if layer_type == "linear_attention": - self.decoder_attention_types.append(0) - elif layer_type == "full_attention": - self.decoder_attention_types.append(1) - else: - raise ValueError(f"Unsupported layer type: {layer_type}") - # Default to full attention - if not self.decoder_attention_types: - self.decoder_attention_types = [1] * config.num_hidden_layers - self.num_layers = config.num_hidden_layers - - self._layer_barrier = False - if get_pp_group().is_first_rank: - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - org_num_embeddings=self.vocab_size, - ) - else: - self.embed_tokens = PPMissingLayer() - - def layer_fn(prefix): - layer_idx = int(prefix.split(".")[-1]) - layer_config = config - layer_config.attention_type = self.decoder_attention_types[layer_idx] - layer_config.layer_idx = layer_idx - - decoder_kwargs = { - "layer_id": layer_idx, - "vllm_config": vllm_config, - } - - if layer_config.attention_type == 0: - decoder_kwargs["linear_layer_id"] = sum( - 1 for i in range(layer_idx) if self.decoder_attention_types[i] == 0 - ) - else: - decoder_kwargs["linear_layer_id"] = None - - if hasattr(config, "num_local_experts") and isinstance( - config.num_local_experts, list - ): - decoder_kwargs["expert_num"] = config.num_local_experts[layer_idx] - elif hasattr(config, "num_local_experts") and isinstance( - config.num_local_experts, int - ): - decoder_kwargs["expert_num"] = config.num_local_experts - else: - decoder_kwargs["expert_num"] = 1 - - return MiniMaxText01DecoderLayer( - layer_config, **decoder_kwargs, prefix=prefix - ) - - self.start_layer, self.end_layer, self.layers = make_layers( - config.num_hidden_layers, layer_fn, prefix=f"{prefix}.layers" - ) - - linear_layer_nums = sum( - 1 - for i in range(config.num_hidden_layers) - if self.decoder_attention_types[i] == 0 - ) - max_slots_number = scheduler_config.max_num_seqs - self.cache_shape = ( - linear_layer_nums, - max_slots_number, - config.num_attention_heads // get_tensor_model_parallel_world_size(), - config.head_dim, - config.head_dim, - ) - _dummy = torch.zeros(1) - self._dtype = _dummy.dtype - del _dummy - - norm_kwargs = {} - if hasattr(config, "rms_norm_eps"): - norm_kwargs["eps"] = config.rms_norm_eps - if get_pp_group().is_last_rank: - self.norm = RMSNorm(config.hidden_size, **norm_kwargs) - else: - self.norm = PPMissingLayer() - self.embed_scale = 1.0 - return - - def _clear_prefill_cache( - self, attn_metadata, minimax_cache_tensors: torch.Tensor, **kwargs - ): - seq_to_slot_maps = {} - seq_id_map = sum(list(kwargs["request_ids_to_seq_ids"].values()), []) - for _, seq_to_slot_map in self.minimax_cache.cache_indices_mapping.items(): - seq_to_slot_maps.update(seq_to_slot_map) - - slots_to_clear = [] - for _prefill_id in range(getattr(attn_metadata, "num_prefills", 0)): - if _prefill_id >= len(seq_id_map): - break - seq_id = seq_id_map[_prefill_id] - if ( - attn_metadata.context_lens_tensor[_prefill_id] == 0 - and seq_id in seq_to_slot_maps - ): - slots_to_clear.append(seq_to_slot_maps[seq_id]) - - if slots_to_clear: - slots_tensor = torch.tensor( - slots_to_clear, device=minimax_cache_tensors.device, dtype=torch.long - ) - minimax_cache_tensors[:, slots_tensor, ...] = 0 - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - def which_layer(name: str) -> int: - if "layers" in name: - after_layer = name.split("layers")[-1] - return int(after_layer.split(".")[1]) - return None - - def is_linear_attn_layer(layer_idx: int) -> bool: - if layer_idx is None or layer_idx >= len(self.decoder_attention_types): - return False - return self.decoder_attention_types[layer_idx] == 0 - - def is_moe_weight(name: str) -> bool: - return "block_sparse_moe" in name and not name.endswith(".bias") - - def get_expert_id(param_name): - pattern = r"layers\.\d+\.block_sparse_moe\.experts\.(\d+)\." - match = re.search(pattern, param_name) - if match: - return match.group(1) - return None - - def load_sparse_moe_weight( - name: str, loaded_weight: torch.Tensor, self - ) -> None: - if isinstance(self.config.num_local_experts, list): - expert_params_mapping = [ - ( - "w13_weight" if weight_name in ["w1", "w3"] else "w2_weight", - f"experts.{expert_id}.{weight_name}.weight", - expert_id, - ) - for expert_id in range(max(self.config.num_local_experts)) - for weight_name in ["w1", "w2", "w3"] - ] - else: - expert_params_mapping = [ - ( - "w13_scale" if weight_name in ["w1", "w3"] else "w2_scale", - f"{expert_id}.{weight_name}.weight_scale", - expert_id, - weight_name, - ) - for expert_id in range(self.config.num_local_experts) - for weight_name in ["w1", "w2", "w3"] - ] + [ - ( - "w13_weight" if weight_name in ["w1", "w3"] else "w2_weight", - f"{expert_id}.{weight_name}.weight", - expert_id, - weight_name, - ) - for expert_id in range(self.config.num_local_experts) - for weight_name in ["w1", "w2", "w3"] - ] - for param_name, weight_name, expert_id, shard_id in expert_params_mapping: - name_expert_id = get_expert_id(name) - if name_expert_id is not None and int(name_expert_id) != int(expert_id): - continue - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name, self): - return - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader = weight_loader_with_alias(name)(weight_loader) - weight_loader( - param, - loaded_weight, - weight_name, - expert_id=expert_id, - shard_id=shard_id, - ) - loaded_params.add(name) - break - else: - if is_pp_missing_parameter(name, self): - return - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader = weight_loader_with_alias(name)(weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return - - def is_shared_mlp_weight(name: str) -> bool: - return "shared_mlp" in name and not name.endswith(".bias") - - def load_shared_mlp_weight( - name: str, loaded_weight: torch.Tensor, self - ) -> None: - if not self.CONCAT_FFN: - if "gate_proj" in name: - name = name.replace("gate_proj", "w1", 1) - elif "up_proj" in name: - name = name.replace("up_proj", "w3", 1) - elif "down_proj" in name: - name = name.replace("down_proj", "w2", 1) - else: - if "gate_proj" in name: - name = name.replace("gate_proj", "gate_up_proj", 1) - loaded_shard_id = 0 - elif "up_proj" in name: - name = name.replace("up_proj", "gate_up_proj", 1) - loaded_shard_id = 1 - if is_pp_missing_parameter(name, self): - return - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader = weight_loader_with_alias(name)(weight_loader) - if not self.CONCAT_FFN: - weight_loader(param, loaded_weight) - else: - if "gate_up_proj" in name: - weight_loader(param, loaded_weight, loaded_shard_id) - elif "down_proj" in name: - weight_loader(param, loaded_weight) - else: - raise AssertionError("MLP weight not in [gate_up_proj, down_proj]") - loaded_params.add(name) - return - - def is_mha_weight(name: str) -> bool: - return "self_attn" in name and not name.endswith(".bias") - - def load_linear_attn_weight( - name: str, loaded_weight: torch.Tensor, self - ) -> None: - if is_pp_missing_parameter(name, self): - return - param = params_dict[name] - - weight_loader = getattr( - param, "weight_loader", MiniMaxText01LinearAttention.weight_direct_load - ) - weight_loader = weight_loader_with_alias(name)(weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return - - def load_flash_attn_weight( - name: str, loaded_weight: torch.Tensor, self - ) -> None: - flash_mha_params_mapping = [ - ("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), - ] - for param_name, weight_name, shard_id in flash_mha_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - if is_pp_missing_parameter(name, self): - return - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader = weight_loader_with_alias(name)(weight_loader) - weight_loader(param, loaded_weight, shard_id) - loaded_params.add(name) - break - else: - if is_pp_missing_parameter(name, self): - return - param = params_dict[name] - - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader = weight_loader_with_alias(name)(weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return - - def is_layer_norm_weight(name: str) -> bool: - return "norm" in name and not name.endswith(".bias") and name in params_dict - - def load_layer_norm_weight( - name: str, loaded_weight: torch.Tensor, self - ) -> None: - if is_pp_missing_parameter(name, self): - return - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader = weight_loader_with_alias(name)(weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return - - def load_basic_weight(name: str, loaded_weight: torch.Tensor, self) -> None: - if is_pp_missing_parameter(name, self): - return - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader = weight_loader_with_alias(name)(weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return - - for name, loaded_weight in weights: - weight_at_layer = which_layer(name) - if weight_at_layer and weight_at_layer >= len(self.decoder_attention_types): - continue - - if is_layer_norm_weight(name): - load_layer_norm_weight(name, loaded_weight, self) - continue - if is_mha_weight(name): - if is_linear_attn_layer(weight_at_layer): - load_linear_attn_weight(name, loaded_weight, self) - else: - load_flash_attn_weight(name, loaded_weight, self) - continue - if is_moe_weight(name): - load_sparse_moe_weight(name, loaded_weight, self) - continue - if is_shared_mlp_weight(name): - load_shared_mlp_weight(name, loaded_weight, self) - continue - - if "rotary_emb.inv_freq" in name: - continue - - load_basic_weight(name, loaded_weight, self) - return loaded_params - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor | IntermediateTensors: - forward_context = get_forward_context() - attn_metadata = forward_context.attn_metadata - - if get_pp_group().is_first_rank: - if inputs_embeds is None: - hidden_states = self.embed_scale * self.embed_tokens(input_ids) - else: - hidden_states = inputs_embeds - residual = None - else: - assert intermediate_tensors is not None - hidden_states = intermediate_tensors["hidden_states"] - residual = intermediate_tensors["residual"] - - for layer in islice(self.layers, self.start_layer, self.end_layer): - hidden_states, residual = layer( - hidden_states=hidden_states, - positions=positions, - attn_metadata=attn_metadata, - residual=residual, - ) - if not get_pp_group().is_last_rank: - return IntermediateTensors( - {"hidden_states": hidden_states, "residual": residual} - ) - if residual is not None: - hidden_states, _ = self.norm(hidden_states, residual) - else: - hidden_states = self.norm(hidden_states) - - return hidden_states - - -class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - config = vllm_config.model_config.hf_config - - self.config = config - - if not hasattr(config, "sliding_window"): - config.sliding_window = None - - self.CONCAT_FFN = True - - if hasattr(vllm_config.model_config, "max_model_len"): - self.config.max_model_len = vllm_config.model_config.max_model_len - self.model = MiniMaxText01Model( - vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") - ) - if get_pp_group().is_last_rank: - self.lm_head = ParallelLMHead( - config.vocab_size, - self.config.hidden_size, - prefix=maybe_prefix(prefix, "lm_head"), - ) - - self.logits_processor = LogitsProcessor( - config.vocab_size, self.config.vocab_size - ) - - else: - self.lm_head = PPMissingLayer() - self.lm_head.float() - flash_layer_count = sum( - 1 for attn_type in self.model.decoder_attention_types if attn_type == 1 - ) - self.kv_cache = [torch.tensor([]) for _ in range(flash_layer_count)] - return - - def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs): - return self.model.minimax_cache.copy_inputs_before_cuda_graphs( - input_buffers, **kwargs - ) - - def get_seqlen_agnostic_capture_inputs(self, batch_size: int): - return self.model.minimax_cache.get_seqlen_agnostic_capture_inputs(batch_size) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.embed_input_ids(input_ids) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs, - ) -> torch.Tensor: - hidden_states = self.model( - input_ids, positions, intermediate_tensors, inputs_embeds, **kwargs - ) - - return hidden_states - - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: - logits = self.logits_processor(self.lm_head, hidden_states.float()) - - return logits - - def make_empty_intermediate_tensors( - self, batch_size: int, dtype: torch.dtype, device: torch.device - ) -> IntermediateTensors: - return IntermediateTensors( - { - "hidden_states": torch.zeros( - (batch_size, self.config.hidden_size), dtype=dtype, device=device - ), - "residual": torch.zeros( - (batch_size, self.config.hidden_size), dtype=dtype, device=device - ), - } - ) - - @classmethod - def get_mamba_state_dtype_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[torch.dtype, torch.dtype]: - return MambaStateDtypeCalculator.linear_attention_state_dtype( - vllm_config.model_config.dtype, - vllm_config.cache_config.mamba_cache_dtype, - ) - - @classmethod - def get_mamba_state_shape_from_config( - cls, - vllm_config: "VllmConfig", - ) -> tuple[tuple[int, ...], ...]: - """Calculate shape for MiniMaxText01LinearAttention cache. - - Args: - vllm_config: vLLM config - - Returns: - Tuple containing: - - state_shape: Shape of the cache - """ - parallel_config = vllm_config.parallel_config - hf_config = vllm_config.model_config.hf_config - - return MambaStateShapeCalculator.linear_attention_state_shape( - num_heads=hf_config.num_attention_heads, - tp_size=parallel_config.tensor_parallel_size, - head_dim=hf_config.head_dim, - ) - - @classmethod - def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc]: - return MambaStateCopyFuncCalculator.linear_attention_state_copy_func() - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/minimax_vl_01.py b/vllm/model_executor/models/minimax_vl_01.py deleted file mode 100644 index ccbd4f98d8b..00000000000 --- a/vllm/model_executor/models/minimax_vl_01.py +++ /dev/null @@ -1,385 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from collections.abc import Iterable, Mapping -from typing import Annotated, Literal, TypeAlias - -import torch -import torch.nn as nn -from transformers import BatchFeature, PretrainedConfig -from transformers.models.llava_next.modeling_llava_next import ( - get_anyres_image_grid_shape, - unpad_image, -) - -from vllm.config import VllmConfig -from vllm.model_executor.layers.activation import get_act_fn -from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear -from vllm.model_executor.layers.quantization import QuantizationConfig -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import MultiModalFieldConfig -from vllm.sequence import IntermediateTensors -from vllm.utils.tensor_schema import TensorSchema, TensorShape - -from .clip import CLIPVisionModel -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP -from .llava import ( - BaseLlavaMultiModalProcessor, - LlavaDummyInputsBuilder, - init_vision_tower_for_llava, -) -from .llava_next import LlavaNextProcessingInfo -from .pixtral import PixtralHFVisionModel -from .siglip import SiglipVisionModel -from .utils import ( - AutoWeightsLoader, - init_vllm_registered_model, - maybe_prefix, -) - - -class MiniMaxVL01ImagePixelInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - np: Number of patches + 1 - - c: Number of channels (3) - - h: Height - - w: Width - - Note that `num_patches` may be different per batch and image, - in which case the data is passed as a list instead of a batched tensor. - """ - - type: Literal["pixel_values"] = "pixel_values" - pixel_values: Annotated[ - torch.Tensor | list[torch.Tensor], - TensorShape("bn", "np", 3, "h", "w", dynamic_dims={"np", "h", "w"}), - ] - - image_sizes: Annotated[torch.Tensor | None, TensorShape("bn", 2)] - # This should be in `(height, width)` format. - - -class MiniMaxVL01ImageEmbeddingInputs(TensorSchema): - """ - Dimensions: - - bn: Batch size * number of images - - ifs: Image feature size - - hs: Hidden size (must match language model backbone) - """ - - type: Literal["image_embeds"] = "image_embeds" - data: Annotated[torch.Tensor, TensorShape("bn", "ifs", "hs")] - - -MiniMaxVL01ImageInputs: TypeAlias = ( - MiniMaxVL01ImagePixelInputs | MiniMaxVL01ImageEmbeddingInputs -) - - -class MiniMaxVL01MultiModalProjector(nn.Module): - def __init__( - self, - vision_hidden_size: int, - text_hidden_size: int, - projector_hidden_act: str, - multimodal_projector_bias: bool, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - - self.linear_1 = ColumnParallelLinear( - vision_hidden_size, - text_hidden_size, - bias=multimodal_projector_bias, - quant_config=quant_config, - prefix=f"{prefix}.linear_1", - ) - self.act = get_act_fn(projector_hidden_act) - self.linear_2 = RowParallelLinear( - text_hidden_size, - text_hidden_size, - bias=multimodal_projector_bias, - quant_config=quant_config, - prefix=f"{prefix}.linear_2", - ) - - def forward(self, image_features: torch.Tensor) -> torch.Tensor: - hidden_states, _ = self.linear_1(image_features) - hidden_states = self.act(hidden_states) - hidden_states, _ = self.linear_2(hidden_states) - return hidden_states - - -class MiniMaxVL01DummyInputsBuilder(LlavaDummyInputsBuilder): - pass - - -class MiniMaxVL01ProcessingInfo(LlavaNextProcessingInfo): - def get_hf_config(self): # Need to override the config type - return self.ctx.get_hf_config(PretrainedConfig) - - def get_hf_processor(self, **kwargs: object): - hf_processor = self.ctx.get_hf_processor(**kwargs) - image_processor = hf_processor.image_processor - image_processor.anyres_preprocess = image_processor.anyres_for_vllm_preprocess - - return hf_processor - - def get_supported_mm_limits(self) -> Mapping[str, int | None]: - return {"image": None} - - -class MiniMaxVL01MultiModalProcessor( - BaseLlavaMultiModalProcessor[MiniMaxVL01ProcessingInfo] -): - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: - processed_outputs = super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) - - pixel_values = processed_outputs.get("pixel_values") - if pixel_values is not None: - # Avoid padding since we need the output for each image to be - # independent of other images for the cache to work correctly - image_sizes = processed_outputs["image_sizes"] - assert len(pixel_values) == len(image_sizes) - - processed_outputs["pixel_values"] = [ - p[:, :h, :w] for p, (h, w) in zip(pixel_values, image_sizes) - ] - - return processed_outputs - - def _get_mm_fields_config( - self, - hf_inputs: BatchFeature, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> Mapping[str, MultiModalFieldConfig]: - return { - "pixel_values": MultiModalFieldConfig.batched("image"), - "image_sizes": MultiModalFieldConfig.batched("image"), - "image_embeds": MultiModalFieldConfig.batched("image"), - } - - -@MULTIMODAL_REGISTRY.register_processor( - MiniMaxVL01MultiModalProcessor, - info=MiniMaxVL01ProcessingInfo, - dummy_inputs=MiniMaxVL01DummyInputsBuilder, -) -class MiniMaxVL01ForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP): - packed_modules_mapping = { - "qkv_proj": ["q_proj", "k_proj", "v_proj"], - "gate_up_proj": ["gate_proj", "up_proj"], - } - - @classmethod - def get_placeholder_str(cls, modality: str, i: int) -> str | None: - if modality.startswith("image"): - return "" - - raise ValueError("Only image modality is supported") - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: - super().__init__() - - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - multimodal_config = vllm_config.model_config.multimodal_config - - self.config = config - self.multimodal_config = multimodal_config - - with self._mark_tower_model(vllm_config, "image"): - self.vision_tower = init_vision_tower_for_llava( - config, - quant_config=quant_config, - require_post_norm=False, - prefix=maybe_prefix(prefix, "vision_tower"), - ) - self.multi_modal_projector = MiniMaxVL01MultiModalProjector( - vision_hidden_size=config.vision_config.hidden_size, - text_hidden_size=config.text_config.hidden_size, - projector_hidden_act=config.projector_hidden_act, - multimodal_projector_bias=True, - quant_config=quant_config, - prefix=maybe_prefix(prefix, "multi_modal_projector"), - ) - self.image_newline = nn.Parameter( - torch.empty(config.text_config.hidden_size) - ) - - with self._mark_language_model(vllm_config): - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - hf_config=config.text_config, - prefix=maybe_prefix(prefix, "language_model"), - ) - - self.vision_feature_layer = config.vision_feature_layer - self.vocab_size = config.text_config.vocab_size - self.pad_token_id = -1 - if self.config.text_config.pad_token_id is not None: - self.pad_token_id = self.config.text_config.pad_token_id - - self.make_empty_intermediate_tensors = ( - self.language_model.make_empty_intermediate_tensors - ) - - def _image_pixels_to_features( - self, - vision_tower: CLIPVisionModel | SiglipVisionModel | PixtralHFVisionModel, - pixel_values: torch.Tensor | list[torch.Tensor], - ) -> torch.Tensor | tuple[torch.Tensor, ...]: - # NOTE: we skip the step to select the vision feature layer since - # this is already done inside the vision tower - feature_select_strategy = self.config.vision_feature_select_strategy - return tuple( - vision_tower(p, feature_select_strategy=feature_select_strategy) - for p in pixel_values - ) - - # adapted from https://huggingface.co/MiniMaxAI/MiniMax-VL-01/blob/main/modeling_minimax_vl_01.py#L616-L631 - def pack_image_features( - self, image_features: list[torch.Tensor], image_sizes: torch.Tensor - ): - new_image_features = [] - for image_idx, image_feature in enumerate(image_features): - if image_feature.shape[0] > 1: - base_image_feature = image_feature[0] - image_feature = image_feature[1:] - height = width = ( - self.config.vision_config.image_size - // self.config.vision_config.patch_size - ) - if height * width != base_image_feature.shape[0]: - raise ValueError( - "The number of patches is not consistent with the image size." - ) - num_patch_height, num_patch_width = get_anyres_image_grid_shape( - image_sizes[image_idx], - self.config.image_grid_pinpoints, - self.config.vision_config.image_size, - ) - - image_feature = image_feature.view( - num_patch_height, num_patch_width, height, width, -1 - ) - image_feature = image_feature.permute(4, 0, 2, 1, 3).contiguous() - image_feature = image_feature.flatten(1, 2).flatten(2, 3) - image_feature = unpad_image(image_feature, image_sizes[image_idx]) - - image_feature = torch.cat( - ( - image_feature, - self.image_newline[:, None, None] - .expand(*image_feature.shape[:-1], 1) - .to(image_feature.dtype), - ), - dim=-1, - ) - image_feature = image_feature.flatten(1, 2).transpose(0, 1) - image_feature = torch.cat((base_image_feature, image_feature), dim=0) - else: - image_feature = image_feature[0] - image_feature = torch.cat( - (image_feature, self.image_newline[None].to(image_feature)), dim=0 - ) - new_image_features.append(image_feature) - return new_image_features - - def _process_image_pixels( - self, - inputs: MiniMaxVL01ImagePixelInputs, - ) -> torch.Tensor | tuple[torch.Tensor, ...]: - pixel_values = inputs["pixel_values"] - return self._image_pixels_to_features(self.vision_tower, pixel_values) - - def _process_image_input( - self, - image_input: MiniMaxVL01ImageInputs, - ) -> torch.Tensor | tuple[torch.Tensor, ...]: - if image_input["type"] == "image_embeds": - return image_input["data"] - - image_features = self._process_image_pixels(image_input) - - if isinstance(image_features, torch.Tensor): - return self.multi_modal_projector(image_features) - - feature_sizes = [image_feature.shape[0] for image_feature in image_features] - - image_embeds = self.multi_modal_projector(torch.cat(image_features)) - image_embeds = torch.split(image_embeds, feature_sizes) - image_sizes = image_input.get("image_sizes") - return self.pack_image_features(image_embeds, image_sizes) - - def _parse_and_validate_image_input( - self, **kwargs: object - ) -> MiniMaxVL01ImageInputs | None: - pixel_values = kwargs.pop("pixel_values", None) - image_sizes = kwargs.pop("image_sizes", None) - image_embeds = kwargs.pop("image_embeds", None) - - if pixel_values is None and image_embeds is None: - return None - - if pixel_values is not None and image_sizes is not None: - return MiniMaxVL01ImagePixelInputs( - type="pixel_values", - pixel_values=pixel_values, - image_sizes=image_sizes, - ) - - if image_embeds is not None: - return MiniMaxVL01ImageEmbeddingInputs( - type="image_embeds", - data=image_embeds, - ) - - raise AssertionError("This line should be unreachable.") - - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: - image_input = self._parse_and_validate_image_input(**kwargs) - if image_input is None: - return [] - - return self._process_image_input(image_input) - - def forward( - self, - input_ids: torch.Tensor | None, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - inputs_embeds: torch.Tensor | None = None, - **kwargs: object, - ) -> torch.Tensor | IntermediateTensors: - if intermediate_tensors is not None: - inputs_embeds = None - - hidden_states = self.language_model.model( - input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds - ) - - return hidden_states - - def compute_logits( - self, - hidden_states: torch.Tensor, - ) -> torch.Tensor | None: - return self.language_model.compute_logits(hidden_states) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loader = AutoWeightsLoader(self) - return loader.load_weights(weights) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 5c023b3f41e..5fb28b1c765 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -159,9 +159,6 @@ _TEXT_GENERATION_MODELS = { "MellumForCausalLM": ("mellum", "MellumForCausalLM"), "MiniCPMForCausalLM": ("minicpm", "MiniCPMForCausalLM"), "MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"), - "MiniMaxForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), - "MiniMaxText01ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), - "MiniMaxM1ForCausalLM": ("minimax_text_01", "MiniMaxText01ForCausalLM"), "MiniMaxM2ForCausalLM": ("minimax_m2", "MiniMaxM2ForCausalLM"), "MiniMaxM3SparseForCausalLM": ( "vllm.models.minimax_m3", @@ -490,10 +487,6 @@ _MULTIMODAL_MODELS = { "vllm.models.minimax_m3", "MiniMaxM3SparseForConditionalGeneration", ), - "MiniMaxVL01ForConditionalGeneration": ( - "minimax_vl_01", - "MiniMaxVL01ForConditionalGeneration", - ), "MiniCPMO": ("minicpmo", "MiniCPMO"), "MiniCPMV": ("minicpmv", "MiniCPMV"), "MiniCPMV4_6ForConditionalGeneration": ( @@ -735,6 +728,10 @@ _PREVIOUSLY_SUPPORTED_MODELS = { "XverseForCausalLM": "0.23.0", "Dots1ForCausalLM": "0.23.0", "BambaForCausalLM": "0.23.0", + "MiniMaxForCausalLM": "0.23.0", + "MiniMaxText01ForCausalLM": "0.23.0", + "MiniMaxM1ForCausalLM": "0.23.0", + "MiniMaxVL01ForConditionalGeneration": "0.23.0", } _OOT_SUPPORTED_MODELS = { diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 7ce1520ffe5..109189a033a 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -130,10 +130,6 @@ _TOOL_PARSERS_TO_REGISTER = { "minimax_m3_tool_parser", "MinimaxM3ToolParser", ), - "minimax": ( - "minimax_tool_parser", - "MinimaxToolParser", - ), "minicpm5": ( "minicpm5xml_tool_parser", "MiniCPM5XMLToolParser", diff --git a/vllm/tool_parsers/minimax_tool_parser.py b/vllm/tool_parsers/minimax_tool_parser.py deleted file mode 100644 index 2a2baa03b0e..00000000000 --- a/vllm/tool_parsers/minimax_tool_parser.py +++ /dev/null @@ -1,852 +0,0 @@ -# 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): - 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 = "" - self.tool_call_end_token = "" - self.tool_call_regex = re.compile( - r"(.*?)|(.*)", re.DOTALL - ) - self.thinking_tag_pattern = r"(.*?)" - 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".*?", "", think_content, flags=re.DOTALL - ) - return f"{cleaned_content}" - - 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("") - close_count = text.count("") - self.in_thinking_tag = open_count > close_count or ( - open_count == close_count and text.endswith("") - ) - - 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"(.*?)", 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